Add init_sql support to DuckDB connection

Author: nicosuave

sidemantic/cli.py

@@ -5,7 +5,7 @@
 import typer
 
 from sidemantic import SemanticLayer, __version__, load_from_directory
-from sidemantic.config import SidemanticConfig, build_connection_string, find_config, load_config
+from sidemantic.config import SidemanticConfig, build_connection_string, find_config, get_init_sql, load_config
 
 
 def version_callback(value: bool):
@@ -348,6 +348,7 @@ def query(
     try:
         # Build connection string from args or config
         connection_str = None
+        init_sql = None
         if connection:
             # Explicit --connection arg provided
             connection_str = connection
@@ -357,6 +358,7 @@ def query(
         elif _loaded_config and _loaded_config.connection:
             # Use connection from config
             connection_str = build_connection_string(_loaded_config)
+            init_sql = get_init_sql(_loaded_config)
         else:
             # Try to find database file in data/
             data_dir = models / "data"
@@ -369,7 +371,9 @@ def query(
         preagg_db = _loaded_config.preagg_database if _loaded_config else None
         preagg_sch = _loaded_config.preagg_schema if _loaded_config else None
         if connection_str:
-            layer = SemanticLayer(connection=connection_str, preagg_database=preagg_db, preagg_schema=preagg_sch)
+            layer = SemanticLayer(
+                connection=connection_str, preagg_database=preagg_db, preagg_schema=preagg_sch, init_sql=init_sql
+            )
         else:
             layer = SemanticLayer(preagg_database=preagg_db, preagg_schema=preagg_sch)
         load_from_directory(layer, str(models))

sidemantic/config.py

@@ -11,6 +11,10 @@ class DuckDBConnection(BaseModel):
 
     type: Literal["duckdb"] = "duckdb"
     path: str = Field(..., description="Path to DuckDB database file or :memory:")
+    init_sql: list[str] | None = Field(
+        default=None,
+        description="SQL statements to run after connecting (e.g., loading extensions, attaching catalogs)",
+    )
 
 
 class PostgreSQLConnection(BaseModel):
@@ -186,7 +190,7 @@ def resolve_paths(self, base_dir: Path | None = None) -> "SidemanticConfig":
             db_p = Path(connection.path)
             if not db_p.is_absolute():
                 db_p = (base / db_p).resolve()
-            connection = DuckDBConnection(type="duckdb", path=str(db_p))
+            connection = DuckDBConnection(type="duckdb", path=str(db_p), init_sql=connection.init_sql)
 
         return SidemanticConfig(
             models_dir=str(models_path),
@@ -263,6 +267,20 @@ def find_config(start_dir: Path | None = None) -> Path | None:
     return None
 
 
+def get_init_sql(config: SidemanticConfig) -> list[str] | None:
+    """Get init_sql statements from config, if any.
+
+    Args:
+        config: Sidemantic configuration
+
+    Returns:
+        List of SQL statements to run after connecting, or None
+    """
+    if config.connection and isinstance(config.connection, DuckDBConnection):
+        return config.connection.init_sql
+    return None
+
+
 def build_connection_string(config: SidemanticConfig) -> str:
     """Build database connection string from config.
 

sidemantic/core/semantic_layer.py

@@ -24,6 +24,7 @@ def __init__(
         use_preaggregations: bool = False,
         preagg_database: str | None = None,
         preagg_schema: str | None = None,
+        init_sql: list[str] | None = None,
     ):
         """Initialize semantic layer.
 
@@ -45,6 +46,8 @@ def __init__(
             use_preaggregations: Enable automatic pre-aggregation routing (default: False)
             preagg_database: Optional database name for pre-aggregation tables
             preagg_schema: Optional schema name for pre-aggregation tables
+            init_sql: SQL statements to run after connecting (DuckDB only, e.g.,
+                loading extensions, attaching catalogs, creating secrets)
         """
         from sidemantic.db.base import BaseDatabaseAdapter
 
@@ -66,10 +69,14 @@ def __init__(
 
                 self.adapter = MotherDuckAdapter.from_url(connection)
                 self.dialect = dialect or "duckdb"
+                # Run init_sql after MotherDuck connection
+                if init_sql:
+                    for stmt in init_sql:
+                        self.adapter.execute(stmt)
             elif connection.startswith("duckdb://"):
                 from sidemantic.db.duckdb import DuckDBAdapter
 
-                self.adapter = DuckDBAdapter.from_url(connection)
+                self.adapter = DuckDBAdapter.from_url(connection, init_sql=init_sql)
                 self.dialect = dialect or "duckdb"
             elif connection.startswith(("postgres://", "postgresql://")):
                 from sidemantic.db.postgres import PostgreSQLAdapter

sidemantic/db/duckdb.py

@@ -14,11 +14,21 @@ class DuckDBAdapter(BaseDatabaseAdapter):
     Wraps DuckDB connection to provide unified adapter interface.
     """
 
-    def __init__(self, path: str = ":memory:", read_only: bool = False, config: dict[str, Any] | None = None):
+    def __init__(
+        self,
+        path: str = ":memory:",
+        read_only: bool = False,
+        config: dict[str, Any] | None = None,
+        init_sql: list[str] | None = None,
+    ):
         """Initialize DuckDB adapter.
 
         Args:
             path: Database file path or ":memory:" for in-memory database
+            read_only: Open database in read-only mode
+            config: DuckDB configuration options
+            init_sql: SQL statements to run immediately after connecting
+                (e.g., loading extensions, attaching catalogs, creating secrets)
         """
         if not read_only and config is None:
             self.conn = duckdb.connect(path)
@@ -27,6 +37,10 @@ def __init__(self, path: str = ":memory:", read_only: bool = False, config: dict
         else:
             self.conn = duckdb.connect(path, read_only=read_only, config=config)
 
+        if init_sql:
+            for stmt in init_sql:
+                self.conn.execute(stmt)
+
     def execute(self, sql: str) -> Any:
         """Execute SQL and return DuckDB relation."""
         return self.conn.execute(sql)
@@ -88,11 +102,12 @@ def raw_connection(self) -> Any:
         return self.conn
 
     @classmethod
-    def from_url(cls, url: str) -> "DuckDBAdapter":
+    def from_url(cls, url: str, init_sql: list[str] | None = None) -> "DuckDBAdapter":
         """Create adapter from connection URL.
 
         Args:
             url: Connection URL (e.g., "duckdb:///:memory:" or "duckdb:///path/to/db.duckdb")
+            init_sql: SQL statements to run after connecting (overrides any init_sql in URL params)
 
         Returns:
             DuckDBAdapter instance
@@ -117,6 +132,7 @@ def from_url(cls, url: str) -> "DuckDBAdapter":
         query = parse_qs(parsed.query)
         read_only = False
         config: dict[str, Any] = {}
+        url_init_sql: list[str] | None = None
 
         def parse_value(value: str) -> Any:
             lowered = value.lower()
@@ -132,10 +148,15 @@ def parse_value(value: str) -> Any:
         for key, values in query.items():
             if not values:
                 continue
+            if key == "init_sql":
+                url_init_sql = values
+                continue
             value = values[-1]
             if key == "read_only":
                 read_only = bool(parse_value(value))
             else:
                 config[key] = parse_value(value)
 
-        return cls(db_path, read_only=read_only, config=config or None)
+        # Parameter init_sql overrides any init_sql found in URL query params
+        effective_init_sql = init_sql if init_sql is not None else url_init_sql
+        return cls(db_path, read_only=read_only, config=config or None, init_sql=effective_init_sql)

tests/db/test_duckdb_adapter.py

@@ -189,3 +189,83 @@ def test_duckdb_valid_schema_names_accepted(schema):
     # Should not raise
     columns = adapter.get_columns("orders", schema=schema)
     assert len(columns) >= 1
+
+
+def test_duckdb_adapter_init_sql():
+    """Test that init_sql statements run after connecting."""
+    adapter = DuckDBAdapter(
+        ":memory:",
+        init_sql=["CREATE TABLE setup_test (id INT, name VARCHAR)"],
+    )
+    result = adapter.execute("SELECT * FROM setup_test")
+    assert result.fetchall() == []
+
+    # Verify columns were created
+    columns = adapter.get_columns("setup_test")
+    col_names = {c["column_name"] for c in columns}
+    assert col_names == {"id", "name"}
+
+
+def test_duckdb_adapter_init_sql_multiple_statements():
+    """Test multiple init_sql statements execute in order."""
+    adapter = DuckDBAdapter(
+        ":memory:",
+        init_sql=[
+            "CREATE TABLE t1 (x INT)",
+            "INSERT INTO t1 VALUES (42)",
+            "CREATE TABLE t2 AS SELECT x * 2 AS doubled FROM t1",
+        ],
+    )
+    result = adapter.execute("SELECT doubled FROM t2")
+    assert result.fetchone()[0] == 84
+
+
+def test_duckdb_adapter_init_sql_none():
+    """Test that init_sql=None is fine (no-op)."""
+    adapter = DuckDBAdapter(":memory:", init_sql=None)
+    result = adapter.execute("SELECT 1")
+    assert result.fetchone()[0] == 1
+
+
+def test_duckdb_adapter_from_url_with_init_sql():
+    """Test from_url passes init_sql through."""
+    adapter = DuckDBAdapter.from_url(
+        "duckdb:///:memory:",
+        init_sql=["CREATE TABLE url_test (val INT)"],
+    )
+    result = adapter.execute("SELECT COUNT(*) FROM url_test")
+    assert result.fetchone()[0] == 0
+
+
+def test_duckdb_adapter_from_url_init_sql_in_query_params():
+    """Test init_sql can be passed via URL query parameters."""
+    adapter = DuckDBAdapter.from_url("duckdb:///:memory:?init_sql=CREATE+TABLE+qs_test+(id+INT)")
+    result = adapter.execute("SELECT COUNT(*) FROM qs_test")
+    assert result.fetchone()[0] == 0
+
+
+def test_duckdb_adapter_from_url_param_overrides_query():
+    """Test that explicit init_sql parameter overrides URL query params."""
+    adapter = DuckDBAdapter.from_url(
+        "duckdb:///:memory:?init_sql=CREATE+TABLE+url_table+(id+INT)",
+        init_sql=["CREATE TABLE param_table (id INT)"],
+    )
+    # param_table should exist (from explicit param)
+    result = adapter.execute("SELECT COUNT(*) FROM param_table")
+    assert result.fetchone()[0] == 0
+
+    # url_table should NOT exist (overridden)
+    with pytest.raises(Exception):
+        adapter.execute("SELECT COUNT(*) FROM url_table")
+
+
+def test_duckdb_semantic_layer_init_sql():
+    """Test init_sql flows through SemanticLayer constructor."""
+    from sidemantic import SemanticLayer
+
+    layer = SemanticLayer(
+        connection="duckdb:///:memory:",
+        init_sql=["CREATE TABLE layer_test (id INT)"],
+    )
+    result = layer.adapter.execute("SELECT COUNT(*) FROM layer_test")
+    assert result.fetchone()[0] == 0

tests/db/test_duckdb_init_sql_config.py

@@ -0,0 +1,96 @@
+"""Tests for init_sql config parsing and plumbing."""
+
+from sidemantic.config import (
+    DuckDBConnection,
+    SidemanticConfig,
+    build_connection_string,
+    get_init_sql,
+    load_config,
+)
+
+
+def test_duckdb_connection_init_sql_field():
+    """Test DuckDBConnection accepts init_sql."""
+    conn = DuckDBConnection(path=":memory:", init_sql=["LOAD httpfs", "LOAD iceberg"])
+    assert conn.init_sql == ["LOAD httpfs", "LOAD iceberg"]
+
+
+def test_duckdb_connection_init_sql_none_by_default():
+    """Test DuckDBConnection init_sql defaults to None."""
+    conn = DuckDBConnection(path=":memory:")
+    assert conn.init_sql is None
+
+
+def test_get_init_sql_returns_statements():
+    """Test get_init_sql extracts init_sql from config."""
+    config = SidemanticConfig(
+        connection=DuckDBConnection(
+            path=":memory:",
+            init_sql=["INSTALL httpfs", "LOAD httpfs"],
+        )
+    )
+    result = get_init_sql(config)
+    assert result == ["INSTALL httpfs", "LOAD httpfs"]
+
+
+def test_get_init_sql_returns_none_when_no_init_sql():
+    """Test get_init_sql returns None when no init_sql configured."""
+    config = SidemanticConfig(connection=DuckDBConnection(path=":memory:"))
+    assert get_init_sql(config) is None
+
+
+def test_get_init_sql_returns_none_for_non_duckdb():
+    """Test get_init_sql returns None for non-DuckDB connections."""
+    config = SidemanticConfig(connection=None)
+    assert get_init_sql(config) is None
+
+
+def test_load_config_with_init_sql(tmp_path):
+    """Test loading YAML config with init_sql."""
+    config_path = tmp_path / "sidemantic.yaml"
+    config_path.write_text(
+        """
+models_dir: .
+connection:
+  type: duckdb
+  path: ":memory:"
+  init_sql:
+    - "INSTALL httpfs"
+    - "LOAD httpfs"
+    - "ATTACH 's3://bucket/db.duckdb' AS remote"
+"""
+    )
+    config = load_config(config_path)
+    assert config.connection is not None
+    init_sql = get_init_sql(config)
+    assert init_sql == [
+        "INSTALL httpfs",
+        "LOAD httpfs",
+        "ATTACH 's3://bucket/db.duckdb' AS remote",
+    ]
+
+
+def test_build_connection_string_ignores_init_sql():
+    """Test that build_connection_string produces URL without init_sql."""
+    config = SidemanticConfig(
+        connection=DuckDBConnection(
+            path=":memory:",
+            init_sql=["LOAD httpfs"],
+        )
+    )
+    url = build_connection_string(config)
+    assert url == "duckdb:///:memory:"
+    assert "init_sql" not in url
+
+
+def test_resolve_paths_preserves_init_sql(tmp_path):
+    """Test that resolve_paths keeps init_sql when resolving DuckDB paths."""
+    config = SidemanticConfig(
+        connection=DuckDBConnection(
+            path="data/warehouse.db",
+            init_sql=["LOAD httpfs"],
+        )
+    )
+    resolved = config.resolve_paths(tmp_path)
+    assert isinstance(resolved.connection, DuckDBConnection)
+    assert resolved.connection.init_sql == ["LOAD httpfs"]