4243 lines
148 KiB
Python
4243 lines
148 KiB
Python
"""Tests for aco.lake modules — catalog, context, engine, transpile.
|
|
|
|
Focuses on pure functions, class construction, validation logic,
|
|
SQL generation helpers, and string processing that can be tested
|
|
WITHOUT external connections (DuckDB files, Iceberg, Trino).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import date, datetime
|
|
from decimal import Decimal
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
import pytest
|
|
from databricks.sdk import WorkspaceClient
|
|
from databricks.sdk.service.catalog import (
|
|
CatalogInfo,
|
|
ColumnInfo,
|
|
SchemaInfo,
|
|
TableInfo,
|
|
VolumeInfo,
|
|
)
|
|
|
|
from aco.lake.catalog import Catalog
|
|
from aco.lake.context import (
|
|
Context,
|
|
DuckDBContext,
|
|
EnterpriseContext,
|
|
IcebergContext,
|
|
ParquetContext,
|
|
TrinoContext,
|
|
)
|
|
from aco.lake.engine import execute
|
|
from aco.lake.transpile import (
|
|
_PYTHON_TO_DUCKDB,
|
|
_clean_duckdb_aliases,
|
|
_python_type_to_duckdb,
|
|
_table_ref_to_view_name,
|
|
_transpile_sql,
|
|
)
|
|
from aco.pipe.base import Pipeline
|
|
from aco.table.base import SQLTable
|
|
|
|
# ── fixtures ──────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.fixture
|
|
def offline_catalog():
|
|
"""Catalog with no Iceberg connection (offline mode)."""
|
|
return Catalog()
|
|
|
|
|
|
@pytest.fixture
|
|
def catalog_with_schema_map():
|
|
"""Catalog with schema and column mapping configured."""
|
|
return Catalog(
|
|
schema_map={
|
|
"core.encounter": "prod.encounters",
|
|
"core.patient": "prod.patients",
|
|
},
|
|
column_map={
|
|
"core.encounter": {
|
|
"encounter_id": "encntr_id",
|
|
"patient_id": "pat_id",
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_fn():
|
|
"""A simple narwhals expression function."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""Passthrough."""
|
|
return df
|
|
|
|
return fn
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Catalog
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestCatalogInit:
|
|
"""Catalog.__init__ stores config correctly."""
|
|
|
|
def test_default_offline_mode(self) -> None:
|
|
cat = Catalog()
|
|
assert cat.catalog_uri == ""
|
|
assert cat.warehouse == ""
|
|
assert cat.properties == {}
|
|
assert cat.schema_map == {}
|
|
assert cat.column_map == {}
|
|
|
|
def test_catalog_uri_stored(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert cat.catalog_uri == "http://nessie:19120/iceberg/"
|
|
|
|
def test_warehouse_stored(self) -> None:
|
|
cat = Catalog(warehouse="s3://lakehouse/")
|
|
assert cat.warehouse == "s3://lakehouse/"
|
|
|
|
def test_properties_stored(self) -> None:
|
|
props = {"credential": "root:secret"}
|
|
cat = Catalog(properties=props)
|
|
assert cat.properties == props
|
|
|
|
def test_properties_default_to_empty_dict(self) -> None:
|
|
cat = Catalog()
|
|
assert cat.properties == {}
|
|
assert isinstance(cat.properties, dict)
|
|
|
|
def test_schema_map_stored(self) -> None:
|
|
smap = {"core.encounter": "prod.encounters"}
|
|
cat = Catalog(schema_map=smap)
|
|
assert cat.schema_map == smap
|
|
|
|
def test_column_map_stored(self) -> None:
|
|
cmap = {
|
|
"core.encounter": {"encounter_id": "encntr_id"},
|
|
}
|
|
cat = Catalog(column_map=cmap)
|
|
assert cat.column_map == cmap
|
|
|
|
def test_iceberg_catalog_initially_none(self) -> None:
|
|
cat = Catalog()
|
|
assert cat._iceberg_catalog is None
|
|
|
|
def test_model_cache_initially_empty(self) -> None:
|
|
cat = Catalog()
|
|
assert cat._model_cache == {}
|
|
|
|
|
|
class TestCatalogPhysicalTableRef:
|
|
"""Catalog.physical_table_ref maps canonical to physical names."""
|
|
|
|
def test_mapped_table_returns_physical(self, catalog_with_schema_map) -> None:
|
|
result = catalog_with_schema_map.physical_table_ref("core.encounter")
|
|
assert result == "prod.encounters"
|
|
|
|
def test_unmapped_table_returns_input(self, catalog_with_schema_map) -> None:
|
|
result = catalog_with_schema_map.physical_table_ref("readmissions.encounter")
|
|
assert result == "readmissions.encounter"
|
|
|
|
def test_empty_schema_map_returns_input(self, offline_catalog) -> None:
|
|
result = offline_catalog.physical_table_ref("core.encounter")
|
|
assert result == "core.encounter"
|
|
|
|
def test_all_mapped_entries(self, catalog_with_schema_map) -> None:
|
|
assert (
|
|
catalog_with_schema_map.physical_table_ref("core.patient")
|
|
== "prod.patients"
|
|
)
|
|
|
|
|
|
class TestCatalogPhysicalColumnName:
|
|
"""Catalog.physical_column_name maps canonical to physical columns."""
|
|
|
|
def test_mapped_column_returns_physical(self, catalog_with_schema_map) -> None:
|
|
result = catalog_with_schema_map.physical_column_name(
|
|
"core.encounter", "encounter_id"
|
|
)
|
|
assert result == "encntr_id"
|
|
|
|
def test_unmapped_column_returns_input(self, catalog_with_schema_map) -> None:
|
|
result = catalog_with_schema_map.physical_column_name(
|
|
"core.encounter", "admission_date"
|
|
)
|
|
assert result == "admission_date"
|
|
|
|
def test_unmapped_table_returns_input(self, catalog_with_schema_map) -> None:
|
|
result = catalog_with_schema_map.physical_column_name(
|
|
"readmissions.encounter", "encounter_id"
|
|
)
|
|
assert result == "encounter_id"
|
|
|
|
def test_empty_column_map_returns_input(self, offline_catalog) -> None:
|
|
result = offline_catalog.physical_column_name("core.encounter", "encounter_id")
|
|
assert result == "encounter_id"
|
|
|
|
def test_multiple_columns_mapped(self, catalog_with_schema_map) -> None:
|
|
assert (
|
|
catalog_with_schema_map.physical_column_name("core.encounter", "patient_id")
|
|
== "pat_id"
|
|
)
|
|
|
|
|
|
class TestCatalogModelLookup:
|
|
"""Catalog.model validation and error handling."""
|
|
|
|
def test_unqualified_name_raises_value_error(self, offline_catalog) -> None:
|
|
with pytest.raises(ValueError, match="qualified"):
|
|
offline_catalog.model("encounter")
|
|
|
|
def test_nonexistent_table_raises_value_error(self, offline_catalog) -> None:
|
|
with pytest.raises(ValueError, match="No SQLTable"):
|
|
offline_catalog.model("nonexistent.table_xyz")
|
|
|
|
def test_model_cache_populated_after_lookup(self) -> None:
|
|
"""After successful model() call, cache should contain it."""
|
|
cat = Catalog()
|
|
# Pre-populate cache manually
|
|
mock_model = type(
|
|
"MockTable",
|
|
(SQLTable,),
|
|
{"__schema__": "test", "__tablename__": "t"},
|
|
)
|
|
cat._model_cache["test.t"] = mock_model
|
|
result = cat.model("test.t")
|
|
assert result is mock_model
|
|
|
|
|
|
class TestCatalogIcebergOfflineGuard:
|
|
"""Iceberg methods must fail gracefully in offline mode."""
|
|
|
|
def test_iceberg_namespaces_raises_in_offline_mode(self, offline_catalog) -> None:
|
|
with pytest.raises(RuntimeError, match="offline"):
|
|
offline_catalog.iceberg_namespaces()
|
|
|
|
def test_iceberg_tables_raises_in_offline_mode(self, offline_catalog) -> None:
|
|
with pytest.raises(RuntimeError, match="offline"):
|
|
offline_catalog.iceberg_tables("core")
|
|
|
|
def test_iceberg_schema_raises_in_offline_mode(self, offline_catalog) -> None:
|
|
with pytest.raises(RuntimeError, match="offline"):
|
|
offline_catalog.iceberg_schema("core.encounter")
|
|
|
|
def test_iceberg_snapshots_raises_in_offline_mode(self, offline_catalog) -> None:
|
|
with pytest.raises(RuntimeError, match="offline"):
|
|
offline_catalog.iceberg_snapshots("core.encounter")
|
|
|
|
def test_iceberg_current_snapshot_raises_in_offline(self, offline_catalog) -> None:
|
|
with pytest.raises(RuntimeError, match="offline"):
|
|
offline_catalog.iceberg_current_snapshot("core.encounter")
|
|
|
|
def test_get_iceberg_catalog_raises_in_offline(self, offline_catalog) -> None:
|
|
with pytest.raises(RuntimeError, match="catalog_uri"):
|
|
offline_catalog._get_iceberg_catalog()
|
|
|
|
|
|
class TestCatalogSchemaDiscovery:
|
|
"""Catalog.schemas() discovers schemas from aco.table packages."""
|
|
|
|
def test_schemas_returns_list(self, offline_catalog) -> None:
|
|
result = offline_catalog.schemas()
|
|
assert isinstance(result, list)
|
|
|
|
def test_schemas_are_sorted(self, offline_catalog) -> None:
|
|
result = offline_catalog.schemas()
|
|
assert result == sorted(result)
|
|
|
|
def test_schemas_contains_known_schemas(self, offline_catalog) -> None:
|
|
"""At least 'core' should exist in aco.table."""
|
|
result = offline_catalog.schemas()
|
|
# The project has aco.table modules — at least some schemas
|
|
assert len(result) > 0
|
|
|
|
def test_schemas_are_unique(self, offline_catalog) -> None:
|
|
result = offline_catalog.schemas()
|
|
assert len(result) == len(set(result))
|
|
|
|
|
|
class TestCatalogTableDiscovery:
|
|
"""Catalog.tables() discovers tables within a schema."""
|
|
|
|
def test_tables_returns_list(self, offline_catalog) -> None:
|
|
schemas = offline_catalog.schemas()
|
|
if schemas:
|
|
result = offline_catalog.tables(schemas[0])
|
|
assert isinstance(result, list)
|
|
|
|
def test_tables_are_qualified(self, offline_catalog) -> None:
|
|
"""Every table ref must contain a dot."""
|
|
schemas = offline_catalog.schemas()
|
|
if schemas:
|
|
tables = offline_catalog.tables(schemas[0])
|
|
for t in tables:
|
|
assert "." in t, f"Unqualified table: {t}"
|
|
|
|
def test_tables_are_sorted(self, offline_catalog) -> None:
|
|
schemas = offline_catalog.schemas()
|
|
if schemas:
|
|
tables = offline_catalog.tables(schemas[0])
|
|
assert tables == sorted(tables)
|
|
|
|
def test_nonexistent_schema_returns_empty(self, offline_catalog) -> None:
|
|
result = offline_catalog.tables("nonexistent_schema_xyz")
|
|
assert result == []
|
|
|
|
|
|
class TestCatalogTableNamespaces:
|
|
"""Catalog._TABLE_NAMESPACES is correctly defined."""
|
|
|
|
def test_namespaces_is_list(self) -> None:
|
|
assert isinstance(Catalog._TABLE_NAMESPACES, list)
|
|
|
|
def test_namespaces_contains_aco(self) -> None:
|
|
assert "aco" in Catalog._TABLE_NAMESPACES
|
|
|
|
def test_namespaces_contains_expected_entries(self) -> None:
|
|
expected = {"aco", "bcda", "ccw", "cms", "pfs"}
|
|
actual = set(Catalog._TABLE_NAMESPACES)
|
|
assert expected == actual
|
|
|
|
|
|
class TestCatalogExceptionBranches:
|
|
"""Cover except Exception: continue branches in catalog methods."""
|
|
|
|
def test_schemas_continues_on_import_error(self) -> None:
|
|
"""schemas() skips modules that fail to import."""
|
|
cat = Catalog()
|
|
# Mock _discover_table_modules to return a bad module
|
|
with patch.object(cat, "_discover_table_modules") as mock_disc:
|
|
mock_disc.return_value = [("nonexistent_ns_xyz", "badmod")]
|
|
result = cat.schemas()
|
|
assert result == []
|
|
|
|
def test_tables_continues_on_import_error(self) -> None:
|
|
"""tables() skips modules that fail to import."""
|
|
cat = Catalog()
|
|
with patch.object(cat, "_discover_table_modules") as mock_disc:
|
|
mock_disc.return_value = [("nonexistent_ns_xyz", "badmod")]
|
|
result = cat.tables("core")
|
|
assert result == []
|
|
|
|
def test_model_continues_on_import_error(self) -> None:
|
|
"""model() skips modules that fail to import."""
|
|
cat = Catalog()
|
|
with patch.object(cat, "_discover_table_modules") as mock_disc:
|
|
mock_disc.return_value = [("nonexistent_ns_xyz", "badmod")]
|
|
with pytest.raises(ValueError, match="No SQLTable"):
|
|
cat.model("core.encounter")
|
|
|
|
def test_model_finds_table_via_search(self) -> None:
|
|
"""model() populates cache when finding a table via search."""
|
|
cat = Catalog()
|
|
# Ensure no cache hit — use a fresh catalog with empty cache
|
|
assert cat._model_cache == {}
|
|
|
|
# Create a mock module with a real SQLTable subclass
|
|
FakeModel = type(
|
|
"FakeModel",
|
|
(SQLTable,),
|
|
{"__schema__": "myns", "__tablename__": "mytbl"},
|
|
)
|
|
mock_mod = MagicMock()
|
|
mock_mod.__name__ = "fake_mod"
|
|
# inspect.getmembers needs real class iteration
|
|
import types
|
|
|
|
mock_mod_real = types.ModuleType("fake_mod")
|
|
mock_mod_real.FakeModel = FakeModel
|
|
mock_mod_real.SQLTable = SQLTable
|
|
|
|
with (
|
|
patch.object(cat, "_discover_table_modules") as mock_disc,
|
|
patch("importlib.import_module") as mock_import,
|
|
):
|
|
mock_disc.return_value = [("myns", "mytbl")]
|
|
mock_import.return_value = mock_mod_real
|
|
|
|
result = cat.model("myns.mytbl")
|
|
assert result is FakeModel
|
|
assert "myns.mytbl" in cat._model_cache
|
|
|
|
def test_discover_table_modules_skips_missing_ns(self) -> None:
|
|
"""_discover_table_modules skips namespaces without table pkg."""
|
|
cat = Catalog()
|
|
# Add a namespace that doesn't exist
|
|
original = cat._TABLE_NAMESPACES[:]
|
|
try:
|
|
cat._TABLE_NAMESPACES.append("totally_fake_ns")
|
|
result = cat._discover_table_modules()
|
|
# Should still return the valid ones
|
|
assert isinstance(result, list)
|
|
# The fake ns should not break anything
|
|
finally:
|
|
cat._TABLE_NAMESPACES[:] = original
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Context
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestContextBase:
|
|
"""Base Context class behavior."""
|
|
|
|
def test_load_raises_not_implemented(self) -> None:
|
|
ctx = Context()
|
|
with pytest.raises(NotImplementedError):
|
|
ctx.load("core.encounter")
|
|
|
|
def test_save_raises_not_implemented(self) -> None:
|
|
ctx = Context()
|
|
df = pl.DataFrame({"x": [1]})
|
|
with pytest.raises(NotImplementedError):
|
|
ctx.save("core.encounter", df)
|
|
|
|
|
|
class TestDuckDBContextInit:
|
|
"""DuckDBContext construction and config validation."""
|
|
|
|
def test_database_path_stored(self) -> None:
|
|
ctx = DuckDBContext(database="test.duckdb")
|
|
assert ctx.database == "test.duckdb"
|
|
|
|
def test_read_only_defaults_true(self) -> None:
|
|
ctx = DuckDBContext(database="test.duckdb")
|
|
assert ctx.read_only is True
|
|
|
|
def test_read_only_can_be_set_false(self) -> None:
|
|
ctx = DuckDBContext(database="test.duckdb", read_only=False)
|
|
assert ctx.read_only is False
|
|
|
|
def test_catalog_defaults_none(self) -> None:
|
|
ctx = DuckDBContext(database="test.duckdb")
|
|
assert ctx.catalog is None
|
|
|
|
def test_catalog_can_be_set(self) -> None:
|
|
cat = Catalog(schema_map={"core.encounter": "prod.enc"})
|
|
ctx = DuckDBContext(database="test.duckdb", catalog=cat)
|
|
assert ctx.catalog is cat
|
|
|
|
def test_connection_initially_none(self) -> None:
|
|
ctx = DuckDBContext(database="test.duckdb")
|
|
assert ctx._connection is None
|
|
|
|
def test_load_rejects_unqualified_ref(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
with pytest.raises(ValueError, match="qualified"):
|
|
ctx.load("encounter")
|
|
|
|
def test_save_rejects_read_only(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
df = pl.DataFrame({"x": [1]})
|
|
with pytest.raises(RuntimeError, match="read-only"):
|
|
ctx.save("core.encounter", df)
|
|
|
|
def test_save_rejects_unqualified_ref(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:", read_only=False)
|
|
df = pl.DataFrame({"x": [1]})
|
|
with pytest.raises(ValueError, match="qualified"):
|
|
ctx.save("encounter", df)
|
|
|
|
|
|
class TestParquetContextInit:
|
|
"""ParquetContext construction and path resolution."""
|
|
|
|
def test_base_path_stored(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
assert ctx.base_path == "./data"
|
|
|
|
def test_default_pattern(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
assert ctx.pattern == "{schema}/{table}.parquet"
|
|
|
|
def test_custom_pattern(self) -> None:
|
|
ctx = ParquetContext(
|
|
base_path="./data",
|
|
pattern="{schema}/{table}/*.parquet",
|
|
)
|
|
assert ctx.pattern == "{schema}/{table}/*.parquet"
|
|
|
|
def test_storage_options_default_empty(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
assert ctx.storage_options == {}
|
|
|
|
def test_storage_options_can_be_set(self) -> None:
|
|
opts = {"key": "AKID", "secret": "SECRET"}
|
|
ctx = ParquetContext(base_path="s3://bucket", storage_options=opts)
|
|
assert ctx.storage_options == opts
|
|
|
|
def test_catalog_defaults_none(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
assert ctx.catalog is None
|
|
|
|
def test_resolve_path_default_pattern(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
path = ctx._resolve_path("core.encounter")
|
|
assert path == "./data/core/encounter.parquet"
|
|
|
|
def test_resolve_path_custom_pattern(self) -> None:
|
|
ctx = ParquetContext(
|
|
base_path="/warehouse",
|
|
pattern="{schema}/{table}/data.parquet",
|
|
)
|
|
path = ctx._resolve_path("readmissions.encounter")
|
|
assert path == ("/warehouse/readmissions/encounter/data.parquet")
|
|
|
|
def test_resolve_path_s3(self) -> None:
|
|
ctx = ParquetContext(base_path="s3://my-bucket/lake")
|
|
path = ctx._resolve_path("core.patient")
|
|
assert path == ("s3://my-bucket/lake/core/patient.parquet")
|
|
|
|
def test_resolve_path_strips_trailing_slash(self) -> None:
|
|
ctx = ParquetContext(base_path="./data/")
|
|
path = ctx._resolve_path("core.encounter")
|
|
assert path == "./data/core/encounter.parquet"
|
|
|
|
def test_resolve_path_with_catalog_mapping(self) -> None:
|
|
cat = Catalog(schema_map={"core.encounter": "prod.encounters"})
|
|
ctx = ParquetContext(base_path="./data", catalog=cat)
|
|
path = ctx._resolve_path("core.encounter")
|
|
assert path == "./data/prod/encounters.parquet"
|
|
|
|
def test_load_rejects_unqualified_ref(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
with pytest.raises(ValueError, match="qualified"):
|
|
ctx.load("encounter")
|
|
|
|
def test_save_rejects_unqualified_ref(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
df = pl.DataFrame({"x": [1]})
|
|
with pytest.raises(ValueError, match="qualified"):
|
|
ctx.save("encounter", df)
|
|
|
|
|
|
class TestIcebergContextInit:
|
|
"""IcebergContext construction and defaults."""
|
|
|
|
def test_catalog_uri_stored(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx.catalog_uri == ("http://nessie:19120/iceberg/")
|
|
|
|
def test_catalog_type_defaults_rest(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx.catalog_type == "rest"
|
|
|
|
def test_warehouse_defaults(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx.warehouse == "s3://lakehouse/"
|
|
|
|
def test_namespace_defaults_empty(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx.namespace == ""
|
|
|
|
def test_ref_defaults_main(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx.ref == "main"
|
|
|
|
def test_properties_defaults_empty(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx.properties == {}
|
|
|
|
def test_custom_properties(self) -> None:
|
|
props = {
|
|
"s3.endpoint": "http://rustfs:9000",
|
|
"s3.path-style-access": "true",
|
|
}
|
|
ctx = IcebergContext(
|
|
catalog_uri="http://nessie:19120/iceberg/",
|
|
properties=props,
|
|
)
|
|
assert ctx.properties == props
|
|
|
|
def test_custom_ref(self) -> None:
|
|
ctx = IcebergContext(
|
|
catalog_uri="http://nessie:19120/iceberg/",
|
|
ref="develop",
|
|
)
|
|
assert ctx.ref == "develop"
|
|
|
|
def test_pyiceberg_catalog_initially_none(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
assert ctx._pyiceberg_catalog is None
|
|
|
|
|
|
class TestTrinoContextInit:
|
|
"""TrinoContext construction and defaults."""
|
|
|
|
def test_host_defaults_trino(self) -> None:
|
|
ctx = TrinoContext()
|
|
assert ctx.host == "trino"
|
|
|
|
def test_port_defaults_8080(self) -> None:
|
|
ctx = TrinoContext()
|
|
assert ctx.port == 8080
|
|
|
|
def test_catalog_defaults_iceberg(self) -> None:
|
|
ctx = TrinoContext()
|
|
assert ctx.catalog == "iceberg"
|
|
|
|
def test_schema_defaults_empty(self) -> None:
|
|
ctx = TrinoContext()
|
|
assert ctx.schema_ == ""
|
|
|
|
def test_custom_host_port(self) -> None:
|
|
ctx = TrinoContext(host="trino-prod", port=443)
|
|
assert ctx.host == "trino-prod"
|
|
assert ctx.port == 443
|
|
|
|
def test_custom_catalog(self) -> None:
|
|
ctx = TrinoContext(catalog="hive")
|
|
assert ctx.catalog == "hive"
|
|
|
|
|
|
class TestEnterpriseContextInit:
|
|
"""EnterpriseContext construction and defaults."""
|
|
|
|
def test_catalog_uri_stored(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://workspace.databricks.com/api")
|
|
assert ctx.catalog_uri == ("https://workspace.databricks.com/api")
|
|
|
|
def test_catalog_type_defaults_rest(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
assert ctx.catalog_type == "rest"
|
|
|
|
def test_warehouse_defaults_empty(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
assert ctx.warehouse == ""
|
|
|
|
def test_dialect_defaults_empty(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
assert ctx.dialect == ""
|
|
|
|
def test_connection_string_defaults_empty(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
assert ctx.connection_string == ""
|
|
|
|
def test_properties_defaults_empty(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
assert ctx.properties == {}
|
|
|
|
def test_custom_dialect(self) -> None:
|
|
ctx = EnterpriseContext(
|
|
catalog_uri="https://example.com/api",
|
|
dialect="databricks",
|
|
)
|
|
assert ctx.dialect == "databricks"
|
|
|
|
def test_custom_warehouse(self) -> None:
|
|
ctx = EnterpriseContext(
|
|
catalog_uri="https://example.com/api",
|
|
warehouse="main",
|
|
)
|
|
assert ctx.warehouse == "main"
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Engine
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestEngineDispatch:
|
|
"""execute() dispatches correctly by context type."""
|
|
|
|
def test_unsupported_context_type_raises(self) -> None:
|
|
"""A plain Context should raise TypeError."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
from aco.express.base import Expr
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = Context()
|
|
|
|
with pytest.raises(TypeError, match="Unsupported"):
|
|
execute(pipeline, ctx)
|
|
|
|
def test_duckdb_context_dispatches_to_direct(self) -> None:
|
|
"""DuckDBContext should go through _execute_direct."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
from aco.express.base import Expr
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = DuckDBContext(database=":memory:")
|
|
|
|
with patch("aco.lake.engine._execute_direct") as mock_direct:
|
|
mock_direct.return_value = {}
|
|
execute(pipeline, ctx)
|
|
mock_direct.assert_called_once()
|
|
|
|
def test_parquet_context_dispatches_to_direct(self) -> None:
|
|
"""ParquetContext should go through _execute_direct."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
from aco.express.base import Expr
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = ParquetContext(base_path="./data")
|
|
|
|
with patch("aco.lake.engine._execute_direct") as mock_direct:
|
|
mock_direct.return_value = {}
|
|
execute(pipeline, ctx)
|
|
mock_direct.assert_called_once()
|
|
|
|
def test_iceberg_context_dispatches_to_direct(self) -> None:
|
|
"""IcebergContext should go through _execute_direct."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
from aco.express.base import Expr
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/iceberg/")
|
|
|
|
with patch("aco.lake.engine._execute_direct") as mock_direct:
|
|
mock_direct.return_value = {}
|
|
execute(pipeline, ctx)
|
|
mock_direct.assert_called_once()
|
|
|
|
def test_trino_context_dispatches_to_transpiled(
|
|
self,
|
|
) -> None:
|
|
"""TrinoContext should go through _execute_transpiled."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
from aco.express.base import Expr
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = TrinoContext()
|
|
|
|
with patch("aco.lake.engine._execute_transpiled") as mock_tp:
|
|
mock_tp.return_value = {}
|
|
execute(pipeline, ctx)
|
|
mock_tp.assert_called_once()
|
|
|
|
def test_enterprise_context_dispatches_to_transpiled(
|
|
self,
|
|
) -> None:
|
|
"""EnterpriseContext -> _execute_transpiled."""
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
from aco.express.base import Expr
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
|
|
with patch("aco.lake.engine._execute_transpiled") as mock_tp:
|
|
mock_tp.return_value = {}
|
|
execute(pipeline, ctx)
|
|
mock_tp.assert_called_once()
|
|
|
|
|
|
class TestExecuteDirectWithMocks:
|
|
"""_execute_direct with mocked context.load/save."""
|
|
|
|
def test_direct_passes_results_back(self) -> None:
|
|
"""Direct execution returns the pipeline cache."""
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_direct
|
|
|
|
source = pl.DataFrame({"x": [1, 2, 3]})
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = MagicMock(spec=DuckDBContext)
|
|
ctx.load = MagicMock(return_value=source)
|
|
ctx.save = MagicMock()
|
|
|
|
results = _execute_direct(pipeline, ctx)
|
|
assert "out.table" in results
|
|
|
|
def test_direct_does_not_save_by_default(self) -> None:
|
|
"""save_outputs=False means ctx.save is never called."""
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_direct
|
|
|
|
source = pl.DataFrame({"x": [1]})
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = MagicMock(spec=DuckDBContext)
|
|
ctx.load = MagicMock(return_value=source)
|
|
ctx.save = MagicMock()
|
|
|
|
_execute_direct(pipeline, ctx, save_outputs=False)
|
|
ctx.save.assert_not_called()
|
|
|
|
def test_direct_saves_when_requested(self) -> None:
|
|
"""save_outputs=True calls ctx.save for each output."""
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_direct
|
|
|
|
source = pl.DataFrame({"x": [1]})
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = MagicMock(spec=DuckDBContext)
|
|
ctx.load = MagicMock(return_value=source)
|
|
ctx.save = MagicMock()
|
|
|
|
_execute_direct(pipeline, ctx, save_outputs=True)
|
|
ctx.save.assert_called_once()
|
|
|
|
def test_direct_output_filter_limits_saves(self) -> None:
|
|
"""Only tables in output_filter get saved."""
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_direct
|
|
|
|
source = pl.DataFrame({"x": [1]})
|
|
|
|
@nw.narwhalify
|
|
def fn_a(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
@nw.narwhalify
|
|
def fn_b(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(
|
|
exprs=[
|
|
Expr(name="a.table", fn=fn_a),
|
|
Expr(name="b.table", fn=fn_b),
|
|
]
|
|
)
|
|
ctx = MagicMock(spec=DuckDBContext)
|
|
ctx.load = MagicMock(return_value=source)
|
|
ctx.save = MagicMock()
|
|
|
|
_execute_direct(
|
|
pipeline,
|
|
ctx,
|
|
save_outputs=True,
|
|
output_filter={"a.table"},
|
|
)
|
|
# Only a.table should be saved
|
|
assert ctx.save.call_count == 1
|
|
call_args = ctx.save.call_args
|
|
assert call_args[0][0] == "a.table"
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Transpile
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestPythonTypeToDuckdb:
|
|
"""_python_type_to_duckdb maps Python annotations to DDL."""
|
|
|
|
def test_str_maps_to_varchar(self) -> None:
|
|
assert _python_type_to_duckdb(str) == "VARCHAR"
|
|
|
|
def test_int_maps_to_bigint(self) -> None:
|
|
assert _python_type_to_duckdb(int) == "BIGINT"
|
|
|
|
def test_float_maps_to_double(self) -> None:
|
|
assert _python_type_to_duckdb(float) == "DOUBLE"
|
|
|
|
def test_bool_maps_to_boolean(self) -> None:
|
|
assert _python_type_to_duckdb(bool) == "BOOLEAN"
|
|
|
|
def test_date_maps_to_date(self) -> None:
|
|
assert _python_type_to_duckdb(date) == "DATE"
|
|
|
|
def test_datetime_maps_to_timestamp(self) -> None:
|
|
assert _python_type_to_duckdb(datetime) == "TIMESTAMP"
|
|
|
|
def test_decimal_maps_to_decimal(self) -> None:
|
|
assert _python_type_to_duckdb(Decimal) == "DECIMAL(18,2)"
|
|
|
|
def test_unknown_type_defaults_to_varchar(self) -> None:
|
|
assert _python_type_to_duckdb(bytes) == "VARCHAR"
|
|
|
|
def test_union_type_str_or_none(self) -> None:
|
|
"""str | None should resolve to VARCHAR."""
|
|
ann = str | None
|
|
result = _python_type_to_duckdb(ann)
|
|
assert result == "VARCHAR"
|
|
|
|
def test_union_type_int_or_none(self) -> None:
|
|
ann = int | None
|
|
result = _python_type_to_duckdb(ann)
|
|
assert result == "BIGINT"
|
|
|
|
def test_union_type_date_or_none(self) -> None:
|
|
ann = date | None
|
|
result = _python_type_to_duckdb(ann)
|
|
assert result == "DATE"
|
|
|
|
def test_union_type_float_or_none(self) -> None:
|
|
ann = float | None
|
|
result = _python_type_to_duckdb(ann)
|
|
assert result == "DOUBLE"
|
|
|
|
def test_union_type_decimal_or_none(self) -> None:
|
|
ann = Decimal | None
|
|
result = _python_type_to_duckdb(ann)
|
|
assert result == "DECIMAL(18,2)"
|
|
|
|
|
|
class TestPythonToDuckdbMapping:
|
|
"""_PYTHON_TO_DUCKDB mapping is complete and correct."""
|
|
|
|
def test_mapping_has_all_expected_types(self) -> None:
|
|
expected = {str, int, float, bool, date, datetime, Decimal}
|
|
assert set(_PYTHON_TO_DUCKDB.keys()) == expected
|
|
|
|
def test_mapping_values_are_strings(self) -> None:
|
|
for v in _PYTHON_TO_DUCKDB.values():
|
|
assert isinstance(v, str)
|
|
|
|
|
|
class TestTableRefToViewName:
|
|
"""_table_ref_to_view_name creates valid DuckDB view names."""
|
|
|
|
def test_basic_conversion(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
result = _table_ref_to_view_name("core.encounter", registry)
|
|
assert result == '"_tv_core_0_encounter"'
|
|
|
|
def test_private_table_conversion(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
result = _table_ref_to_view_name("readmissions._int_encounter", registry)
|
|
assert result == ('"_tv_readmissions_0__int_encounter"')
|
|
|
|
def test_registry_populated(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
_table_ref_to_view_name("core.encounter", registry)
|
|
assert "_tv_core_0_encounter" in registry
|
|
assert registry["_tv_core_0_encounter"] == ("core.encounter")
|
|
|
|
def test_registry_private_table(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
_table_ref_to_view_name("readmissions._int_encounter", registry)
|
|
key = "_tv_readmissions_0__int_encounter"
|
|
assert registry[key] == "readmissions._int_encounter"
|
|
|
|
def test_multiple_registrations(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
_table_ref_to_view_name("core.encounter", registry)
|
|
_table_ref_to_view_name("core.patient", registry)
|
|
assert len(registry) == 2
|
|
|
|
def test_view_name_starts_with_prefix(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
result = _table_ref_to_view_name("some.table", registry)
|
|
# Unquoted name should start with _tv_
|
|
inner = result.strip('"')
|
|
assert inner.startswith("_tv_")
|
|
|
|
def test_dot_replaced_with_separator(self) -> None:
|
|
registry: dict[str, str] = {}
|
|
result = _table_ref_to_view_name("schema.table", registry)
|
|
inner = result.strip('"')
|
|
assert "." not in inner
|
|
assert "_0_" in inner
|
|
|
|
def test_roundtrip_via_registry(self) -> None:
|
|
"""Registry must allow lossless reversal."""
|
|
registry: dict[str, str] = {}
|
|
refs = [
|
|
"core.encounter",
|
|
"core._stg_encounter",
|
|
"readmissions._int_encounter",
|
|
"claims_preprocessing.member_months",
|
|
]
|
|
for ref in refs:
|
|
_table_ref_to_view_name(ref, registry)
|
|
|
|
# All original refs recoverable from registry
|
|
recovered = set(registry.values())
|
|
assert recovered == set(refs)
|
|
|
|
|
|
class TestCleanDuckdbAliases:
|
|
"""_clean_duckdb_aliases removes DuckDB internal names."""
|
|
|
|
def test_unnamed_relation_renamed(self) -> None:
|
|
"""unnamed_relation_<hex> aliases become t1, t2, etc."""
|
|
import sqlglot
|
|
|
|
sql = "SELECT * FROM (SELECT 1) AS unnamed_relation_abc123def"
|
|
tree = sqlglot.parse_one(sql, read="duckdb")
|
|
_clean_duckdb_aliases(tree)
|
|
result = tree.sql(dialect="duckdb")
|
|
assert "unnamed_relation" not in result
|
|
assert "t1" in result
|
|
|
|
def test_multiple_unnamed_relations(self) -> None:
|
|
import sqlglot
|
|
|
|
sql = (
|
|
"SELECT a.x, b.y FROM "
|
|
"(SELECT 1 AS x) AS unnamed_relation_aaa "
|
|
"JOIN "
|
|
"(SELECT 2 AS y) AS unnamed_relation_bbb "
|
|
"ON 1 = 1"
|
|
)
|
|
tree = sqlglot.parse_one(sql, read="duckdb")
|
|
_clean_duckdb_aliases(tree)
|
|
result = tree.sql(dialect="duckdb")
|
|
assert "unnamed_relation" not in result
|
|
assert "t1" in result
|
|
assert "t2" in result
|
|
|
|
def test_row_index_renamed(self) -> None:
|
|
"""row_index_<hex> columns become _dedup_row."""
|
|
import sqlglot
|
|
|
|
sql = "SELECT row_index_abc123 FROM my_table"
|
|
tree = sqlglot.parse_one(sql, read="duckdb")
|
|
_clean_duckdb_aliases(tree)
|
|
result = tree.sql(dialect="duckdb")
|
|
assert "row_index_abc123" not in result
|
|
assert "_dedup_row" in result
|
|
|
|
def test_normal_aliases_untouched(self) -> None:
|
|
"""Regular aliases must not be renamed."""
|
|
import sqlglot
|
|
|
|
sql = "SELECT * FROM my_table AS t"
|
|
tree = sqlglot.parse_one(sql, read="duckdb")
|
|
_clean_duckdb_aliases(tree)
|
|
result = tree.sql(dialect="duckdb")
|
|
assert "my_table" in result
|
|
assert " t" in result or "AS t" in result
|
|
|
|
|
|
class TestUnnamedRelationRegex:
|
|
"""The _UNNAMED_RE pattern matches DuckDB internal names."""
|
|
|
|
def test_matches_valid_hex(self) -> None:
|
|
from aco.lake.transpile import _UNNAMED_RE
|
|
|
|
assert _UNNAMED_RE.fullmatch("unnamed_relation_abc123")
|
|
|
|
def test_matches_long_hex(self) -> None:
|
|
from aco.lake.transpile import _UNNAMED_RE
|
|
|
|
assert _UNNAMED_RE.fullmatch("unnamed_relation_0123456789abcdef")
|
|
|
|
def test_no_match_regular_name(self) -> None:
|
|
from aco.lake.transpile import _UNNAMED_RE
|
|
|
|
assert not _UNNAMED_RE.fullmatch("my_table")
|
|
|
|
def test_no_match_partial(self) -> None:
|
|
from aco.lake.transpile import _UNNAMED_RE
|
|
|
|
assert not _UNNAMED_RE.fullmatch("unnamed_relation_")
|
|
|
|
|
|
class TestRowIndexRegex:
|
|
"""The _ROW_INDEX_RE pattern matches DuckDB row index names."""
|
|
|
|
def test_matches_valid_hex(self) -> None:
|
|
from aco.lake.transpile import _ROW_INDEX_RE
|
|
|
|
assert _ROW_INDEX_RE.fullmatch("row_index_abc123")
|
|
|
|
def test_no_match_regular_name(self) -> None:
|
|
from aco.lake.transpile import _ROW_INDEX_RE
|
|
|
|
assert not _ROW_INDEX_RE.fullmatch("row_number")
|
|
|
|
def test_no_match_empty_suffix(self) -> None:
|
|
from aco.lake.transpile import _ROW_INDEX_RE
|
|
|
|
assert not _ROW_INDEX_RE.fullmatch("row_index_")
|
|
|
|
|
|
class TestTranspileSql:
|
|
"""_transpile_sql generates correct target-dialect SQL."""
|
|
|
|
def test_select_mode_returns_bare_select(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="select",
|
|
view_registry={},
|
|
)
|
|
assert "SELECT" in result
|
|
assert "INSERT" not in result
|
|
assert "CREATE" not in result
|
|
|
|
def test_insert_mode_wraps_with_insert(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="insert",
|
|
view_registry={},
|
|
)
|
|
assert result.startswith("INSERT INTO")
|
|
assert "core.encounter" in result
|
|
|
|
def test_ctas_mode_wraps_with_create(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="ctas",
|
|
view_registry={},
|
|
)
|
|
assert "CREATE OR REPLACE TABLE" in result
|
|
|
|
def test_view_mode_wraps_with_view(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="view",
|
|
view_registry={},
|
|
)
|
|
assert "CREATE OR REPLACE VIEW" in result
|
|
|
|
def test_catalog_prefix_added(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="homelab",
|
|
output_mode="select",
|
|
view_registry={},
|
|
)
|
|
assert "homelab" in result
|
|
|
|
def test_catalog_prefix_in_insert_target(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="out.table",
|
|
target_dialect="databricks",
|
|
catalog="homelab",
|
|
output_mode="insert",
|
|
view_registry={},
|
|
)
|
|
assert "homelab.out.table" in result
|
|
|
|
def test_view_registry_rewrites_temp_views(self) -> None:
|
|
"""Temp view names in SQL should be rewritten back."""
|
|
registry = {"_tv_core_0_encounter": "core.encounter"}
|
|
sql = 'SELECT * FROM "_tv_core_0_encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="out.table",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="select",
|
|
view_registry=registry,
|
|
)
|
|
# The rewritten SQL should reference core.encounter
|
|
assert "_tv_" not in result
|
|
|
|
def test_no_catalog_prefix_when_empty(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="select",
|
|
view_registry={},
|
|
)
|
|
# Should not have double-dotted prefix
|
|
assert "..encounter" not in result
|
|
|
|
def test_unknown_output_mode_returns_select(self) -> None:
|
|
"""Fallback for unknown output_mode is bare SELECT."""
|
|
sql = "SELECT 1 AS x"
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="out.table",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="unknown_mode",
|
|
view_registry={},
|
|
)
|
|
assert "SELECT" in result
|
|
assert "INSERT" not in result
|
|
assert "CREATE" not in result
|
|
|
|
|
|
class TestTranspileSqlDialects:
|
|
"""_transpile_sql respects different target dialects."""
|
|
|
|
def test_trino_dialect(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="trino",
|
|
catalog="",
|
|
output_mode="select",
|
|
view_registry={},
|
|
)
|
|
assert "SELECT" in result
|
|
|
|
def test_snowflake_dialect(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="snowflake",
|
|
catalog="",
|
|
output_mode="select",
|
|
view_registry={},
|
|
)
|
|
assert "SELECT" in result
|
|
|
|
def test_databricks_dialect(self) -> None:
|
|
sql = 'SELECT * FROM "core"."encounter"'
|
|
result = _transpile_sql(
|
|
sql,
|
|
expr_name="core.encounter",
|
|
target_dialect="databricks",
|
|
catalog="",
|
|
output_mode="select",
|
|
view_registry={},
|
|
)
|
|
assert "SELECT" in result
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# WriteMode literal type
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestWriteMode:
|
|
"""WriteMode literal accepts only 'append' and 'replace'."""
|
|
|
|
def test_append_is_valid(self) -> None:
|
|
from aco.lake.context import WriteMode
|
|
|
|
mode: WriteMode = "append"
|
|
assert mode == "append"
|
|
|
|
def test_replace_is_valid(self) -> None:
|
|
from aco.lake.context import WriteMode
|
|
|
|
mode: WriteMode = "replace"
|
|
assert mode == "replace"
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Module-level exports (__init__.py)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestLakeExports:
|
|
"""aco.lake exposes the expected public API."""
|
|
|
|
def test_catalog_exported(self) -> None:
|
|
from aco.lake import Catalog
|
|
|
|
assert Catalog is not None
|
|
|
|
def test_context_exported(self) -> None:
|
|
from aco.lake import Context
|
|
|
|
assert Context is not None
|
|
|
|
def test_duckdb_context_exported(self) -> None:
|
|
from aco.lake import DuckDBContext
|
|
|
|
assert DuckDBContext is not None
|
|
|
|
def test_iceberg_context_exported(self) -> None:
|
|
from aco.lake import IcebergContext
|
|
|
|
assert IcebergContext is not None
|
|
|
|
def test_trino_context_exported(self) -> None:
|
|
from aco.lake import TrinoContext
|
|
|
|
assert TrinoContext is not None
|
|
|
|
def test_enterprise_context_exported(self) -> None:
|
|
from aco.lake import EnterpriseContext
|
|
|
|
assert EnterpriseContext is not None
|
|
|
|
def test_parquet_context_exported(self) -> None:
|
|
from aco.lake import ParquetContext
|
|
|
|
assert ParquetContext is not None
|
|
|
|
def test_write_mode_exported(self) -> None:
|
|
from aco.lake import WriteMode
|
|
|
|
assert WriteMode is not None
|
|
|
|
def test_execute_exported(self) -> None:
|
|
from aco.lake import execute
|
|
|
|
assert callable(execute)
|
|
|
|
def test_transpile_exported(self) -> None:
|
|
from aco.lake import transpile
|
|
|
|
assert callable(transpile)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Context hierarchy
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestContextHierarchy:
|
|
"""All concrete contexts inherit from Context."""
|
|
|
|
def test_duckdb_is_context(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
assert isinstance(ctx, Context)
|
|
|
|
def test_parquet_is_context(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
assert isinstance(ctx, Context)
|
|
|
|
def test_iceberg_is_context(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://localhost:19120/")
|
|
assert isinstance(ctx, Context)
|
|
|
|
def test_trino_is_context(self) -> None:
|
|
ctx = TrinoContext()
|
|
assert isinstance(ctx, Context)
|
|
|
|
def test_enterprise_is_context(self) -> None:
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com/api")
|
|
assert isinstance(ctx, Context)
|
|
|
|
def test_all_contexts_are_pydantic_models(self) -> None:
|
|
"""All contexts should be Pydantic BaseModel subclasses."""
|
|
from pydantic import BaseModel
|
|
|
|
for cls in [
|
|
DuckDBContext,
|
|
ParquetContext,
|
|
IcebergContext,
|
|
TrinoContext,
|
|
EnterpriseContext,
|
|
]:
|
|
assert issubclass(cls, BaseModel)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# __init__.py ImportError path (lines 64-65)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestInitImportErrorPath:
|
|
"""The try/except ImportError in __init__.py gracefully skips
|
|
sync/unity imports when their deps are unavailable."""
|
|
|
|
def test_import_error_path_skips_sync_unity(self) -> None:
|
|
"""When databricks.sdk is not importable, sync/unity
|
|
symbols are simply absent from the namespace."""
|
|
import importlib
|
|
import sys
|
|
|
|
# Force the ImportError path by temporarily making
|
|
# aco.lake.sync raise ImportError on next import.
|
|
saved_sync = sys.modules.pop("aco.lake.sync", None)
|
|
saved_unity = sys.modules.pop("aco.lake.unity", None)
|
|
saved_init = sys.modules.pop("aco.lake", None)
|
|
|
|
# Patch builtins.__import__ to fail for aco.lake.sync
|
|
original_import = (
|
|
__builtins__.__import__
|
|
if hasattr(__builtins__, "__import__")
|
|
else __import__
|
|
)
|
|
|
|
def failing_import(name, *args, **kwargs):
|
|
if name in ("aco.lake.sync", "aco.lake.unity"):
|
|
raise ImportError(f"mocked: {name}")
|
|
# For relative imports, check the package argument
|
|
if len(args) >= 3 and args[2] is not None:
|
|
# globals, locals, fromlist
|
|
fromlist = args[2]
|
|
if isinstance(fromlist, (list, tuple)):
|
|
for item in fromlist:
|
|
if item in ("sync", "unity"):
|
|
# Check if this is aco.lake doing a
|
|
# relative import
|
|
pkg = args[0] if args else {}
|
|
pkg_name = (
|
|
pkg.get("__name__", "") if isinstance(pkg, dict) else ""
|
|
)
|
|
if "aco.lake" in str(pkg_name) or name in (
|
|
".sync",
|
|
".unity",
|
|
):
|
|
raise ImportError(f"mocked: {item}")
|
|
return original_import(name, *args, **kwargs)
|
|
|
|
try:
|
|
with patch("builtins.__import__", side_effect=failing_import):
|
|
# Re-import aco.lake — the try/except should catch
|
|
mod = importlib.import_module("aco.lake")
|
|
# Core exports must still work
|
|
assert hasattr(mod, "Catalog")
|
|
assert hasattr(mod, "Context")
|
|
assert hasattr(mod, "DuckDBContext")
|
|
assert hasattr(mod, "execute")
|
|
assert hasattr(mod, "transpile")
|
|
finally:
|
|
# Restore modules
|
|
if saved_sync is not None:
|
|
sys.modules["aco.lake.sync"] = saved_sync
|
|
if saved_unity is not None:
|
|
sys.modules["aco.lake.unity"] = saved_unity
|
|
if saved_init is not None:
|
|
sys.modules["aco.lake"] = saved_init
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# DuckDBContext load/save with mocked DuckDB (lines 203-343)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestDuckDBContextLoadMocked:
|
|
"""DuckDBContext.load with mocked DuckDB connection."""
|
|
|
|
def test_load_returns_narwhals_df(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
mock_con = MagicMock()
|
|
mock_con.execute.return_value.pl.return_value = pl.DataFrame({"x": [1, 2]})
|
|
ctx._connection = mock_con
|
|
|
|
result = ctx.load("core.encounter")
|
|
assert nw.to_native(result).shape == (2, 1)
|
|
mock_con.execute.assert_called_once()
|
|
|
|
def test_load_with_catalog_mapping(self) -> None:
|
|
cat = Catalog(
|
|
schema_map={"core.encounter": "prod.encounters"},
|
|
column_map={"core.encounter": {"encounter_id": "encntr_id"}},
|
|
)
|
|
ctx = DuckDBContext(database=":memory:", catalog=cat)
|
|
mock_con = MagicMock()
|
|
mock_con.execute.return_value.pl.return_value = pl.DataFrame({"encntr_id": [1]})
|
|
ctx._connection = mock_con
|
|
|
|
result = ctx.load("core.encounter")
|
|
# Should rename encntr_id -> encounter_id
|
|
native = nw.to_native(result)
|
|
assert "encounter_id" in native.columns
|
|
|
|
def test_load_with_catalog_no_rename_needed(self) -> None:
|
|
"""When catalog has column_map but no columns match, skip."""
|
|
cat = Catalog(column_map={"core.encounter": {"foo": "bar"}})
|
|
ctx = DuckDBContext(database=":memory:", catalog=cat)
|
|
mock_con = MagicMock()
|
|
mock_con.execute.return_value.pl.return_value = pl.DataFrame({"x": [1]})
|
|
ctx._connection = mock_con
|
|
|
|
result = ctx.load("core.encounter")
|
|
native = nw.to_native(result)
|
|
assert "x" in native.columns
|
|
|
|
def test_load_raises_runtime_error_on_failure(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
mock_con = MagicMock()
|
|
mock_con.execute.side_effect = Exception("table not found")
|
|
ctx._connection = mock_con
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to load"):
|
|
ctx.load("core.encounter")
|
|
|
|
|
|
class TestDuckDBContextSaveMocked:
|
|
"""DuckDBContext.save with mocked DuckDB connection."""
|
|
|
|
def test_save_replace_mode(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:", read_only=False)
|
|
mock_con = MagicMock()
|
|
ctx._connection = mock_con
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
# Should call: CREATE SCHEMA, DROP TABLE, CREATE TABLE
|
|
assert len(mock_con.execute.call_args_list) >= 3
|
|
|
|
def test_save_append_mode_insert(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:", read_only=False)
|
|
mock_con = MagicMock()
|
|
ctx._connection = mock_con
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="append")
|
|
# CREATE SCHEMA + INSERT INTO
|
|
assert mock_con.execute.call_count >= 2
|
|
|
|
def test_save_append_fallback_to_create(self) -> None:
|
|
"""When INSERT fails (table doesn't exist), falls back to CREATE."""
|
|
ctx = DuckDBContext(database=":memory:", read_only=False)
|
|
mock_con = MagicMock()
|
|
|
|
call_count = [0]
|
|
|
|
def side_effect(sql):
|
|
call_count[0] += 1
|
|
if "INSERT INTO" in sql:
|
|
raise Exception("table does not exist")
|
|
return MagicMock()
|
|
|
|
mock_con.execute.side_effect = side_effect
|
|
ctx._connection = mock_con
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="append")
|
|
# Should have: CREATE SCHEMA, INSERT (fails), CREATE TABLE
|
|
assert mock_con.execute.call_count >= 3
|
|
|
|
def test_save_with_catalog_mapping(self) -> None:
|
|
cat = Catalog(
|
|
schema_map={"core.encounter": "prod.encounters"},
|
|
column_map={"core.encounter": {"encounter_id": "encntr_id"}},
|
|
)
|
|
ctx = DuckDBContext(database=":memory:", read_only=False, catalog=cat)
|
|
mock_con = MagicMock()
|
|
ctx._connection = mock_con
|
|
|
|
df = nw.from_native(pl.DataFrame({"encounter_id": [1]}))
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
# Should use physical table name prod.encounters
|
|
sql_calls = [str(c[0][0]) for c in mock_con.execute.call_args_list]
|
|
assert any("prod" in s for s in sql_calls)
|
|
|
|
def test_save_raises_on_outer_exception(self) -> None:
|
|
"""RuntimeError raised when the outer try block fails."""
|
|
ctx = DuckDBContext(database=":memory:", read_only=False)
|
|
mock_con = MagicMock()
|
|
# CREATE SCHEMA itself raises
|
|
mock_con.execute.side_effect = Exception("disk full")
|
|
ctx._connection = mock_con
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with pytest.raises(RuntimeError, match="Failed to save"):
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
|
|
|
|
class TestDuckDBContextConnection:
|
|
"""DuckDBContext._get_connection and __del__."""
|
|
|
|
def test_get_connection_creates_connection(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
with patch("duckdb.connect") as mock_connect:
|
|
mock_connect.return_value = MagicMock()
|
|
con = ctx._get_connection()
|
|
mock_connect.assert_called_once_with(database=":memory:", read_only=True)
|
|
assert con is not None
|
|
|
|
def test_get_connection_caches(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
mock_con = MagicMock()
|
|
ctx._connection = mock_con
|
|
result = ctx._get_connection()
|
|
assert result is mock_con
|
|
|
|
def test_del_closes_connection(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
mock_con = MagicMock()
|
|
ctx._connection = mock_con
|
|
ctx.__del__()
|
|
mock_con.close.assert_called_once()
|
|
|
|
def test_del_handles_close_exception(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
mock_con = MagicMock()
|
|
mock_con.close.side_effect = Exception("already closed")
|
|
ctx._connection = mock_con
|
|
# Should not raise
|
|
ctx.__del__()
|
|
|
|
def test_del_with_no_connection(self) -> None:
|
|
ctx = DuckDBContext(database=":memory:")
|
|
# No connection set — should not raise
|
|
ctx.__del__()
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# ParquetContext load/save with mocked polars (lines 429-516)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestParquetContextLoadMocked:
|
|
"""ParquetContext.load with mocked polars."""
|
|
|
|
def test_load_returns_narwhals_df(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
with patch("polars.read_parquet") as mock_read:
|
|
mock_read.return_value = pl.DataFrame({"x": [1]})
|
|
result = ctx.load("core.encounter")
|
|
native = nw.to_native(result)
|
|
assert native.shape == (1, 1)
|
|
|
|
def test_load_with_catalog_column_mapping(self) -> None:
|
|
cat = Catalog(column_map={"core.enc": {"encounter_id": "enc_id"}})
|
|
ctx = ParquetContext(base_path="./data", catalog=cat)
|
|
with patch("polars.read_parquet") as mock_read:
|
|
mock_read.return_value = pl.DataFrame({"enc_id": [1]})
|
|
result = ctx.load("core.enc")
|
|
native = nw.to_native(result)
|
|
assert "encounter_id" in native.columns
|
|
|
|
def test_load_raises_on_failure(self) -> None:
|
|
ctx = ParquetContext(base_path="./data")
|
|
with patch("polars.read_parquet") as mock_read:
|
|
mock_read.side_effect = FileNotFoundError("nope")
|
|
with pytest.raises(RuntimeError, match="Failed to load"):
|
|
ctx.load("core.encounter")
|
|
|
|
|
|
class TestParquetContextSaveMocked:
|
|
"""ParquetContext.save with mocked polars."""
|
|
|
|
def test_save_replace_mode(self) -> None:
|
|
ctx = ParquetContext(base_path="/tmp/test_parquet_ctx")
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with (
|
|
patch("polars.read_parquet") as mock_read,
|
|
patch.object(pl.DataFrame, "write_parquet") as mock_write,
|
|
patch("pathlib.Path.mkdir"),
|
|
):
|
|
mock_read.side_effect = FileNotFoundError("no file")
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
mock_write.assert_called_once()
|
|
|
|
def test_save_append_mode_concat(self) -> None:
|
|
ctx = ParquetContext(base_path="/tmp/test_parquet_ctx")
|
|
df = nw.from_native(pl.DataFrame({"x": [2]}))
|
|
with (
|
|
patch("polars.read_parquet") as mock_read,
|
|
patch.object(pl.DataFrame, "write_parquet") as mock_write,
|
|
patch("pathlib.Path.mkdir"),
|
|
):
|
|
mock_read.return_value = pl.DataFrame({"x": [1]})
|
|
ctx.save("core.encounter", df, mode="append")
|
|
mock_write.assert_called_once()
|
|
|
|
def test_save_append_first_write(self) -> None:
|
|
"""First write in append mode — no existing file."""
|
|
ctx = ParquetContext(base_path="/tmp/test_parquet_ctx")
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with (
|
|
patch("polars.read_parquet") as mock_read,
|
|
patch.object(pl.DataFrame, "write_parquet") as mock_write,
|
|
patch("pathlib.Path.mkdir"),
|
|
):
|
|
mock_read.side_effect = FileNotFoundError("first write")
|
|
ctx.save("core.encounter", df, mode="append")
|
|
mock_write.assert_called_once()
|
|
|
|
def test_save_with_catalog_column_mapping(self) -> None:
|
|
cat = Catalog(column_map={"core.enc": {"encounter_id": "enc_id"}})
|
|
ctx = ParquetContext(base_path="/tmp/test_pq", catalog=cat)
|
|
df = nw.from_native(pl.DataFrame({"encounter_id": [1]}))
|
|
with (
|
|
patch("polars.read_parquet") as mock_read,
|
|
patch.object(pl.DataFrame, "write_parquet"),
|
|
patch("pathlib.Path.mkdir"),
|
|
):
|
|
mock_read.side_effect = FileNotFoundError("no")
|
|
ctx.save("core.enc", df, mode="replace")
|
|
|
|
def test_save_s3_path_no_mkdir(self) -> None:
|
|
"""S3 paths should not create local dirs."""
|
|
ctx = ParquetContext(base_path="s3://bucket/lake")
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with (
|
|
patch("polars.read_parquet") as mock_read,
|
|
patch.object(pl.DataFrame, "write_parquet"),
|
|
patch("pathlib.Path.mkdir") as mock_mkdir,
|
|
):
|
|
mock_read.side_effect = FileNotFoundError("no")
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
mock_mkdir.assert_not_called()
|
|
|
|
def test_save_raises_on_write_failure(self) -> None:
|
|
ctx = ParquetContext(base_path="/tmp/test_pq")
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with (
|
|
patch("polars.read_parquet") as mock_read,
|
|
patch.object(
|
|
pl.DataFrame, "write_parquet", side_effect=Exception("write fail")
|
|
),
|
|
patch("pathlib.Path.mkdir"),
|
|
):
|
|
mock_read.side_effect = FileNotFoundError("no")
|
|
with pytest.raises(RuntimeError, match="Failed to save"):
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# IcebergContext load/save (lines 591-747)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestIcebergContextLoadMocked:
|
|
"""IcebergContext.load with mocked PyIceberg."""
|
|
|
|
def test_load_returns_narwhals_df(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_table = MagicMock()
|
|
import pyarrow as pa
|
|
|
|
arrow_tbl = pa.table({"x": [1, 2]})
|
|
mock_table.scan.return_value.to_arrow.return_value = arrow_tbl
|
|
mock_catalog.load_table.return_value = mock_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
result = ctx.load("core.encounter")
|
|
native = nw.to_native(result)
|
|
assert native.shape == (2, 1)
|
|
|
|
def test_load_with_catalog_column_mapping(self) -> None:
|
|
cat = Catalog(column_map={"core.enc": {"eid": "enc_id"}})
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/", catalog=cat)
|
|
mock_catalog = MagicMock()
|
|
mock_table = MagicMock()
|
|
import pyarrow as pa
|
|
|
|
arrow_tbl = pa.table({"enc_id": [1]})
|
|
mock_table.scan.return_value.to_arrow.return_value = arrow_tbl
|
|
mock_catalog.load_table.return_value = mock_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
result = ctx.load("core.enc")
|
|
native = nw.to_native(result)
|
|
assert "eid" in native.columns
|
|
|
|
def test_load_with_physical_table_mapping(self) -> None:
|
|
cat = Catalog(
|
|
schema_map={"core.encounter": "prod.encounters"},
|
|
)
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/", catalog=cat)
|
|
mock_catalog = MagicMock()
|
|
mock_table = MagicMock()
|
|
import pyarrow as pa
|
|
|
|
arrow_tbl = pa.table({"x": [1]})
|
|
mock_table.scan.return_value.to_arrow.return_value = arrow_tbl
|
|
mock_catalog.load_table.return_value = mock_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
ctx.load("core.encounter")
|
|
# Should call with physical table ref parts
|
|
mock_catalog.load_table.assert_called_once_with(("prod", "encounters"))
|
|
|
|
def test_load_raises_runtime_error(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_catalog.load_table.side_effect = Exception("not found")
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to load"):
|
|
ctx.load("core.encounter")
|
|
|
|
|
|
class TestIcebergContextSaveMocked:
|
|
"""IcebergContext.save with mocked PyIceberg."""
|
|
|
|
def test_save_append_existing_table(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_catalog.load_table.return_value = mock_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="append")
|
|
mock_table.append.assert_called_once()
|
|
|
|
def test_save_replace_existing_table(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_catalog.load_table.return_value = mock_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="replace")
|
|
mock_table.overwrite.assert_called_once()
|
|
|
|
def test_save_creates_new_table(self) -> None:
|
|
"""When table doesn't exist, create namespace + table."""
|
|
import sys
|
|
|
|
# Inject mock pyiceberg modules since they're not installed
|
|
mock_schema_mod = MagicMock()
|
|
mock_types_mod = MagicMock()
|
|
saved = {k: sys.modules.get(k) for k in ["pyiceberg.schema", "pyiceberg.types"]}
|
|
sys.modules["pyiceberg.schema"] = mock_schema_mod
|
|
sys.modules["pyiceberg.types"] = mock_types_mod
|
|
|
|
try:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_catalog.load_table.side_effect = Exception("not found")
|
|
mock_new_table = MagicMock()
|
|
mock_catalog.create_table.return_value = mock_new_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="append")
|
|
mock_catalog.create_namespace.assert_called_once()
|
|
mock_catalog.create_table.assert_called_once()
|
|
mock_new_table.append.assert_called_once()
|
|
finally:
|
|
for k, v in saved.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
def test_save_creates_table_namespace_exists(self) -> None:
|
|
"""When namespace already exists, swallow the error."""
|
|
import sys
|
|
|
|
mock_schema_mod = MagicMock()
|
|
mock_types_mod = MagicMock()
|
|
saved = {k: sys.modules.get(k) for k in ["pyiceberg.schema", "pyiceberg.types"]}
|
|
sys.modules["pyiceberg.schema"] = mock_schema_mod
|
|
sys.modules["pyiceberg.types"] = mock_types_mod
|
|
|
|
try:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_catalog.load_table.side_effect = Exception("not found")
|
|
mock_catalog.create_namespace.side_effect = Exception("exists")
|
|
mock_new_table = MagicMock()
|
|
mock_catalog.create_table.return_value = mock_new_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx.save("core.encounter", df, mode="append")
|
|
# Should still proceed
|
|
mock_catalog.create_table.assert_called_once()
|
|
finally:
|
|
for k, v in saved.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
def test_save_with_catalog_mapping(self) -> None:
|
|
cat = Catalog(
|
|
schema_map={"core.encounter": "prod.encounters"},
|
|
column_map={"core.encounter": {"encounter_id": "enc_id"}},
|
|
)
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/", catalog=cat)
|
|
mock_catalog = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_catalog.load_table.return_value = mock_table
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
df = nw.from_native(pl.DataFrame({"encounter_id": [1]}))
|
|
ctx.save("core.encounter", df, mode="append")
|
|
# Should use physical table
|
|
mock_catalog.load_table.assert_called_with(("prod", "encounters"))
|
|
|
|
def test_save_raises_runtime_error(self) -> None:
|
|
"""Outer exception handling wraps in RuntimeError."""
|
|
import sys
|
|
|
|
mock_schema_mod = MagicMock()
|
|
mock_types_mod = MagicMock()
|
|
saved = {k: sys.modules.get(k) for k in ["pyiceberg.schema", "pyiceberg.types"]}
|
|
sys.modules["pyiceberg.schema"] = mock_schema_mod
|
|
sys.modules["pyiceberg.types"] = mock_types_mod
|
|
|
|
try:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_catalog = MagicMock()
|
|
mock_catalog.load_table.side_effect = Exception("not found")
|
|
mock_catalog.create_table.side_effect = Exception("perm denied")
|
|
ctx._pyiceberg_catalog = mock_catalog
|
|
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with pytest.raises(RuntimeError, match="Failed to save"):
|
|
ctx.save("core.encounter", df, mode="append")
|
|
finally:
|
|
for k, v in saved.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
|
|
class TestIcebergContextGetCatalog:
|
|
"""IcebergContext._get_pyiceberg_catalog."""
|
|
|
|
def _inject_pyiceberg(self):
|
|
"""Inject mock pyiceberg.catalog into sys.modules."""
|
|
import sys
|
|
|
|
mock_catalog_mod = MagicMock()
|
|
saved = {}
|
|
for k in ["pyiceberg", "pyiceberg.catalog"]:
|
|
saved[k] = sys.modules.get(k)
|
|
sys.modules.setdefault("pyiceberg", MagicMock())
|
|
sys.modules["pyiceberg.catalog"] = mock_catalog_mod
|
|
return mock_catalog_mod, saved
|
|
|
|
def _restore_pyiceberg(self, saved):
|
|
import sys
|
|
|
|
for k, v in saved.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
def test_creates_catalog_on_first_call(self) -> None:
|
|
mock_mod, saved = self._inject_pyiceberg()
|
|
try:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_mod.load_catalog.return_value = MagicMock()
|
|
result = ctx._get_pyiceberg_catalog()
|
|
mock_mod.load_catalog.assert_called_once()
|
|
assert result is not None
|
|
finally:
|
|
self._restore_pyiceberg(saved)
|
|
|
|
def test_caches_catalog(self) -> None:
|
|
ctx = IcebergContext(catalog_uri="http://nessie:19120/")
|
|
mock_cat = MagicMock()
|
|
ctx._pyiceberg_catalog = mock_cat
|
|
result = ctx._get_pyiceberg_catalog()
|
|
assert result is mock_cat
|
|
|
|
def test_custom_ref_added_to_config(self) -> None:
|
|
mock_mod, saved = self._inject_pyiceberg()
|
|
try:
|
|
ctx = IcebergContext(
|
|
catalog_uri="http://nessie:19120/",
|
|
ref="develop",
|
|
)
|
|
mock_mod.load_catalog.return_value = MagicMock()
|
|
ctx._get_pyiceberg_catalog()
|
|
call_kwargs = mock_mod.load_catalog.call_args[1]
|
|
assert call_kwargs["ref"] == "develop"
|
|
finally:
|
|
self._restore_pyiceberg(saved)
|
|
|
|
def test_main_ref_not_in_config(self) -> None:
|
|
mock_mod, saved = self._inject_pyiceberg()
|
|
try:
|
|
ctx = IcebergContext(
|
|
catalog_uri="http://nessie:19120/",
|
|
ref="main",
|
|
)
|
|
mock_mod.load_catalog.return_value = MagicMock()
|
|
ctx._get_pyiceberg_catalog()
|
|
call_kwargs = mock_mod.load_catalog.call_args[1]
|
|
assert "ref" not in call_kwargs
|
|
finally:
|
|
self._restore_pyiceberg(saved)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Engine: _execute_transpiled (lines 241-264)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestExecuteTranspiled:
|
|
"""_execute_transpiled calls transpile with correct parameters."""
|
|
|
|
def test_enterprise_context_uses_dialect(self) -> None:
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_transpiled
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = EnterpriseContext(
|
|
catalog_uri="https://example.com",
|
|
dialect="snowflake",
|
|
warehouse="main",
|
|
)
|
|
|
|
with patch("aco.lake.transpile.transpile") as mock_tp:
|
|
mock_tp.return_value = {"out.table": "SELECT 1"}
|
|
result = _execute_transpiled(pipeline, ctx)
|
|
mock_tp.assert_called_once()
|
|
kwargs = mock_tp.call_args
|
|
assert kwargs[1]["target_dialect"] == "snowflake"
|
|
assert kwargs[1]["catalog"] == "main"
|
|
assert "out.table" in result
|
|
|
|
def test_trino_context_uses_trino_dialect(self) -> None:
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_transpiled
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = TrinoContext(catalog="iceberg")
|
|
|
|
with patch("aco.lake.transpile.transpile") as mock_tp:
|
|
mock_tp.return_value = {"out.table": "SELECT 1"}
|
|
_execute_transpiled(pipeline, ctx)
|
|
kwargs = mock_tp.call_args
|
|
assert kwargs[1]["target_dialect"] == "trino"
|
|
assert kwargs[1]["catalog"] == "iceberg"
|
|
|
|
def test_enterprise_no_dialect_defaults_databricks(self) -> None:
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_transpiled
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = EnterpriseContext(catalog_uri="https://example.com")
|
|
|
|
with patch("aco.lake.transpile.transpile") as mock_tp:
|
|
mock_tp.return_value = {}
|
|
_execute_transpiled(pipeline, ctx)
|
|
kwargs = mock_tp.call_args
|
|
assert kwargs[1]["target_dialect"] == "databricks"
|
|
|
|
def test_replace_mode_uses_ctas(self) -> None:
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_transpiled
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = EnterpriseContext(
|
|
catalog_uri="https://example.com",
|
|
dialect="databricks",
|
|
)
|
|
|
|
with patch("aco.lake.transpile.transpile") as mock_tp:
|
|
mock_tp.return_value = {}
|
|
_execute_transpiled(pipeline, ctx, mode="replace")
|
|
kwargs = mock_tp.call_args
|
|
assert kwargs[1]["output_mode"] == "ctas"
|
|
|
|
|
|
class TestExecuteDirectSaveFailure:
|
|
"""_execute_direct handles save failures gracefully."""
|
|
|
|
def test_save_failure_prints_warning(self) -> None:
|
|
from aco.express.base import Expr
|
|
from aco.lake.engine import _execute_direct
|
|
|
|
source = pl.DataFrame({"x": [1]})
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="out.table", fn=fn)])
|
|
ctx = MagicMock(spec=DuckDBContext)
|
|
ctx.load = MagicMock(return_value=source)
|
|
ctx.save = MagicMock(side_effect=Exception("save failed"))
|
|
|
|
# Should not raise — just prints warning
|
|
results = _execute_direct(pipeline, ctx, save_outputs=True)
|
|
assert "out.table" in results
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Transpile: transpile() full entry point (lines 157-224)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestTranspileFullEntry:
|
|
"""transpile() generates SQL for pipeline expressions."""
|
|
|
|
def test_transpile_with_explicit_connection(self) -> None:
|
|
import duckdb
|
|
|
|
from aco.express.base import Expr
|
|
from aco.lake.transpile import transpile
|
|
|
|
con = duckdb.connect(":memory:")
|
|
con.execute('CREATE SCHEMA "test"')
|
|
con.execute('CREATE TABLE "test"."src" (x BIGINT, y VARCHAR)')
|
|
|
|
@nw.narwhalify
|
|
def my_expr(test__src):
|
|
return test__src
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="test.out", fn=my_expr)])
|
|
result = transpile(
|
|
pipeline,
|
|
con,
|
|
target_dialect="databricks",
|
|
catalog="homelab",
|
|
output_mode="insert",
|
|
)
|
|
con.close()
|
|
assert "test.out" in result
|
|
assert "INSERT INTO" in result["test.out"]
|
|
assert "homelab" in result["test.out"]
|
|
|
|
def test_transpile_error_in_expr_fn(self) -> None:
|
|
"""When an expression function raises, record error."""
|
|
import duckdb
|
|
|
|
from aco.express.base import Expr
|
|
from aco.lake.transpile import transpile
|
|
|
|
con = duckdb.connect(":memory:")
|
|
con.execute('CREATE SCHEMA "test"')
|
|
con.execute('CREATE TABLE "test"."src" (x BIGINT)')
|
|
|
|
def bad_fn(test__src):
|
|
raise ValueError("intentional error")
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="test.out", fn=bad_fn)])
|
|
result = transpile(pipeline, con, target_dialect="databricks")
|
|
con.close()
|
|
assert "test.out" in result
|
|
assert "ERROR" in result["test.out"]
|
|
|
|
def test_transpile_error_fallback_loads_table(self) -> None:
|
|
"""When expr fails but the table exists, cache it for
|
|
downstream."""
|
|
import duckdb
|
|
|
|
from aco.express.base import Expr
|
|
from aco.lake.transpile import transpile
|
|
|
|
con = duckdb.connect(":memory:")
|
|
con.execute('CREATE SCHEMA "test"')
|
|
con.execute('CREATE TABLE "test"."out" (x BIGINT)')
|
|
con.execute('CREATE TABLE "test"."src" (x BIGINT)')
|
|
|
|
def bad_fn(test__src):
|
|
raise ValueError("fail")
|
|
|
|
@nw.narwhalify
|
|
def downstream(test__out):
|
|
return test__out
|
|
|
|
pipeline = Pipeline(
|
|
exprs=[
|
|
Expr(name="test.out", fn=bad_fn),
|
|
Expr(name="test.final", fn=downstream),
|
|
]
|
|
)
|
|
result = transpile(pipeline, con, target_dialect="databricks")
|
|
con.close()
|
|
assert "ERROR" in result["test.out"]
|
|
# downstream should still work using the fallback table
|
|
assert "test.final" in result
|
|
|
|
def test_transpile_auto_creates_connection(self) -> None:
|
|
"""When con=None, _schema_only_connection is called."""
|
|
from aco.express.base import Expr
|
|
from aco.lake.transpile import transpile
|
|
|
|
@nw.narwhalify
|
|
def fn(core__encounter):
|
|
return core__encounter
|
|
|
|
pipeline = Pipeline(exprs=[Expr(name="test.out", fn=fn)])
|
|
|
|
with patch("aco.lake.transpile._schema_only_connection") as mock_soc:
|
|
mock_con = MagicMock()
|
|
mock_rel = MagicMock()
|
|
mock_rel.sql_query.return_value = "SELECT * FROM t"
|
|
mock_con.sql.return_value = mock_rel
|
|
mock_soc.return_value = mock_con
|
|
transpile(pipeline, target_dialect="databricks")
|
|
mock_soc.assert_called_once()
|
|
mock_con.close.assert_called_once()
|
|
|
|
def test_load_relation(self) -> None:
|
|
"""_load_relation creates a DuckDB relation."""
|
|
import duckdb
|
|
|
|
from aco.lake.transpile import _load_relation
|
|
|
|
con = duckdb.connect(":memory:")
|
|
con.execute('CREATE SCHEMA "core"')
|
|
con.execute('CREATE TABLE "core"."encounter" (x BIGINT)')
|
|
result = _load_relation(con, "core.encounter")
|
|
assert result is not None
|
|
con.close()
|
|
|
|
|
|
class TestSchemaOnlyConnection:
|
|
"""_schema_only_connection creates DDL from aco.table models."""
|
|
|
|
def test_creates_connection_with_schemas(self) -> None:
|
|
from aco.lake.transpile import _schema_only_connection
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["test"]
|
|
mock_cat.tables.return_value = ["test.t1"]
|
|
|
|
# Create a mock model with fields
|
|
mock_model = MagicMock()
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
"name": FieldInfo(annotation=str),
|
|
}
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
con = _schema_only_connection()
|
|
# Should have created the schema and table
|
|
# Verify by querying
|
|
result = con.execute(
|
|
"SELECT * FROM information_schema.tables WHERE table_schema = 'test'"
|
|
).fetchall()
|
|
assert len(result) > 0
|
|
con.close()
|
|
|
|
def test_skips_empty_model(self) -> None:
|
|
from aco.lake.transpile import _schema_only_connection
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["test"]
|
|
mock_cat.tables.return_value = ["test.empty"]
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {}
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
con = _schema_only_connection()
|
|
# Table should not be created (no columns)
|
|
result = con.execute(
|
|
"SELECT * FROM information_schema.tables WHERE table_name = 'empty'"
|
|
).fetchall()
|
|
assert len(result) == 0
|
|
con.close()
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Catalog Iceberg methods with mocks (lines 341-535)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestCatalogIcebergMethods:
|
|
"""Catalog Iceberg methods with mocked PyIceberg catalog."""
|
|
|
|
def test_get_iceberg_catalog_loads_on_first_call(self) -> None:
|
|
import sys
|
|
|
|
mock_catalog_mod = MagicMock()
|
|
saved = {}
|
|
for k in ["pyiceberg", "pyiceberg.catalog"]:
|
|
saved[k] = sys.modules.get(k)
|
|
sys.modules.setdefault("pyiceberg", MagicMock())
|
|
sys.modules["pyiceberg.catalog"] = mock_catalog_mod
|
|
|
|
try:
|
|
cat = Catalog(
|
|
catalog_uri="http://nessie:19120/",
|
|
warehouse="s3://lake/",
|
|
)
|
|
mock_catalog_mod.load_catalog.return_value = MagicMock()
|
|
result = cat._get_iceberg_catalog()
|
|
mock_catalog_mod.load_catalog.assert_called_once_with(
|
|
"default",
|
|
uri="http://nessie:19120/",
|
|
warehouse="s3://lake/",
|
|
)
|
|
assert result is not None
|
|
finally:
|
|
for k, v in saved.items():
|
|
if v is None:
|
|
sys.modules.pop(k, None)
|
|
else:
|
|
sys.modules[k] = v
|
|
|
|
def test_get_iceberg_catalog_caches(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_cat = MagicMock()
|
|
cat._iceberg_catalog = mock_cat
|
|
result = cat._get_iceberg_catalog()
|
|
assert result is mock_cat
|
|
|
|
def test_iceberg_namespaces(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
mock_ice.list_namespaces.return_value = [
|
|
("core",),
|
|
("readmissions",),
|
|
]
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
result = cat.iceberg_namespaces()
|
|
assert result == ["core", "readmissions"]
|
|
|
|
def test_iceberg_tables(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
mock_ice.list_tables.return_value = [
|
|
("core", "encounter"),
|
|
("core", "patient"),
|
|
]
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
result = cat.iceberg_tables("core")
|
|
assert result == ["core.encounter", "core.patient"]
|
|
|
|
def test_iceberg_schema(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_schema = MagicMock()
|
|
mock_table.schema.return_value = mock_schema
|
|
mock_ice.load_table.return_value = mock_table
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
result = cat.iceberg_schema("core.encounter")
|
|
assert result is mock_schema
|
|
mock_ice.load_table.assert_called_with(("core", "encounter"))
|
|
|
|
def test_iceberg_schema_with_mapping(self) -> None:
|
|
cat = Catalog(
|
|
catalog_uri="http://nessie:19120/",
|
|
schema_map={"core.encounter": "prod.encounters"},
|
|
)
|
|
mock_ice = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_ice.load_table.return_value = mock_table
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
cat.iceberg_schema("core.encounter")
|
|
mock_ice.load_table.assert_called_with(("prod", "encounters"))
|
|
|
|
def test_iceberg_schema_unqualified_raises(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
with pytest.raises(ValueError, match="qualified"):
|
|
cat.iceberg_schema("encounter")
|
|
|
|
def test_iceberg_snapshots(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_table.snapshots.return_value = ["snap1", "snap2"]
|
|
mock_ice.load_table.return_value = mock_table
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
result = cat.iceberg_snapshots("core.encounter")
|
|
assert result == ["snap1", "snap2"]
|
|
|
|
def test_iceberg_current_snapshot(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
mock_table = MagicMock()
|
|
mock_snap = MagicMock()
|
|
mock_table.current_snapshot.return_value = mock_snap
|
|
mock_ice.load_table.return_value = mock_table
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
result = cat.iceberg_current_snapshot("core.encounter")
|
|
assert result is mock_snap
|
|
|
|
|
|
class TestCatalogValidate:
|
|
"""Catalog.validate compares schema definitions vs Iceberg."""
|
|
|
|
def test_validate_finds_mismatches(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
# tables() returns schema tables
|
|
with (
|
|
patch.object(cat, "tables") as mock_tables,
|
|
patch.object(cat, "iceberg_tables") as mock_ice_tables,
|
|
patch.object(cat, "columns") as mock_cols,
|
|
patch.object(cat, "iceberg_schema") as mock_ice_schema,
|
|
):
|
|
mock_tables.return_value = ["test.a", "test.b"]
|
|
mock_ice_tables.return_value = ["test.b", "test.c"]
|
|
# For common table "test.b"
|
|
mock_cols.return_value = ["col1", "col2"]
|
|
mock_ice_field_1 = MagicMock()
|
|
mock_ice_field_1.name = "col2"
|
|
mock_ice_field_2 = MagicMock()
|
|
mock_ice_field_2.name = "col3"
|
|
mock_ice_schema.return_value.fields = [
|
|
mock_ice_field_1,
|
|
mock_ice_field_2,
|
|
]
|
|
|
|
result = cat.validate("test")
|
|
assert result["schema"] == "test"
|
|
assert "test.a" in result["missing_in_iceberg"]
|
|
assert "test.c" in result["missing_in_schema"]
|
|
assert "test.b" in result["column_mismatches"]
|
|
mismatches = result["column_mismatches"]["test.b"]
|
|
assert "col1" in mismatches["missing_in_iceberg"]
|
|
assert "col3" in mismatches["missing_in_schema"]
|
|
|
|
def test_validate_column_error(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
with (
|
|
patch.object(cat, "tables") as mock_tables,
|
|
patch.object(cat, "iceberg_tables") as mock_ice_tables,
|
|
patch.object(cat, "columns") as mock_cols,
|
|
patch.object(cat, "iceberg_schema") as mock_ice_schema,
|
|
):
|
|
mock_tables.return_value = ["test.a"]
|
|
mock_ice_tables.return_value = ["test.a"]
|
|
mock_cols.return_value = ["col1"]
|
|
mock_ice_schema.side_effect = Exception("timeout")
|
|
|
|
result = cat.validate("test")
|
|
assert "test.a" in result["column_mismatches"]
|
|
assert "error" in result["column_mismatches"]["test.a"]
|
|
|
|
def test_validate_no_mismatches(self) -> None:
|
|
cat = Catalog(catalog_uri="http://nessie:19120/")
|
|
mock_ice = MagicMock()
|
|
cat._iceberg_catalog = mock_ice
|
|
|
|
with (
|
|
patch.object(cat, "tables") as mock_tables,
|
|
patch.object(cat, "iceberg_tables") as mock_ice_tables,
|
|
patch.object(cat, "columns") as mock_cols,
|
|
patch.object(cat, "iceberg_schema") as mock_ice_schema,
|
|
):
|
|
mock_tables.return_value = ["test.a"]
|
|
mock_ice_tables.return_value = ["test.a"]
|
|
mock_cols.return_value = ["col1"]
|
|
mock_field = MagicMock()
|
|
mock_field.name = "col1"
|
|
mock_ice_schema.return_value.fields = [mock_field]
|
|
|
|
result = cat.validate("test")
|
|
assert result["column_mismatches"] == {}
|
|
|
|
def test_columns_delegates_to_model(self) -> None:
|
|
cat = Catalog()
|
|
mock_model = MagicMock()
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model.model_fields = {
|
|
"b_col": FieldInfo(annotation=str),
|
|
"a_col": FieldInfo(annotation=int),
|
|
}
|
|
cat._model_cache["test.t"] = mock_model
|
|
|
|
result = cat.columns("test.t")
|
|
assert result == ["a_col", "b_col"]
|
|
|
|
def test_discover_table_modules(self) -> None:
|
|
"""_discover_table_modules returns (ns, module_name) pairs."""
|
|
cat = Catalog()
|
|
result = cat._discover_table_modules()
|
|
assert isinstance(result, list)
|
|
# At minimum aco.table should produce some modules
|
|
assert len(result) > 0
|
|
assert all(isinstance(r, tuple) and len(r) == 2 for r in result)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Sync module (lines 114-425)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestSyncResolveTable:
|
|
"""_resolve_tables builds table list from source."""
|
|
|
|
def test_explicit_tables_returned_sorted(self) -> None:
|
|
from aco.lake.sync import _resolve_tables
|
|
|
|
result = _resolve_tables(MagicMock(), tables=["b.t", "a.t"], schemas=None)
|
|
assert result == ["a.t", "b.t"]
|
|
|
|
def test_auto_discovery_from_duckdb(self) -> None:
|
|
from aco.lake.sync import _resolve_tables
|
|
|
|
mock_ctx = MagicMock(spec=DuckDBContext)
|
|
mock_con = MagicMock()
|
|
mock_con.execute.return_value.fetchall.return_value = [
|
|
("core.encounter",),
|
|
("core.patient",),
|
|
]
|
|
mock_ctx._get_connection.return_value = mock_con
|
|
|
|
result = _resolve_tables(mock_ctx, tables=None, schemas=None)
|
|
assert result == ["core.encounter", "core.patient"]
|
|
|
|
def test_auto_discovery_with_schema_filter(self) -> None:
|
|
from aco.lake.sync import _resolve_tables
|
|
|
|
mock_ctx = MagicMock(spec=DuckDBContext)
|
|
mock_con = MagicMock()
|
|
mock_con.execute.return_value.fetchall.return_value = [
|
|
("core.encounter",),
|
|
("readmissions.enc",),
|
|
]
|
|
mock_ctx._get_connection.return_value = mock_con
|
|
|
|
result = _resolve_tables(mock_ctx, tables=None, schemas=["core"])
|
|
assert result == ["core.encounter"]
|
|
|
|
def test_auto_discovery_non_duckdb_raises(self) -> None:
|
|
from aco.lake.sync import _resolve_tables
|
|
|
|
mock_ctx = MagicMock(spec=ParquetContext)
|
|
with pytest.raises(ValueError, match="DuckDBContext"):
|
|
_resolve_tables(mock_ctx, tables=None, schemas=None)
|
|
|
|
|
|
class TestSyncTableGeneric:
|
|
"""_sync_table_generic transfers one table."""
|
|
|
|
def test_copies_table(self) -> None:
|
|
from aco.lake.sync import _sync_table_generic
|
|
|
|
src = MagicMock()
|
|
dst = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1, 2]}))
|
|
src.load.return_value = df
|
|
|
|
ev = _sync_table_generic(src, dst, "core.encounter")
|
|
assert ev.status == "copied"
|
|
assert ev.rows == 2
|
|
dst.save.assert_called_once()
|
|
|
|
def test_skips_empty_table(self) -> None:
|
|
from aco.lake.sync import _sync_table_generic
|
|
|
|
src = MagicMock()
|
|
dst = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": pl.Series([], dtype=pl.Int64)}))
|
|
src.load.return_value = df
|
|
|
|
ev = _sync_table_generic(src, dst, "core.encounter", skip_empty=True)
|
|
assert ev.status == "skipped"
|
|
|
|
def test_load_failure(self) -> None:
|
|
from aco.lake.sync import _sync_table_generic
|
|
|
|
src = MagicMock()
|
|
dst = MagicMock()
|
|
src.load.side_effect = Exception("no table")
|
|
|
|
ev = _sync_table_generic(src, dst, "core.encounter")
|
|
assert ev.status == "failed"
|
|
assert "load" in ev.error
|
|
|
|
def test_save_failure(self) -> None:
|
|
from aco.lake.sync import _sync_table_generic
|
|
|
|
src = MagicMock()
|
|
dst = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
dst.save.side_effect = Exception("write error")
|
|
|
|
ev = _sync_table_generic(src, dst, "core.encounter")
|
|
assert ev.status == "failed"
|
|
assert "save" in ev.error
|
|
|
|
|
|
class TestSyncTableDatabricks:
|
|
"""_sync_table_databricks with mocked workspace client."""
|
|
|
|
def test_copies_successfully(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1, 2]}))
|
|
src.load.return_value = df
|
|
|
|
ws = MagicMock()
|
|
resp_status = MagicMock()
|
|
resp_status.state = "SUCCEEDED"
|
|
resp_status.error = None
|
|
resp = MagicMock()
|
|
resp.status = resp_status
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "copied"
|
|
assert ev.rows == 2
|
|
ws.files.upload.assert_called_once()
|
|
|
|
def test_skips_empty(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": pl.Series([], dtype=pl.Int64)}))
|
|
src.load.return_value = df
|
|
|
|
ws = MagicMock()
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "skipped"
|
|
|
|
def test_load_failure(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
src.load.side_effect = Exception("load error")
|
|
ws = MagicMock()
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "failed"
|
|
assert "load" in ev.error
|
|
|
|
def test_upload_failure(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
ws = MagicMock()
|
|
ws.files.upload.side_effect = Exception("upload error")
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "failed"
|
|
assert "upload" in ev.error
|
|
|
|
def test_copy_into_failure(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
ws = MagicMock()
|
|
ws.statement_execution.execute_statement.side_effect = Exception("sql error")
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "failed"
|
|
assert "COPY INTO" in ev.error
|
|
|
|
def test_copy_into_non_succeeded_state(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
ws = MagicMock()
|
|
resp = MagicMock()
|
|
resp.status.state = "FAILED"
|
|
resp.status.error.message = "bad schema"
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "failed"
|
|
assert "COPY INTO FAILED" in ev.error
|
|
|
|
def test_copy_into_non_succeeded_no_error_msg(self) -> None:
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
ws = MagicMock()
|
|
resp = MagicMock()
|
|
resp.status.state = "FAILED"
|
|
resp.status.error = None
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "failed"
|
|
|
|
def test_underscore_table_name_prefix(self) -> None:
|
|
"""Tables starting with _ get data_ prefix in volume path."""
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
ws = MagicMock()
|
|
resp = MagicMock()
|
|
resp.status.state = "SUCCEEDED"
|
|
resp.status.error = None
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"schema._hidden",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "copied"
|
|
upload_call = ws.files.upload.call_args
|
|
vol_path = upload_call[0][0]
|
|
assert "data__hidden" in vol_path
|
|
|
|
def test_cleanup_failure_ignored(self) -> None:
|
|
"""Cleanup delete failure should not affect result."""
|
|
from aco.lake.sync import _sync_table_databricks
|
|
|
|
src = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
ws = MagicMock()
|
|
resp = MagicMock()
|
|
resp.status.state = "SUCCEEDED"
|
|
resp.status.error = None
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
ws.files.delete.side_effect = Exception("delete error")
|
|
|
|
ev = _sync_table_databricks(
|
|
src,
|
|
ws,
|
|
"core.encounter",
|
|
"homelab",
|
|
"wh123",
|
|
"/Volumes/homelab/default/_sync_staging",
|
|
)
|
|
assert ev.status == "copied"
|
|
|
|
|
|
class TestEnsureStagingVolume:
|
|
"""_ensure_staging_volume creates volume if needed."""
|
|
|
|
def test_volume_exists(self) -> None:
|
|
from aco.lake.sync import _ensure_staging_volume
|
|
|
|
ws = MagicMock()
|
|
result = _ensure_staging_volume(ws, "homelab")
|
|
assert result == "/Volumes/homelab/default/_sync_staging"
|
|
ws.volumes.read.assert_called_once()
|
|
|
|
def test_volume_not_exists_creates(self) -> None:
|
|
from aco.lake.sync import _ensure_staging_volume
|
|
|
|
ws = MagicMock()
|
|
ws.volumes.read.side_effect = Exception("not found")
|
|
|
|
with patch("databricks.sdk.service.catalog.VolumeType") as MockVT:
|
|
MockVT.MANAGED = "MANAGED"
|
|
result = _ensure_staging_volume(ws, "homelab")
|
|
ws.volumes.create.assert_called_once()
|
|
assert result == "/Volumes/homelab/default/_sync_staging"
|
|
|
|
def test_volume_create_race_condition(self) -> None:
|
|
"""Both read and create fail — still returns path."""
|
|
from aco.lake.sync import _ensure_staging_volume
|
|
|
|
ws = MagicMock()
|
|
ws.volumes.read.side_effect = Exception("not found")
|
|
ws.volumes.create.side_effect = Exception("race")
|
|
|
|
with patch("databricks.sdk.service.catalog.VolumeType") as MockVT:
|
|
MockVT.MANAGED = "MANAGED"
|
|
result = _ensure_staging_volume(ws, "homelab")
|
|
assert result == "/Volumes/homelab/default/_sync_staging"
|
|
|
|
|
|
class TestHandleEvent:
|
|
"""_handle_event updates report and invokes callback."""
|
|
|
|
def test_copied_event(self) -> None:
|
|
from aco.lake.sync import SyncEvent, SyncReport, _handle_event
|
|
|
|
report = SyncReport()
|
|
ev = SyncEvent(
|
|
table_ref="core.encounter",
|
|
status="copied",
|
|
rows=100,
|
|
bytes=5000,
|
|
elapsed_seconds=1.5,
|
|
)
|
|
_handle_event(ev, report, "collect", None)
|
|
assert "core.encounter" in report.tables_synced
|
|
assert report.total_rows == 100
|
|
assert report.total_bytes == 5000
|
|
|
|
def test_skipped_event(self) -> None:
|
|
from aco.lake.sync import SyncEvent, SyncReport, _handle_event
|
|
|
|
report = SyncReport()
|
|
ev = SyncEvent(table_ref="core.empty", status="skipped")
|
|
_handle_event(ev, report, "collect", None)
|
|
assert "core.empty" in report.tables_skipped
|
|
|
|
def test_failed_event_collect(self) -> None:
|
|
from aco.lake.sync import SyncEvent, SyncReport, _handle_event
|
|
|
|
report = SyncReport()
|
|
ev = SyncEvent(
|
|
table_ref="core.fail",
|
|
status="failed",
|
|
error="oops",
|
|
)
|
|
_handle_event(ev, report, "collect", None)
|
|
assert len(report.tables_failed) == 1
|
|
assert report.tables_failed[0].table_ref == "core.fail"
|
|
|
|
def test_failed_event_raise(self) -> None:
|
|
from aco.lake.sync import SyncEvent, SyncReport, _handle_event
|
|
|
|
report = SyncReport()
|
|
ev = SyncEvent(
|
|
table_ref="core.fail",
|
|
status="failed",
|
|
error="oops",
|
|
)
|
|
with pytest.raises(RuntimeError, match="Sync failed"):
|
|
_handle_event(ev, report, "raise", None)
|
|
|
|
def test_failed_event_skip(self) -> None:
|
|
from aco.lake.sync import SyncEvent, SyncReport, _handle_event
|
|
|
|
report = SyncReport()
|
|
ev = SyncEvent(
|
|
table_ref="core.fail",
|
|
status="failed",
|
|
error="oops",
|
|
)
|
|
_handle_event(ev, report, "skip", None)
|
|
assert len(report.tables_failed) == 0
|
|
|
|
def test_progress_callback(self) -> None:
|
|
from aco.lake.sync import SyncEvent, SyncReport, _handle_event
|
|
|
|
report = SyncReport()
|
|
ev = SyncEvent(
|
|
table_ref="core.encounter",
|
|
status="copied",
|
|
rows=10,
|
|
)
|
|
callback = MagicMock()
|
|
_handle_event(ev, report, "collect", callback)
|
|
callback.assert_called_once_with(ev)
|
|
|
|
|
|
class TestWidenForDelta:
|
|
"""_widen_for_delta upcasts narrow types."""
|
|
|
|
def test_int8_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1, 2], dtype=pl.Int8)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_int16_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1, 2], dtype=pl.Int16)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_int32_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1], dtype=pl.Int32)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_uint8_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1], dtype=pl.UInt8)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_uint16_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1], dtype=pl.UInt16)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_uint32_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1], dtype=pl.UInt32)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_uint64_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1], dtype=pl.UInt64)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_float32_to_float64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1.0], dtype=pl.Float32)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Float64
|
|
|
|
def test_decimal_scale0_to_int64(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame(
|
|
{"x": pl.Series([1], dtype=pl.Decimal(precision=10, scale=0))}
|
|
)
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_decimal_nonzero_scale_to_decimal38(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame(
|
|
{"x": pl.Series(["1.23"], dtype=pl.Decimal(precision=10, scale=2))}
|
|
)
|
|
result = _widen_for_delta(df)
|
|
assert isinstance(result["x"].dtype, pl.Decimal)
|
|
assert result["x"].dtype.precision == 38
|
|
|
|
def test_int64_unchanged(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series([1], dtype=pl.Int64)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Int64
|
|
|
|
def test_no_cast_needed(self) -> None:
|
|
from aco.lake.sync import _widen_for_delta
|
|
|
|
df = pl.DataFrame({"x": pl.Series(["a"], dtype=pl.Utf8)})
|
|
result = _widen_for_delta(df)
|
|
assert result["x"].dtype == pl.Utf8
|
|
|
|
|
|
class TestSyncEntryPoint:
|
|
"""sync() function dispatches to the right path."""
|
|
|
|
def test_sync_databricks_path_empty(self) -> None:
|
|
from aco.lake.sync import sync
|
|
|
|
src = MagicMock(spec=DuckDBContext)
|
|
|
|
with (
|
|
patch("aco.lake.unity.UnityClient.from_env") as mock_from_env,
|
|
patch("aco.lake.sync._ensure_staging_volume") as mock_vol,
|
|
):
|
|
mock_client = MagicMock()
|
|
mock_from_env.return_value = mock_client
|
|
mock_client._ws = MagicMock()
|
|
mock_vol.return_value = "/Volumes/homelab/default/_sync"
|
|
|
|
report = sync(src, warehouse_id="wh123", tables=[])
|
|
assert isinstance(report, object)
|
|
|
|
def test_sync_databricks_path_with_tables(self) -> None:
|
|
from aco.lake.sync import sync
|
|
|
|
src = MagicMock(spec=DuckDBContext)
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
|
|
with (
|
|
patch("aco.lake.unity.UnityClient.from_env") as mock_from_env,
|
|
patch("aco.lake.sync._ensure_staging_volume") as mock_vol,
|
|
):
|
|
mock_client = MagicMock()
|
|
mock_from_env.return_value = mock_client
|
|
mock_ws = MagicMock()
|
|
mock_client._ws = mock_ws
|
|
mock_vol.return_value = "/Volumes/homelab/default/_sync"
|
|
|
|
# Mock successful COPY INTO
|
|
resp = MagicMock()
|
|
resp.status.state = "SUCCEEDED"
|
|
resp.status.error = None
|
|
mock_ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
report = sync(
|
|
src,
|
|
warehouse_id="wh123",
|
|
tables=["core.encounter"],
|
|
)
|
|
assert "core.encounter" in report.tables_synced
|
|
assert report.total_rows == 1
|
|
|
|
def test_sync_generic_path(self) -> None:
|
|
from aco.lake.sync import sync
|
|
|
|
src = MagicMock()
|
|
dst = MagicMock()
|
|
df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
src.load.return_value = df
|
|
|
|
report = sync(src, dst, tables=["core.encounter"])
|
|
# Should call _sync_table_generic for each table
|
|
assert report is not None
|
|
assert "core.encounter" in report.tables_synced
|
|
|
|
def test_sync_no_target_no_warehouse_raises(self) -> None:
|
|
from aco.lake.sync import sync
|
|
|
|
src = MagicMock()
|
|
with pytest.raises(ValueError, match="Provide either"):
|
|
sync(src, tables=["core.encounter"])
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Unity module (lines 57-638)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestUnityCatalogModel:
|
|
"""UnityCatalog Pydantic model and from_sdk."""
|
|
|
|
def test_from_sdk(self) -> None:
|
|
from aco.lake.unity import UnityCatalog
|
|
|
|
info = MagicMock(spec=CatalogInfo)
|
|
info.name = "main"
|
|
info.comment = "test catalog"
|
|
info.owner = "admin"
|
|
info.storage_root = "s3://root"
|
|
info.metastore_id = "ms-123"
|
|
info.full_name = "main"
|
|
info.created_at = 1000
|
|
info.updated_at = 2000
|
|
info.catalog_type = MagicMock()
|
|
info.isolation_mode = MagicMock()
|
|
|
|
result = UnityCatalog.from_sdk(info)
|
|
assert result.name == "main"
|
|
assert result.comment == "test catalog"
|
|
assert result.owner == "admin"
|
|
|
|
def test_from_sdk_none_values(self) -> None:
|
|
from aco.lake.unity import UnityCatalog
|
|
|
|
info = MagicMock(spec=CatalogInfo)
|
|
info.name = None
|
|
info.comment = None
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.metastore_id = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
info.catalog_type = None
|
|
info.isolation_mode = None
|
|
|
|
result = UnityCatalog.from_sdk(info)
|
|
assert result.name == ""
|
|
assert result.catalog_type is None
|
|
assert result.isolation_mode is None
|
|
|
|
|
|
class TestUnitySchemaModel:
|
|
"""UnitySchema Pydantic model and from_sdk."""
|
|
|
|
def test_from_sdk(self) -> None:
|
|
from aco.lake.unity import UnitySchema
|
|
|
|
info = MagicMock(spec=SchemaInfo)
|
|
info.name = "core"
|
|
info.catalog_name = "main"
|
|
info.comment = "core schema"
|
|
info.owner = "admin"
|
|
info.storage_root = None
|
|
info.full_name = "main.core"
|
|
info.created_at = 1000
|
|
info.updated_at = 2000
|
|
|
|
result = UnitySchema.from_sdk(info)
|
|
assert result.name == "core"
|
|
assert result.catalog_name == "main"
|
|
|
|
def test_from_sdk_none_name(self) -> None:
|
|
from aco.lake.unity import UnitySchema
|
|
|
|
info = MagicMock(spec=SchemaInfo)
|
|
info.name = None
|
|
info.catalog_name = None
|
|
info.comment = None
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
|
|
result = UnitySchema.from_sdk(info)
|
|
assert result.name == ""
|
|
assert result.catalog_name == ""
|
|
|
|
|
|
class TestUnityTableColumnModel:
|
|
"""UnityTableColumn Pydantic model and from_sdk."""
|
|
|
|
def test_from_sdk(self) -> None:
|
|
from aco.lake.unity import UnityTableColumn
|
|
|
|
info = MagicMock(spec=ColumnInfo)
|
|
info.name = "encounter_id"
|
|
info.type_text = "string"
|
|
info.type_name = MagicMock()
|
|
info.comment = "Primary key"
|
|
info.nullable = False
|
|
info.position = 0
|
|
|
|
result = UnityTableColumn.from_sdk(info)
|
|
assert result.name == "encounter_id"
|
|
assert result.nullable is False
|
|
|
|
def test_from_sdk_none_values(self) -> None:
|
|
from aco.lake.unity import UnityTableColumn
|
|
|
|
info = MagicMock(spec=ColumnInfo)
|
|
info.name = None
|
|
info.type_text = None
|
|
info.type_name = None
|
|
info.comment = None
|
|
info.nullable = None
|
|
info.position = None
|
|
|
|
result = UnityTableColumn.from_sdk(info)
|
|
assert result.name == ""
|
|
assert result.type_text == ""
|
|
assert result.type_name == ""
|
|
assert result.nullable is True
|
|
assert result.position == 0
|
|
|
|
|
|
class TestUnityTableModel:
|
|
"""UnityTable Pydantic model and from_sdk."""
|
|
|
|
def test_from_sdk(self) -> None:
|
|
from aco.lake.unity import UnityTable
|
|
|
|
col_info = MagicMock(spec=ColumnInfo)
|
|
col_info.name = "id"
|
|
col_info.type_text = "long"
|
|
col_info.type_name = MagicMock()
|
|
col_info.comment = None
|
|
col_info.nullable = True
|
|
col_info.position = 0
|
|
|
|
info = MagicMock(spec=TableInfo)
|
|
info.name = "encounter"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "core"
|
|
info.table_type = MagicMock()
|
|
info.data_source_format = MagicMock()
|
|
info.columns = [col_info]
|
|
info.comment = "Encounter table"
|
|
info.storage_location = "s3://tables/enc"
|
|
info.owner = "admin"
|
|
info.full_name = "main.core.encounter"
|
|
info.created_at = 1000
|
|
info.updated_at = 2000
|
|
|
|
result = UnityTable.from_sdk(info)
|
|
assert result.name == "encounter"
|
|
assert len(result.columns) == 1
|
|
assert result.columns[0].name == "id"
|
|
|
|
def test_from_sdk_no_columns(self) -> None:
|
|
from aco.lake.unity import UnityTable
|
|
|
|
info = MagicMock(spec=TableInfo)
|
|
info.name = "empty"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "core"
|
|
info.table_type = None
|
|
info.data_source_format = None
|
|
info.columns = None
|
|
info.comment = None
|
|
info.storage_location = None
|
|
info.owner = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
|
|
result = UnityTable.from_sdk(info)
|
|
assert result.name == "empty"
|
|
assert result.columns == []
|
|
assert result.table_type == "MANAGED"
|
|
|
|
|
|
class TestUnityVolumeModel:
|
|
"""UnityVolume Pydantic model and from_sdk."""
|
|
|
|
def test_from_sdk(self) -> None:
|
|
from aco.lake.unity import UnityVolume
|
|
|
|
info = MagicMock(spec=VolumeInfo)
|
|
info.name = "staging"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "default"
|
|
info.volume_type = MagicMock()
|
|
info.storage_location = "s3://volumes/staging"
|
|
info.comment = "Staging volume"
|
|
info.owner = "admin"
|
|
info.full_name = "main.default.staging"
|
|
info.created_at = 1000
|
|
info.updated_at = 2000
|
|
|
|
result = UnityVolume.from_sdk(info)
|
|
assert result.name == "staging"
|
|
assert result.catalog_name == "main"
|
|
|
|
def test_from_sdk_none_values(self) -> None:
|
|
from aco.lake.unity import UnityVolume
|
|
|
|
info = MagicMock(spec=VolumeInfo)
|
|
info.name = None
|
|
info.catalog_name = None
|
|
info.schema_name = None
|
|
info.volume_type = None
|
|
info.storage_location = None
|
|
info.comment = None
|
|
info.owner = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
|
|
result = UnityVolume.from_sdk(info)
|
|
assert result.name == ""
|
|
assert result.volume_type == "MANAGED"
|
|
|
|
|
|
class TestUnityClientInit:
|
|
"""UnityClient construction."""
|
|
|
|
def test_init_stores_workspace(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock(spec=WorkspaceClient)
|
|
client = UnityClient(ws)
|
|
assert client._ws is ws
|
|
|
|
def test_host_property(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock(spec=WorkspaceClient)
|
|
ws.config.host = "https://example.databricks.com"
|
|
client = UnityClient(ws)
|
|
assert client.host == "https://example.databricks.com"
|
|
|
|
|
|
class TestUnityClientFromEnv:
|
|
"""UnityClient.from_env reads credentials."""
|
|
|
|
def test_from_env_pat_auth(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
env = {
|
|
"DATABRICKS_HOST": "https://example.databricks.com",
|
|
"DATABRICKS_TOKEN": "dapi123",
|
|
}
|
|
with (
|
|
patch.dict("os.environ", env, clear=False),
|
|
patch("os.path.exists", return_value=False),
|
|
patch("aco.lake.unity.WorkspaceClient") as MockWS,
|
|
):
|
|
MockWS.return_value = MagicMock()
|
|
UnityClient.from_env()
|
|
MockWS.assert_called_once_with(
|
|
host="https://example.databricks.com",
|
|
token="dapi123",
|
|
)
|
|
|
|
def test_from_env_oauth_auth(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
env = {
|
|
"DATABRICKS_HOST": "https://example.databricks.com",
|
|
"DATABRICKS_CLIENT_ID": "client123",
|
|
"DATABRICKS_CLIENT_SECRET": "secret456",
|
|
}
|
|
# Remove DATABRICKS_TOKEN if present
|
|
with (
|
|
patch.dict("os.environ", env, clear=False),
|
|
patch("os.path.exists", return_value=False),
|
|
patch("aco.lake.unity.WorkspaceClient") as MockWS,
|
|
):
|
|
# Ensure TOKEN is not set
|
|
import os
|
|
|
|
os.environ.pop("DATABRICKS_TOKEN", None)
|
|
MockWS.return_value = MagicMock()
|
|
UnityClient.from_env()
|
|
MockWS.assert_called_once_with(
|
|
host="https://example.databricks.com",
|
|
client_id="client123",
|
|
client_secret="secret456",
|
|
)
|
|
|
|
def test_from_env_no_host_raises(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
with (
|
|
patch.dict(
|
|
"os.environ",
|
|
{},
|
|
clear=True,
|
|
),
|
|
patch("os.path.exists", return_value=False),
|
|
):
|
|
# Clear relevant env vars
|
|
import os
|
|
|
|
for k in list(os.environ):
|
|
if k.startswith("DATABRICKS_"):
|
|
del os.environ[k]
|
|
with pytest.raises(RuntimeError, match="DATABRICKS_HOST"):
|
|
UnityClient.from_env()
|
|
|
|
def test_from_env_no_token_no_client_raises(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
env = {
|
|
"DATABRICKS_HOST": "https://example.databricks.com",
|
|
}
|
|
with (
|
|
patch.dict("os.environ", env, clear=False),
|
|
patch("os.path.exists", return_value=False),
|
|
):
|
|
import os
|
|
|
|
os.environ.pop("DATABRICKS_TOKEN", None)
|
|
os.environ.pop("DATABRICKS_CLIENT_ID", None)
|
|
os.environ.pop("DATABRICKS_CLIENT_SECRET", None)
|
|
with pytest.raises(RuntimeError, match="DATABRICKS_TOKEN"):
|
|
UnityClient.from_env()
|
|
|
|
def test_from_env_reads_dotenv(self, tmp_path) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
dotenv = tmp_path / ".env"
|
|
dotenv.write_text(
|
|
"DATABRICKS_HOST=https://test.databricks.com\n"
|
|
"DATABRICKS_TOKEN=dapi_test\n"
|
|
"# comment\n"
|
|
"\n"
|
|
"OTHER_VAR=ignored\n"
|
|
)
|
|
|
|
with (
|
|
patch.dict("os.environ", {}, clear=False),
|
|
patch("aco.lake.unity.WorkspaceClient") as MockWS,
|
|
):
|
|
import os
|
|
|
|
os.environ.pop("DATABRICKS_HOST", None)
|
|
os.environ.pop("DATABRICKS_TOKEN", None)
|
|
os.environ.pop("DATABRICKS_CLIENT_ID", None)
|
|
os.environ.pop("DATABRICKS_CLIENT_SECRET", None)
|
|
MockWS.return_value = MagicMock()
|
|
UnityClient.from_env(dotenv_path=str(dotenv))
|
|
# Should have loaded DATABRICKS_HOST and TOKEN
|
|
assert os.environ.get("DATABRICKS_HOST") == ("https://test.databricks.com")
|
|
assert os.environ.get("DATABRICKS_TOKEN") == "dapi_test"
|
|
# OTHER_VAR should NOT be set (not DATABRICKS_ prefix)
|
|
MockWS.assert_called_once()
|
|
|
|
|
|
class TestUnityClientCatalogMethods:
|
|
"""UnityClient catalog CRUD methods."""
|
|
|
|
def test_list_catalogs(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=CatalogInfo)
|
|
info.name = "main"
|
|
info.comment = None
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.metastore_id = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
info.catalog_type = None
|
|
info.isolation_mode = None
|
|
ws.catalogs.list.return_value = [info]
|
|
|
|
client = UnityClient(ws)
|
|
result = client.list_catalogs()
|
|
assert len(result) == 1
|
|
assert result[0].name == "main"
|
|
|
|
def test_get_catalog(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=CatalogInfo)
|
|
info.name = "main"
|
|
info.comment = None
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.metastore_id = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
info.catalog_type = None
|
|
info.isolation_mode = None
|
|
ws.catalogs.get.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
result = client.get_catalog("main")
|
|
assert result.name == "main"
|
|
ws.catalogs.get.assert_called_with("main")
|
|
|
|
def test_create_catalog(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=CatalogInfo)
|
|
info.name = "new_cat"
|
|
info.comment = "new"
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.metastore_id = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
info.catalog_type = None
|
|
info.isolation_mode = None
|
|
ws.catalogs.create.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
result = client.create_catalog("new_cat", comment="new")
|
|
assert result.name == "new_cat"
|
|
|
|
def test_update_catalog(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=CatalogInfo)
|
|
info.name = "main"
|
|
info.comment = "updated"
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.metastore_id = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
info.catalog_type = None
|
|
info.isolation_mode = None
|
|
ws.catalogs.update.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
result = client.update_catalog("main", comment="updated")
|
|
assert result.comment == "updated"
|
|
|
|
def test_delete_catalog(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
client.delete_catalog("main", force=True)
|
|
ws.catalogs.delete.assert_called_with("main", force=True)
|
|
|
|
|
|
class TestUnityClientSchemaMethods:
|
|
"""UnityClient schema CRUD methods."""
|
|
|
|
def test_list_schemas(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=SchemaInfo)
|
|
info.name = "core"
|
|
info.catalog_name = "main"
|
|
info.comment = None
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.schemas.list.return_value = [info]
|
|
|
|
client = UnityClient(ws)
|
|
result = client.list_schemas("main")
|
|
assert len(result) == 1
|
|
assert result[0].name == "core"
|
|
|
|
def test_get_schema(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=SchemaInfo)
|
|
info.name = "core"
|
|
info.catalog_name = "main"
|
|
info.comment = None
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.schemas.get.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
client.get_schema("main", "core")
|
|
ws.schemas.get.assert_called_with("main.core")
|
|
|
|
def test_create_schema(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=SchemaInfo)
|
|
info.name = "core"
|
|
info.catalog_name = "main"
|
|
info.comment = "Core schema"
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.schemas.create.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
result = client.create_schema("main", "core", comment="Core schema")
|
|
assert result.name == "core"
|
|
|
|
def test_update_schema(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=SchemaInfo)
|
|
info.name = "core"
|
|
info.catalog_name = "main"
|
|
info.comment = "updated"
|
|
info.owner = None
|
|
info.storage_root = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.schemas.update.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
client.update_schema("main", "core", comment="updated")
|
|
ws.schemas.update.assert_called_with("main.core", comment="updated", owner=None)
|
|
|
|
def test_delete_schema(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
client.delete_schema("main", "core", force=True)
|
|
ws.schemas.delete.assert_called_with("main.core", force=True)
|
|
|
|
|
|
class TestUnityClientTableMethods:
|
|
"""UnityClient table CRUD methods."""
|
|
|
|
def test_list_tables(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=TableInfo)
|
|
info.name = "encounter"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "core"
|
|
info.table_type = None
|
|
info.data_source_format = None
|
|
info.columns = None
|
|
info.comment = None
|
|
info.storage_location = None
|
|
info.owner = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.tables.list.return_value = [info]
|
|
|
|
client = UnityClient(ws)
|
|
result = client.list_tables("main", "core")
|
|
assert len(result) == 1
|
|
assert result[0].name == "encounter"
|
|
|
|
def test_get_table(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=TableInfo)
|
|
info.name = "encounter"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "core"
|
|
info.table_type = None
|
|
info.data_source_format = None
|
|
info.columns = None
|
|
info.comment = None
|
|
info.storage_location = None
|
|
info.owner = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.tables.get.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
client.get_table("main", "core", "encounter")
|
|
ws.tables.get.assert_called_with("main.core.encounter")
|
|
|
|
def test_delete_table(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
client.delete_table("main", "core", "encounter")
|
|
ws.tables.delete.assert_called_with("main.core.encounter")
|
|
|
|
|
|
class TestUnityClientVolumeMethods:
|
|
"""UnityClient volume CRUD methods."""
|
|
|
|
def test_list_volumes(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=VolumeInfo)
|
|
info.name = "staging"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "default"
|
|
info.volume_type = None
|
|
info.storage_location = None
|
|
info.comment = None
|
|
info.owner = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.volumes.list.return_value = [info]
|
|
|
|
client = UnityClient(ws)
|
|
result = client.list_volumes("main", "default")
|
|
assert len(result) == 1
|
|
assert result[0].name == "staging"
|
|
|
|
def test_create_volume(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
info = MagicMock(spec=VolumeInfo)
|
|
info.name = "newvol"
|
|
info.catalog_name = "main"
|
|
info.schema_name = "default"
|
|
info.volume_type = MagicMock()
|
|
info.storage_location = None
|
|
info.comment = None
|
|
info.owner = None
|
|
info.full_name = None
|
|
info.created_at = None
|
|
info.updated_at = None
|
|
ws.volumes.create.return_value = info
|
|
|
|
client = UnityClient(ws)
|
|
with patch("databricks.sdk.service.catalog.VolumeType") as MockVT:
|
|
MockVT.return_value = "MANAGED"
|
|
result = client.create_volume("main", "default", "newvol")
|
|
assert result.name == "newvol"
|
|
|
|
def test_delete_volume(self) -> None:
|
|
from aco.lake.unity import UnityClient
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
client.delete_volume("main", "default", "staging")
|
|
ws.volumes.delete.assert_called_with("main.default.staging")
|
|
|
|
|
|
class TestPythonTypeToDatabricks:
|
|
"""_python_type_to_databricks maps annotations to Databricks types."""
|
|
|
|
def test_str(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(str) == ("STRING", "string")
|
|
|
|
def test_int(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(int) == ("LONG", "long")
|
|
|
|
def test_float(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(float) == ("DOUBLE", "double")
|
|
|
|
def test_bool(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(bool) == ("BOOLEAN", "boolean")
|
|
|
|
def test_date(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(date) == ("DATE", "date")
|
|
|
|
def test_datetime(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(datetime) == (
|
|
"TIMESTAMP_NTZ",
|
|
"timestamp_ntz",
|
|
)
|
|
|
|
def test_decimal(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(Decimal) == ("DECIMAL", "decimal(38,2)")
|
|
|
|
def test_unknown_fallback(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(bytes) == ("STRING", "string")
|
|
|
|
def test_union_type(self) -> None:
|
|
from aco.lake.unity import _python_type_to_databricks
|
|
|
|
assert _python_type_to_databricks(int | None) == ("LONG", "long")
|
|
|
|
|
|
class TestSqlTableToColumnInfos:
|
|
"""_sql_table_to_column_infos converts model fields."""
|
|
|
|
def test_converts_fields(self) -> None:
|
|
from aco.lake.unity import _sql_table_to_column_infos
|
|
|
|
MockModel = type(
|
|
"MockModel",
|
|
(SQLTable,),
|
|
{
|
|
"__schema__": "test",
|
|
"__tablename__": "t",
|
|
"__annotations__": {
|
|
"id": int,
|
|
"name": str,
|
|
},
|
|
},
|
|
)
|
|
|
|
result = _sql_table_to_column_infos(MockModel)
|
|
assert len(result) >= 0 # May or may not have fields
|
|
# Test with explicit model_fields
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int, description="Primary key"),
|
|
"name": FieldInfo(annotation=str),
|
|
}
|
|
|
|
result = _sql_table_to_column_infos(mock_model)
|
|
assert len(result) == 2
|
|
assert result[0].name == "id"
|
|
assert result[0].comment == "Primary key"
|
|
assert result[1].name == "name"
|
|
|
|
|
|
class TestModelToDdl:
|
|
"""_model_to_ddl generates DDL from SQLTable model."""
|
|
|
|
def test_generates_ddl(self) -> None:
|
|
from pydantic.fields import FieldInfo
|
|
|
|
from aco.lake.unity import _model_to_ddl
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
"name": FieldInfo(annotation=str),
|
|
}
|
|
mock_model.__doc__ = "Test table doc"
|
|
|
|
result = _model_to_ddl("homelab", "core", "encounter", mock_model)
|
|
assert "CREATE TABLE IF NOT EXISTS" in result
|
|
assert "`homelab`.`core`.`encounter`" in result
|
|
assert "LONG" in result.upper() or "long" in result
|
|
assert "COMMENT" in result
|
|
|
|
def test_generates_ddl_no_doc(self) -> None:
|
|
from pydantic.fields import FieldInfo
|
|
|
|
from aco.lake.unity import _model_to_ddl
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = ""
|
|
|
|
result = _model_to_ddl("homelab", "core", "t", mock_model)
|
|
assert "CREATE TABLE IF NOT EXISTS" in result
|
|
assert "COMMENT" not in result
|
|
|
|
def test_generates_ddl_long_doc_truncated(self) -> None:
|
|
from pydantic.fields import FieldInfo
|
|
|
|
from aco.lake.unity import _model_to_ddl
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = "x" * 2000
|
|
|
|
result = _model_to_ddl("homelab", "core", "t", mock_model)
|
|
assert "..." in result
|
|
|
|
|
|
class TestSetupCatalogFromSchemas:
|
|
"""setup_catalog_from_schemas creates schemas and tables."""
|
|
|
|
def test_dry_run(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = "Test"
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
setup_catalog_from_schemas(client, "homelab", "wh123", dry_run=True)
|
|
# Dry run should not call SDK
|
|
ws.statement_execution.execute_statement.assert_not_called()
|
|
|
|
def test_creates_schemas_and_tables(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
|
|
schema_info = MagicMock(spec=SchemaInfo)
|
|
schema_info.name = "default"
|
|
schema_info.catalog_name = "homelab"
|
|
schema_info.comment = None
|
|
schema_info.owner = None
|
|
schema_info.storage_root = None
|
|
schema_info.full_name = None
|
|
schema_info.created_at = None
|
|
schema_info.updated_at = None
|
|
ws.schemas.list.return_value = [schema_info]
|
|
|
|
create_info = MagicMock(spec=SchemaInfo)
|
|
create_info.name = "core"
|
|
create_info.catalog_name = "homelab"
|
|
create_info.comment = None
|
|
create_info.owner = None
|
|
create_info.storage_root = None
|
|
create_info.full_name = None
|
|
create_info.created_at = None
|
|
create_info.updated_at = None
|
|
ws.schemas.create.return_value = create_info
|
|
|
|
resp = MagicMock()
|
|
resp.status.state = "SUCCEEDED"
|
|
resp.status.error = None
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = "Test"
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert "core" in report["schemas_created"]
|
|
assert "core.encounter" in report["tables_created"]
|
|
|
|
def test_schema_create_error(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
ws.schemas.list.return_value = []
|
|
ws.schemas.create.side_effect = Exception("perm denied")
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = ""
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert len(report["errors"]) > 0
|
|
|
|
def test_table_create_error(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
|
|
schema_info = MagicMock(spec=SchemaInfo)
|
|
schema_info.name = "core"
|
|
schema_info.catalog_name = "homelab"
|
|
schema_info.comment = None
|
|
schema_info.owner = None
|
|
schema_info.storage_root = None
|
|
schema_info.full_name = None
|
|
schema_info.created_at = None
|
|
schema_info.updated_at = None
|
|
ws.schemas.list.return_value = [schema_info]
|
|
|
|
ws.statement_execution.execute_statement.side_effect = Exception("sql error")
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = ""
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert len(report["errors"]) > 0
|
|
assert any("table" in e for e in report["errors"])
|
|
|
|
def test_table_ddl_non_succeeded(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
|
|
schema_info = MagicMock(spec=SchemaInfo)
|
|
schema_info.name = "core"
|
|
schema_info.catalog_name = "homelab"
|
|
schema_info.comment = None
|
|
schema_info.owner = None
|
|
schema_info.storage_root = None
|
|
schema_info.full_name = None
|
|
schema_info.created_at = None
|
|
schema_info.updated_at = None
|
|
ws.schemas.list.return_value = [schema_info]
|
|
|
|
resp = MagicMock()
|
|
resp.status.state = "FAILED"
|
|
resp.status.error.message = "bad ddl"
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = ""
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert len(report["errors"]) > 0
|
|
|
|
def test_table_ddl_non_succeeded_no_error_msg(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
|
|
schema_info = MagicMock(spec=SchemaInfo)
|
|
schema_info.name = "core"
|
|
schema_info.catalog_name = "homelab"
|
|
schema_info.comment = None
|
|
schema_info.owner = None
|
|
schema_info.storage_root = None
|
|
schema_info.full_name = None
|
|
schema_info.created_at = None
|
|
schema_info.updated_at = None
|
|
ws.schemas.list.return_value = [schema_info]
|
|
|
|
resp = MagicMock()
|
|
resp.status.state = "FAILED"
|
|
resp.status.error = None
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = ""
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert len(report["errors"]) > 0
|
|
|
|
def test_model_lookup_error(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
ws.schemas.list.return_value = []
|
|
|
|
create_info = MagicMock(spec=SchemaInfo)
|
|
create_info.name = "core"
|
|
create_info.catalog_name = "homelab"
|
|
create_info.comment = None
|
|
create_info.owner = None
|
|
create_info.storage_root = None
|
|
create_info.full_name = None
|
|
create_info.created_at = None
|
|
create_info.updated_at = None
|
|
ws.schemas.create.return_value = create_info
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
mock_cat.model.side_effect = ValueError("not found")
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert len(report["errors"]) > 0
|
|
assert any("model" in e for e in report["errors"])
|
|
|
|
def test_list_existing_schemas_error(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
ws.schemas.list.side_effect = Exception("api error")
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = []
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False
|
|
)
|
|
assert len(report["errors"]) > 0
|
|
|
|
def test_existing_schema_skipped(self) -> None:
|
|
from aco.lake.unity import UnityClient, setup_catalog_from_schemas
|
|
|
|
ws = MagicMock()
|
|
client = UnityClient(ws)
|
|
|
|
schema_info = MagicMock(spec=SchemaInfo)
|
|
schema_info.name = "core"
|
|
schema_info.catalog_name = "homelab"
|
|
schema_info.comment = None
|
|
schema_info.owner = None
|
|
schema_info.storage_root = None
|
|
schema_info.full_name = None
|
|
schema_info.created_at = None
|
|
schema_info.updated_at = None
|
|
ws.schemas.list.return_value = [schema_info]
|
|
|
|
resp = MagicMock()
|
|
resp.status.state = "SUCCEEDED"
|
|
resp.status.error = None
|
|
ws.statement_execution.execute_statement.return_value = resp
|
|
|
|
with patch("aco.lake.catalog.Catalog") as MockCat:
|
|
mock_cat = MockCat.return_value
|
|
mock_cat.schemas.return_value = ["core"]
|
|
mock_cat.tables.return_value = ["core.encounter"]
|
|
|
|
from pydantic.fields import FieldInfo
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_fields = {
|
|
"id": FieldInfo(annotation=int),
|
|
}
|
|
mock_model.__doc__ = ""
|
|
mock_cat.model.return_value = mock_model
|
|
|
|
report = setup_catalog_from_schemas(
|
|
client, "homelab", "wh123", dry_run=False, skip_existing=True
|
|
)
|
|
assert "core" in report["schemas_skipped"]
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
# Transpile: _clean_duckdb_aliases identifier branch (line 324-325)
|
|
# ═══════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class TestCleanDuckdbAliasesIdentifierBranch:
|
|
"""Cover the Identifier node branch for unnamed_relation refs."""
|
|
|
|
def test_unnamed_relation_identifier_rewritten(self) -> None:
|
|
import sqlglot
|
|
|
|
# BFS walk order: TableAlias is visited before the
|
|
# Identifier in the CROSS JOIN subquery, so `seen`
|
|
# dict is populated when the Identifier is encountered.
|
|
sql = (
|
|
"SELECT * FROM "
|
|
"(SELECT 1 AS x) AS unnamed_relation_abc123def "
|
|
"CROSS JOIN "
|
|
"(SELECT unnamed_relation_abc123def.x) AS sub"
|
|
)
|
|
tree = sqlglot.parse_one(sql, read="duckdb")
|
|
_clean_duckdb_aliases(tree)
|
|
result = tree.sql(dialect="duckdb")
|
|
assert "unnamed_relation" not in result
|
|
assert "t1" in result
|