homelab/stack
1
0
Fork 0
Code Issues 20 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

implement P4 lakehouse: deploy, load, and validate Iceberg tables (fix #14, fix #15, fix #16, fix #17)

Browse Source 
- src/aco/lake/deploy.py: deploy_schemas() creates Iceberg namespaces
  and tables from SQLTable models via PyIceberg REST Catalog
- src/aco/lake/load.py: load_to_iceberg() copies DuckDB tables to
  Iceberg with append/replace modes
- src/cli/lake.py: wire deploy, load, validate commands with
  --catalog-type (nessie/polaris), --schema, --dry-run, --mode options
- Catalog.validate() compares SQLTable definitions vs Iceberg metadata
- 20 new tests (type resolution, dry run, mock deploy/load, CLI help)
- All 11339 tests pass
This commit is contained in:
kert
2026-03-12 17:06:58 -04:00
parent 8047640c3e
commit 4f67b4f9fd
5 changed files with 631 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
src/aco/lake/deploy.py Normal file
View File

@@ -0,0 +1,190 @@
"""Deploy Iceberg schemas and tables from SQLTable models.
Reads all SQLTable models from ``aco.table``, ``bcda.table``, etc.,
creates Iceberg namespaces and tables in the configured catalog
(Nessie or Polaris), matching the column schemas defined in Pydantic.
Usage::
from aco.lake.deploy import deploy_schemas
stats = deploy_schemas(catalog_uri="http://nessie:19120/iceberg/")
"""
from __future__ import annotations
import logging
from typing import Any
log = logging.getLogger(__name__)
# Python type name → PyIceberg type mapping
_TYPE_MAP: dict[str, str] = {
"str": "string",
"int": "long",
"float": "double",
"bool": "boolean",
"date": "date",
"datetime": "timestamp",
"Decimal": "decimal(18,2)",
}
def _resolve_iceberg_type(annotation: Any) -> str:
"""Map a Python type annotation to an Iceberg type string."""
import types
if isinstance(annotation, types.UnionType):
args = [a for a in annotation.__args__ if a is not type(None)]
if args:
return _resolve_iceberg_type(args[0])
if isinstance(annotation, type):
name = annotation.__name__
return _TYPE_MAP.get(name, "string")
return "string"
def deploy_schemas(
catalog_uri: str | None = None,
warehouse: str | None = None,
*,
catalog_type: str = "nessie",
schemas: list[str] | None = None,
dry_run: bool = False,
properties: dict[str, str] | None = None,
) -> dict[str, list[str]]:
"""Create Iceberg namespaces and tables from SQLTable models.
Parameters
----------
catalog_uri : str, optional
REST catalog endpoint. Defaults to stack.toml config.
warehouse : str, optional
Warehouse location. Defaults to stack.toml config.
catalog_type : str
Which catalog config to use: ``nessie`` or ``polaris``.
schemas : list[str], optional
Only deploy these schemas. If None, deploys all discovered schemas.
dry_run : bool
If True, only report what would be created without making changes.
properties : dict, optional
Extra catalog properties (S3 creds, etc.).
Returns
-------
dict[str, list[str]]
Mapping of schema → list of created table names.
"""
from aco.lake.catalog import Catalog
from conf import cfg
if catalog_uri is None:
lake_cfg = cfg.lake
if catalog_type == "polaris":
catalog_uri = lake_cfg.polaris.catalog_uri
else:
catalog_uri = lake_cfg.nessie.catalog_uri
if warehouse is None:
warehouse = cfg.lake.warehouse
cat = Catalog(
catalog_uri=catalog_uri,
warehouse=warehouse,
properties=properties or {},
)
target_schemas = schemas or cat.schemas()
results: dict[str, list[str]] = {}
for schema_name in target_schemas:
tables = cat.tables(schema_name)
if not tables:
continue
created: list[str] = []
if not dry_run:
_ensure_namespace(cat, schema_name)
for table_ref in tables:
model = cat.model(table_ref)
if dry_run:
log.info(
"Would create %s (%d columns)", table_ref, len(model.model_fields)
)
created.append(table_ref)
continue
try:
_create_iceberg_table(cat, table_ref, model)
created.append(table_ref)
log.info("Created %s", table_ref)
except Exception as e:
if "already exists" in str(e).lower():
log.info("Exists: %s", table_ref)
created.append(table_ref)
else:
log.warning("Failed to create %s: %s", table_ref, e)
results[schema_name] = created
return results
def _ensure_namespace(cat: Any, namespace: str) -> None:
"""Create an Iceberg namespace if it doesn't exist."""
iceberg_cat = cat._get_iceberg_catalog()
try:
iceberg_cat.create_namespace(namespace)
except Exception:
pass # Already exists
def _create_iceberg_table(cat: Any, table_ref: str, model: type) -> None:
"""Create an Iceberg table from a SQLTable model."""
from pyiceberg.schema import Schema
from pyiceberg.types import (
BooleanType,
DateType,
DecimalType,
DoubleType,
LongType,
NestedField,
StringType,
TimestampType,
)
type_map = {
"string": StringType(),
"long": LongType(),
"double": DoubleType(),
"boolean": BooleanType(),
"date": DateType(),
"timestamp": TimestampType(),
"decimal(18,2)": DecimalType(18, 2),
}
fields = []
for i, (name, info) in enumerate(model.model_fields.items()):
iceberg_type_name = _resolve_iceberg_type(info.annotation)
iceberg_type = type_map.get(iceberg_type_name, StringType())
fields.append(
NestedField(
field_id=i + 1,
name=name,
field_type=iceberg_type,
required=False,
)
)
schema = Schema(*fields)
parts = table_ref.split(".")
namespace = tuple(parts[:-1])
table_name = parts[-1]
iceberg_cat = cat._get_iceberg_catalog()
iceberg_cat.create_table(
identifier=(*namespace, table_name),
schema=schema,
)

106
src/aco/lake/load.py Normal file
View File

@@ -0,0 +1,106 @@
"""Load data from DuckDB into Iceberg tables.
Reads tables from the local DuckDB database and writes them to the
configured Iceberg catalog (Nessie or Polaris) via PyIceberg.
Usage::
from aco.lake.load import load_to_iceberg
stats = load_to_iceberg(catalog_type="nessie", schemas=["core"])
"""
from __future__ import annotations
import logging
log = logging.getLogger(__name__)
def load_to_iceberg(
catalog_uri: str | None = None,
warehouse: str | None = None,
*,
catalog_type: str = "nessie",
schemas: list[str] | None = None,
mode: str = "append",
properties: dict[str, str] | None = None,
) -> dict[str, list[str]]:
"""Copy tables from DuckDB to Iceberg.
Parameters
----------
catalog_uri : str, optional
REST catalog endpoint. Defaults to stack.toml config.
warehouse : str, optional
Warehouse location. Defaults to stack.toml config.
catalog_type : str
Which catalog config to use: ``nessie`` or ``polaris``.
schemas : list[str], optional
Only load these schemas. If None, loads all discovered schemas.
mode : str
Write mode: ``append`` or ``replace``.
properties : dict, optional
Extra catalog properties (S3 creds, etc.).
Returns
-------
dict[str, list[str]]
Mapping of schema -> list of loaded table names.
"""
from aco.lake.catalog import Catalog
from aco.lake.context import DuckDBContext, IcebergContext
from conf import cfg
if catalog_uri is None:
lake_cfg = cfg.lake
if catalog_type == "polaris":
catalog_uri = lake_cfg.polaris.catalog_uri
else:
catalog_uri = lake_cfg.nessie.catalog_uri
if warehouse is None:
warehouse = cfg.lake.warehouse
cat = Catalog(
catalog_uri=catalog_uri,
warehouse=warehouse,
properties=properties or {},
)
# Source: DuckDB
db_path = str(cfg.path("db.aco"))
duckdb_ctx = DuckDBContext(database=db_path, read_only=True)
# Destination: Iceberg
iceberg_ctx = IcebergContext(
catalog_uri=catalog_uri,
warehouse=warehouse,
properties=properties or {},
)
target_schemas = schemas or cat.schemas()
results: dict[str, list[str]] = {}
for schema_name in target_schemas:
tables = cat.tables(schema_name)
if not tables:
continue
loaded: list[str] = []
for table_ref in tables:
try:
df = duckdb_ctx.load(table_ref)
row_count = len(df)
if row_count == 0:
log.info("Skipping empty table %s", table_ref)
continue
iceberg_ctx.save(table_ref, df, mode=mode)
loaded.append(table_ref)
log.info("Loaded %s (%d rows)", table_ref, row_count)
except Exception as e:
log.warning("Failed to load %s: %s", table_ref, e)
results[schema_name] = loaded
return results

92
src/cli/lake.py
View File

@@ -2,18 +2,102 @@
from __future__ import annotations
from typing import Optional
import typer
app = typer.Typer(no_args_is_help=True)
@app.command()
def deploy() -> None:
def deploy(
catalog_type: str = typer.Option("nessie", help="Catalog type: nessie or polaris."),
schema: Optional[list[str]] = typer.Option(None, help="Only deploy these schemas."),
dry_run: bool = typer.Option(False, "--dry-run", help="Report without creating."),
) -> None:
"""Create Iceberg schemas and tables from SQLTable models."""
typer.echo("lake deploy [stub]")
from aco.lake.deploy import deploy_schemas
results = deploy_schemas(
catalog_type=catalog_type,
schemas=schema,
dry_run=dry_run,
)
total = sum(len(tables) for tables in results.values())
verb = "Would create" if dry_run else "Deployed"
for schema_name, tables in sorted(results.items()):
for t in tables:
typer.echo(f" {t}")
typer.echo(f"{verb} {total} tables across {len(results)} schemas.")
@app.command()
def load() -> None:
def load(
catalog_type: str = typer.Option("nessie", help="Catalog type: nessie or polaris."),
schema: Optional[list[str]] = typer.Option(None, help="Only load these schemas."),
mode: str = typer.Option("append", help="Write mode: append or replace."),
) -> None:
"""Populate Iceberg tables with data from DuckDB."""
typer.echo("lake load [stub]")
from aco.lake.load import load_to_iceberg
results = load_to_iceberg(
catalog_type=catalog_type,
schemas=schema,
mode=mode,
)
total = sum(len(tables) for tables in results.values())
for tables in results.values():
for t in tables:
typer.echo(f" {t}")
typer.echo(f"Loaded {total} tables across {len(results)} schemas.")
@app.command()
def validate(
catalog_type: str = typer.Option("nessie", help="Catalog type: nessie or polaris."),
schema: Optional[list[str]] = typer.Option(
None, help="Only validate these schemas."
),
) -> None:
"""Validate Iceberg tables match SQLTable model definitions."""
from aco.lake.catalog import Catalog
from conf import cfg
lake_cfg = cfg.lake
if catalog_type == "polaris":
catalog_uri = lake_cfg.polaris.catalog_uri
else:
catalog_uri = lake_cfg.nessie.catalog_uri
cat = Catalog(
catalog_uri=catalog_uri,
warehouse=cfg.lake.warehouse,
)
target_schemas = schema or cat.schemas()
issues = 0
for schema_name in target_schemas:
report = cat.validate(schema_name)
if report["missing_in_iceberg"]:
typer.echo(f" {schema_name}: missing in Iceberg:")
for t in report["missing_in_iceberg"]:
typer.echo(f" - {t}")
issues += len(report["missing_in_iceberg"])
if report["missing_in_schema"]:
typer.echo(f" {schema_name}: extra in Iceberg (no model):")
for t in report["missing_in_schema"]:
typer.echo(f" - {t}")
issues += len(report["missing_in_schema"])
if report["column_mismatches"]:
for t, mm in report["column_mismatches"].items():
typer.echo(f" {t}: column mismatch: {mm}")
issues += len(report["column_mismatches"])
if issues == 0:
typer.echo("All schemas valid.")
else:
typer.echo(f"{issues} issues found.")
raise typer.Exit(code=1)

234
tests/aco/test_lake_deploy_load.py Normal file
View File

@@ -0,0 +1,234 @@
"""Tests for aco.lake.deploy and aco.lake.load — Iceberg schema/data ops.
Tests focus on pure logic (type resolution, schema building) and mock
external Iceberg/DuckDB connections since those aren't available in CI.
"""
from __future__ import annotations
from datetime import date, datetime
from decimal import Decimal
from unittest.mock import MagicMock, patch
from aco.lake.deploy import _resolve_iceberg_type
class TestResolveIcebergType:
def test_str(self):
assert _resolve_iceberg_type(str) == "string"
def test_int(self):
assert _resolve_iceberg_type(int) == "long"
def test_float(self):
assert _resolve_iceberg_type(float) == "double"
def test_bool(self):
assert _resolve_iceberg_type(bool) == "boolean"
def test_date(self):
assert _resolve_iceberg_type(date) == "date"
def test_datetime(self):
assert _resolve_iceberg_type(datetime) == "timestamp"
def test_decimal(self):
assert _resolve_iceberg_type(Decimal) == "decimal(18,2)"
def test_unknown_type_defaults_to_string(self):
assert _resolve_iceberg_type(bytes) == "string"
def test_none_defaults_to_string(self):
assert _resolve_iceberg_type(None) == "string"
def test_union_type_unwraps(self):
union = str | None
assert _resolve_iceberg_type(union) == "string"
def test_union_int_none(self):
union = int | None
assert _resolve_iceberg_type(union) == "long"
def test_string_annotation_defaults(self):
assert _resolve_iceberg_type("SomeForwardRef") == "string"
class TestDeploySchemasDryRun:
"""Test deploy_schemas with dry_run=True — no Iceberg connection needed."""
def test_dry_run_discovers_tables(self):
from aco.lake.deploy import deploy_schemas
# When catalog_uri and warehouse are passed explicitly,
# cfg is never accessed, so no mock needed
results = deploy_schemas(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
dry_run=True,
schemas=["core"],
)
assert "core" in results
assert len(results["core"]) > 0
def test_dry_run_all_schemas(self):
from aco.lake.deploy import deploy_schemas
results = deploy_schemas(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
dry_run=True,
)
assert len(results) > 0
total = sum(len(t) for t in results.values())
assert total > 0
class TestDeployLive:
"""Test deploy_schemas with mocked _create_iceberg_table (pyiceberg not installed)."""
def test_creates_tables(self):
from aco.lake.deploy import deploy_schemas
created_tables = []
def fake_create(cat, table_ref, model):
created_tables.append(table_ref)
with (
patch("aco.lake.deploy._create_iceberg_table", side_effect=fake_create),
patch("aco.lake.deploy._ensure_namespace"),
):
results = deploy_schemas(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
schemas=["core"],
)
assert "core" in results
assert len(results["core"]) > 0
assert len(created_tables) > 0
def test_handles_already_exists(self):
from aco.lake.deploy import deploy_schemas
with (
patch(
"aco.lake.deploy._create_iceberg_table",
side_effect=Exception("Table already exists"),
),
patch("aco.lake.deploy._ensure_namespace"),
):
results = deploy_schemas(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
schemas=["core"],
)
assert "core" in results
assert len(results["core"]) > 0
def test_handles_creation_failure(self):
from aco.lake.deploy import deploy_schemas
with (
patch(
"aco.lake.deploy._create_iceberg_table",
side_effect=Exception("Connection refused"),
),
patch("aco.lake.deploy._ensure_namespace"),
):
results = deploy_schemas(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
schemas=["core"],
)
assert "core" in results
assert len(results["core"]) == 0
class TestLoadToIceberg:
"""Test load_to_iceberg with mocked DuckDB and Iceberg contexts."""
def test_loads_tables(self):
import narwhals as nw
import polars as pl
from aco.lake.load import load_to_iceberg
mock_df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]}))
mock_duckdb = MagicMock()
mock_duckdb.load.return_value = mock_df
mock_iceberg = MagicMock()
with (
patch("conf.cfg") as mock_cfg,
patch("aco.lake.context.DuckDBContext", return_value=mock_duckdb),
patch("aco.lake.context.IcebergContext", return_value=mock_iceberg),
):
mock_cfg.path.return_value = "/tmp/fake.duckdb"
results = load_to_iceberg(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
schemas=["core"],
)
assert "core" in results
assert len(results["core"]) > 0
assert mock_iceberg.save.call_count > 0
def test_skips_empty_tables(self):
import narwhals as nw
import polars as pl
from aco.lake.load import load_to_iceberg
empty_df = nw.from_native(pl.DataFrame({"a": pl.Series([], dtype=pl.Int64)}))
mock_duckdb = MagicMock()
mock_duckdb.load.return_value = empty_df
mock_iceberg = MagicMock()
with (
patch("conf.cfg") as mock_cfg,
patch("aco.lake.context.DuckDBContext", return_value=mock_duckdb),
patch("aco.lake.context.IcebergContext", return_value=mock_iceberg),
):
mock_cfg.path.return_value = "/tmp/fake.duckdb"
results = load_to_iceberg(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
schemas=["core"],
)
mock_iceberg.save.assert_not_called()
assert "core" in results
assert len(results["core"]) == 0
def test_handles_load_failure(self):
from aco.lake.load import load_to_iceberg
mock_duckdb = MagicMock()
mock_duckdb.load.side_effect = RuntimeError("Table not found")
mock_iceberg = MagicMock()
with (
patch("conf.cfg") as mock_cfg,
patch("aco.lake.context.DuckDBContext", return_value=mock_duckdb),
patch("aco.lake.context.IcebergContext", return_value=mock_iceberg),
):
mock_cfg.path.return_value = "/tmp/fake.duckdb"
results = load_to_iceberg(
catalog_uri="http://fake:19120/iceberg/",
warehouse="s3://fake/",
schemas=["core"],
)
assert "core" in results
assert len(results["core"]) == 0

19
tests/cli/test_cli.py
View File

@@ -83,15 +83,22 @@ class TestBib:
class TestLake:
def test_lake_deploy(self) -> None:
result = runner.invoke(app, ["lake", "deploy"])
def test_lake_deploy_help(self) -> None:
result = runner.invoke(app, ["lake", "deploy", "--help"])
assert result.exit_code == 0
assert "lake deploy" in result.output
assert "catalog-type" in result.output
assert "dry-run" in result.output
def test_lake_load(self) -> None:
result = runner.invoke(app, ["lake", "load"])
def test_lake_load_help(self) -> None:
result = runner.invoke(app, ["lake", "load", "--help"])
assert result.exit_code == 0
assert "lake load" in result.output
assert "catalog-type" in result.output
assert "mode" in result.output
def test_lake_validate_help(self) -> None:
result = runner.invoke(app, ["lake", "validate", "--help"])
assert result.exit_code == 0
assert "catalog-type" in result.output
class TestDb: