stack/tests/conf/test_connect.py

"""Tests for conf.connect — connection factories."""

from __future__ import annotations

import os
import sqlite3
from unittest.mock import MagicMock, patch

import pytest

from conf import connect, path

_HAS_ACO_DB = path("db.aco").exists()


class TestDuckdb:
    @pytest.mark.skipif(not _HAS_ACO_DB, reason="data/aco.duckdb not present")
    def test_returns_connection(self):
        con = connect.duckdb()
        assert con is not None
        result = con.execute("SELECT 1 AS x").fetchone()
        assert result[0] == 1
        con.close()

    @pytest.mark.skipif(not _HAS_ACO_DB, reason="data/aco.duckdb not present")
    def test_read_only_default(self):
        con = connect.duckdb()
        with pytest.raises(Exception, match="read-only"):
            con.execute("CREATE TABLE _test_ro (x INT)")
        con.close()

    def test_custom_db_name(self, tmp_path, monkeypatch):
        """Verify connect.duckdb() resolves a custom [db] key."""
        # Create a real DuckDB file so it opens without sqlite_scanner
        import duckdb as _duckdb

        db_file = tmp_path / "custom.duckdb"
        _duckdb.connect(str(db_file)).close()

        # Patch cfg.db to point "custom" at our temp file
        monkeypatch.setattr(
            "conf.connect.path",
            lambda key: db_file if key == "db.custom" else path(key),
        )
        con = connect.duckdb("custom")
        result = con.execute("SELECT 1 AS x").fetchone()
        assert result[0] == 1
        con.close()


class TestBib:
    def test_returns_store(self):
        store = connect.bib()
        assert hasattr(store, "list_items")
        assert hasattr(store, "list_tags")
        store.close()


class TestZotero:
    def test_returns_sqlite_connection(self):
        try:
            con = connect.zotero()
            assert isinstance(con, sqlite3.Connection)
            con.close()
        except Exception:
            pytest.skip("Zotero database not available")


class TestTrino:
    def _mock_trino(self):
        """Create a mock trino module with dbapi.connect."""
        mock_mod = MagicMock()
        return mock_mod

    def test_creates_connection(self):
        mock_trino = self._mock_trino()
        with patch.dict("sys.modules", {"trino": mock_trino}):
            connect.trino()
            mock_trino.dbapi.connect.assert_called_once()
            kw = mock_trino.dbapi.connect.call_args.kwargs
            assert kw["host"] == "trino"
            assert kw["port"] == 8080
            assert kw["catalog"] == "iceberg"

    def test_custom_catalog(self):
        mock_trino = self._mock_trino()
        with patch.dict("sys.modules", {"trino": mock_trino}):
            connect.trino(catalog="memory")
            assert mock_trino.dbapi.connect.call_args.kwargs["catalog"] == "memory"

    def test_custom_user(self):
        mock_trino = self._mock_trino()
        with patch.dict("sys.modules", {"trino": mock_trino}):
            connect.trino(user="testuser")
            assert mock_trino.dbapi.connect.call_args.kwargs["user"] == "testuser"

    def test_custom_schema(self):
        mock_trino = self._mock_trino()
        with patch.dict("sys.modules", {"trino": mock_trino}):
            connect.trino(schema="myschema")
            assert mock_trino.dbapi.connect.call_args.kwargs["schema"] == "myschema"


class TestNessie:
    def test_creates_client(self):
        client = connect.nessie(base_url="http://localhost:19120/api/v2")
        assert client.base_url == "http://localhost:19120/api/v2"
        client.close()

    def test_default_url_from_config(self):
        client = connect.nessie()
        assert "nessie" in client.base_url
        client.close()


class TestPolaris:
    def test_creates_client_without_secret(self):
        client = connect.polaris(base_url="http://localhost:8181/api/catalog")
        assert client.base_url == "http://localhost:8181/api/catalog"
        client.close()

    def test_default_url_from_config(self):
        client = connect.polaris()
        assert "polaris" in client.base_url
        client.close()


class TestS3:
    def test_creates_client(self):
        client = connect.s3()
        assert client.bucket == "lakehouse"
        client.close()

    def test_custom_bucket(self):
        client = connect.s3(bucket="exports")
        assert client.bucket == "exports"
        client.close()

    def test_reads_env_vars(self):
        with patch.dict(
            "os.environ",
            {
                "RUSTFS_ENDPOINT": "http://test:9000",
                "RUSTFS_ACCESS_KEY": "testkey",
                "RUSTFS_SECRET_KEY": "testsecret",
            },
        ):
            client = connect.s3()
            assert client.endpoint == "http://test:9000"
            assert client.access_key == "testkey"
            assert client.secret_key == "testsecret"
            client.close()


class TestObstore:
    def test_creates_s3store_and_sets_env(self):
        mock_store = MagicMock()
        env = {
            "RUSTFS_ENDPOINT": "http://test:9000",
            "RUSTFS_ACCESS_KEY": "ak",
            "RUSTFS_SECRET_KEY": "sk",
        }
        with patch("conf.connect.cfg") as mock_cfg:
            mock_cfg.s3.endpoint = "http://rustfs:9000"
            with patch.dict("os.environ", env, clear=False):
                with patch(
                    "obstore.store.S3Store", return_value=mock_store
                ) as mock_cls:
                    result = connect.obstore(bucket="mybucket")
                    mock_cls.assert_called_once_with("mybucket")
                    assert result is mock_store
                    assert os.environ["AWS_ENDPOINT_URL"] == "http://test:9000"
                    assert os.environ["AWS_ACCESS_KEY_ID"] == "ak"
                    assert os.environ["AWS_SECRET_ACCESS_KEY"] == "sk"
                    assert os.environ["AWS_ALLOW_HTTP"] == "true"

    def test_default_bucket(self):
        mock_store = MagicMock()
        with patch("conf.connect.cfg") as mock_cfg:
            mock_cfg.s3.endpoint = "http://rustfs:9000"
            with patch("obstore.store.S3Store", return_value=mock_store) as mock_cls:
                connect.obstore()
                assert mock_cls.call_args.args[0] == "lakehouse"


class TestTheme:
    def test_activates_altair_theme(self):
        connect.theme()
        import altair as alt

        assert alt.theme.active == "fhirworx"


class TestDuckdbBatch:
    def _patch_path(self, monkeypatch, db_file):
        monkeypatch.setattr(
            "conf.connect.path",
            lambda key: db_file if key == "db.custom" else path(key),
        )

    def test_yields_writable_connection_and_closes(self, tmp_path, monkeypatch):
        import duckdb as _duckdb

        db_file = tmp_path / "custom.duckdb"
        _duckdb.connect(str(db_file)).close()
        self._patch_path(monkeypatch, db_file)

        with connect.duckdb_batch("custom") as con:
            con.execute("CREATE TABLE t (x INT)")
            con.execute("INSERT INTO t VALUES (1)")
            kept = con
        # closed on exit
        with pytest.raises(Exception):
            kept.execute("SELECT 1")
        # write persisted
        con2 = connect.duckdb("custom")
        assert con2.execute("SELECT count(*) FROM t").fetchone()[0] == 1
        con2.close()

    def test_lock_held_raises_actionable_error(self, tmp_path, monkeypatch):
        """A connection held by another process blocks the writer — after
        retries the error names the holders instead of a raw IOException
        (#508). Same-process connections share one DuckDB instance, so the
        holder must be a subprocess."""
        import subprocess
        import sys

        import duckdb as _duckdb

        db_file = tmp_path / "custom.duckdb"
        _duckdb.connect(str(db_file)).close()
        self._patch_path(monkeypatch, db_file)

        holder = subprocess.Popen(
            [
                sys.executable,
                "-c",
                "import time, duckdb; "
                f"con = duckdb.connect({str(db_file)!r}); "
                "print('locked', flush=True); time.sleep(60)",
            ],
            stdout=subprocess.PIPE,
            text=True,
        )
        try:
            assert holder.stdout is not None
            assert holder.stdout.readline().strip() == "locked"
            with pytest.raises(RuntimeError, match="(?s)write lock.*Holders"):
                with connect.duckdb_batch("custom", retries=2, backoff_s=0.05):
                    pass
        finally:
            holder.kill()
            holder.wait()

    def test_closes_on_exception_in_body(self, tmp_path, monkeypatch):
        import duckdb as _duckdb

        db_file = tmp_path / "custom.duckdb"
        _duckdb.connect(str(db_file)).close()
        self._patch_path(monkeypatch, db_file)

        with pytest.raises(ValueError, match="boom"):
            with connect.duckdb_batch("custom") as con:
                raise ValueError("boom")
        # lock released — a fresh writer can open immediately
        with connect.duckdb_batch("custom", retries=1) as con:
            con.execute("SELECT 1")


class TestReplica:
    def _patch_path(self, monkeypatch, db_file):
        monkeypatch.setattr(
            "conf.connect.path",
            lambda key: db_file if key == "db.custom" else path(key),
        )

    def _make_primary(self, tmp_path):
        import duckdb as _duckdb

        db_file = tmp_path / "custom.duckdb"
        con = _duckdb.connect(str(db_file))
        con.execute("CREATE TABLE t AS SELECT 1 AS x")
        con.close()
        return db_file

    def test_publish_creates_consistent_snapshot(self, tmp_path, monkeypatch):
        db_file = self._make_primary(tmp_path)
        self._patch_path(monkeypatch, db_file)

        dst = connect.publish_replica("custom")
        assert dst.name == "custom.ro.duckdb"
        import duckdb as _duckdb

        con = _duckdb.connect(str(dst), read_only=True)
        assert con.execute("SELECT x FROM t").fetchone()[0] == 1
        con.close()

    def test_read_only_resolves_to_replica(self, tmp_path, monkeypatch):
        db_file = self._make_primary(tmp_path)
        self._patch_path(monkeypatch, db_file)
        connect.publish_replica("custom")

        # diverge the primary so we can tell which file we read
        with connect.duckdb_batch("custom") as con:
            con.execute("UPDATE t SET x = 2")

        assert connect.duckdb("custom").execute("SELECT x FROM t").fetchone()[0] == 1
        # explicit opt-out reads the live primary
        assert (
            connect.duckdb("custom", replica=False)
            .execute("SELECT x FROM t")
            .fetchone()[0]
            == 2
        )
        # env kill-switch
        monkeypatch.setenv("STACK_DUCKDB_REPLICA", "0")
        assert connect.duckdb("custom").execute("SELECT x FROM t").fetchone()[0] == 2

    def test_no_replica_falls_back_to_primary(self, tmp_path, monkeypatch):
        db_file = self._make_primary(tmp_path)
        self._patch_path(monkeypatch, db_file)
        assert connect.duckdb("custom").execute("SELECT x FROM t").fetchone()[0] == 1

    def test_write_opens_never_use_replica(self, tmp_path, monkeypatch):
        db_file = self._make_primary(tmp_path)
        self._patch_path(monkeypatch, db_file)
        connect.publish_replica("custom")
        con = connect.duckdb("custom", read_only=False)
        con.execute("UPDATE t SET x = 3")
        con.close()
        assert (
            connect.duckdb("custom", replica=False)
            .execute("SELECT x FROM t")
            .fetchone()[0]
            == 3
        )

    def test_republish_refreshes_atomically(self, tmp_path, monkeypatch):
        db_file = self._make_primary(tmp_path)
        self._patch_path(monkeypatch, db_file)
        connect.publish_replica("custom")
        with connect.duckdb_batch("custom") as con:
            con.execute("UPDATE t SET x = 9")
        connect.publish_replica("custom")
        assert connect.duckdb("custom").execute("SELECT x FROM t").fetchone()[0] == 9


class TestDucklakeFactory:
    def test_returns_caller_owned_lake_connection(self, tmp_path, monkeypatch):
        """connect.ducklake() with a file catalog: connection survives the
        internal context's GC and resolves lake tables as the default db."""
        from aco.lake import DuckLakeContext

        # Seed a lake with a file catalog + local data path
        seed = DuckLakeContext(
            catalog_dsn=str(tmp_path / "cat.ducklake"),
            data_path=str(tmp_path / "data") + "/",
            read_only=False,
        )
        import narwhals as nw
        import polars as pl

        seed.save(
            "opps.addendum_b",
            nw.from_native(pl.DataFrame({"hcpcs": ["Q4186"], "year": [2026]})),
            mode="replace",
        )
        seed.close()

        from types import SimpleNamespace

        _dl = SimpleNamespace(
            catalog=str(tmp_path / "cat.ducklake"),
            catalog_ro=str(tmp_path / "cat.ducklake"),
            data_path=str(tmp_path / "data") + "/",
            s3_endpoint="",
        )
        # cfg.lake materializes fresh objects per access — patch the whole
        # cfg reference connect.ducklake() reads.
        monkeypatch.setattr(
            connect, "cfg", SimpleNamespace(lake=SimpleNamespace(ducklake=_dl))
        )
        import gc

        con = connect.ducklake()
        gc.collect()  # the internal context must not close our connection
        assert con.execute("SELECT count(*) FROM opps.addendum_b").fetchone()[0] == 1
        with pytest.raises(Exception, match='type "CREATE"'):
            con.execute("CREATE TABLE opps.nope (x INT)")
        con.close()