feat: ducklake materialization for data pipelines (#9689)

python-client/wmill/wmill/client.py

@@ -2289,6 +2289,131 @@ class DucklakeClient:
             )
         )
 
+    def _qualified(self, table: str, schema: str = None) -> str:
+        return f'dl."{schema}"."{table}"' if schema else f"dl.{table}"
+
+    def _materialize_finish(self, sql, table, schema, partition, partition_col):
+        """Return the materialize query; in a pipeline (WM_PIPELINE) append a
+        summary read and record materialized_partition state after a successful
+        run so SDK-materialized slices appear in the grid like `// materialize`
+        ones. Outside a pipeline it stays a plain query (no recording)."""
+        bind = {} if partition is None else {"_wm_partition": partition}
+        if os.environ.get("WM_PIPELINE") != "true":
+            return self.query(sql, **bind)
+        t = self._qualified(table, schema)
+        where = f" WHERE {partition_col} = $_wm_partition" if partition is not None else ""
+        summary = (
+            f"\nSELECT (SELECT count(*) FROM {t}{where}) AS rows, "
+            f"(SELECT max(snapshot_id) FROM ducklake_snapshots('dl')) AS snapshot_id;"
+        )
+        q = self.query(sql + summary, **bind)
+        # Asset path mirrors the `// materialize` engine: <lake>/<schema>.<table>
+        # for an explicit schema, else <lake>/<table>. Dropping the schema would
+        # hide the row from the grid and collide distinct schemas under one key.
+        asset_path = f"{self.name}/{schema}.{table}" if schema else f"{self.name}/{table}"
+        return _RecordingSqlQuery(q, self.client, asset_path, partition or "")
+
+    def upsert_partition(
+        self,
+        table: str,
+        select_sql: str,
+        partition: str = None,
+        unique_key: str = None,
+        partition_col: str = "_wm_partition",
+        schema: str = None,
+    ):
+        """Idempotently materialize the rows of `select_sql` into ducklake
+        `table` for one `partition` (or the whole table when `partition` is
+        None). Client-side equivalent of the `// materialize` engine: with
+        `unique_key` it upserts within the slice (delete-by-key + insert);
+        without it, it replaces (whole table → CREATE OR REPLACE; partition →
+        delete the partition + insert). Re-running the same slice is safe — the
+        backfill / failure-recovery contract.
+
+        The partition value is bound as a DuckDB arg (never string-interpolated)
+        so it cannot inject SQL. `select_sql` is trusted (your own query).
+        """
+        t = self._qualified(table, schema)
+        # Whole-table (no partition): no partition column; replace rebuilds the
+        # table with CREATE OR REPLACE, merge upserts the whole table by key.
+        if partition is None:
+            if unique_key:
+                sql = (
+                    f"CREATE TABLE IF NOT EXISTS {t} AS SELECT * FROM ({select_sql}) WHERE false;\n"
+                    f"BEGIN TRANSACTION;\n"
+                    f"DELETE FROM {t} WHERE {unique_key} IN (SELECT {unique_key} FROM ({select_sql}));\n"
+                    f"INSERT INTO {t} SELECT * FROM ({select_sql});\n"
+                    f"COMMIT;"
+                )
+            else:
+                sql = f"CREATE OR REPLACE TABLE {t} AS SELECT * FROM ({select_sql});"
+            return self._materialize_finish(sql, table, schema, partition, partition_col)
+        src = f"SELECT *, $_wm_partition AS {partition_col} FROM ({select_sql})"
+        if unique_key:
+            # Upsert via delete-by-key + insert (not MERGE — DuckLake's MERGE
+            # fails writing the first rows of a fresh partition).
+            body = (
+                f"DELETE FROM {t} WHERE {partition_col} = $_wm_partition "
+                f"AND {unique_key} IN (SELECT {unique_key} FROM ({select_sql}));\n"
+                f"INSERT INTO {t} {src};"
+            )
+        else:
+            body = (
+                f"DELETE FROM {t} WHERE {partition_col} = $_wm_partition;\n"
+                f"INSERT INTO {t} {src};"
+            )
+        sql = (
+            f"CREATE TABLE IF NOT EXISTS {t} AS "
+            f"SELECT *, CAST(NULL AS VARCHAR) AS {partition_col} FROM ({select_sql}) WHERE false;\n"
+            f"ALTER TABLE {t} SET PARTITIONED BY ({partition_col});\n"
+            f"BEGIN TRANSACTION;\n{body}\nCOMMIT;"
+        )
+        return self._materialize_finish(sql, table, schema, partition, partition_col)
+
+    def append_partition(
+        self,
+        table: str,
+        select_sql: str,
+        partition: str = None,
+        partition_col: str = "_wm_partition",
+        schema: str = None,
+    ):
+        """INSERT-only materialization (no dedup / no replace) for an immutable
+        event-log table — for one `partition`, or the whole table when
+        `partition` is None. NOTE: unlike `upsert_partition`, re-running the same
+        slice duplicates rows — use only for append-only sources."""
+        t = self._qualified(table, schema)
+        # Whole-table (no partition): insert into the bare table, no partition col.
+        if partition is None:
+            sql = (
+                f"CREATE TABLE IF NOT EXISTS {t} AS SELECT * FROM ({select_sql}) WHERE false;\n"
+                f"INSERT INTO {t} SELECT * FROM ({select_sql});"
+            )
+            return self._materialize_finish(sql, table, schema, partition, partition_col)
+        sql = (
+            f"CREATE TABLE IF NOT EXISTS {t} AS "
+            f"SELECT *, CAST(NULL AS VARCHAR) AS {partition_col} FROM ({select_sql}) WHERE false;\n"
+            f"ALTER TABLE {t} SET PARTITIONED BY ({partition_col});\n"
+            f"INSERT INTO {t} SELECT *, $_wm_partition AS {partition_col} FROM ({select_sql});"
+        )
+        return self._materialize_finish(sql, table, schema, partition, partition_col)
+
+    def read(
+        self,
+        table: str,
+        partition: str = None,
+        partition_col: str = "_wm_partition",
+        schema: str = None,
+    ):
+        """Read a materialized ducklake table, optionally a single partition."""
+        t = self._qualified(table, schema)
+        if partition is not None:
+            return self.query(
+                f"SELECT * FROM {t} WHERE {partition_col} = $_wm_partition",
+                _wm_partition=partition,
+            )
+        return self.query(f"SELECT * FROM {t}")
+
 class SqlQuery:
     """Query result handler for DataTable and DuckLake queries."""
 
@@ -2337,6 +2462,60 @@ class SqlQuery:
         """
         self.fetch_one()
 
+
+class _RecordingSqlQuery:
+    """Wraps a ducklake materialize query so that, on a successful run, the
+    trailing summary (row count + snapshot id) is captured and the
+    materialized_partition state is recorded (best-effort). Only used in pipeline
+    context — outside it the helpers return a plain SqlQuery. Mirrors SqlQuery's
+    terminal methods so `.execute()` / `.fetch_one()` behave the same."""
+
+    def __init__(self, inner, client, asset_path, partition):
+        self._inner = inner
+        self._client = client
+        self._asset_path = asset_path
+        self._partition = partition
+        self.sql = inner.sql
+
+    def execute(self):
+        self._run()
+
+    def fetch_one(self):
+        return self._run()
+
+    def fetch(self, result_collection=None):
+        return self._run()
+
+    def _run(self):
+        try:
+            row = self._inner.fetch_one()
+        except Exception as e:
+            self._record("failed", None, None, str(e))
+            raise
+        snap = row.get("snapshot_id") if isinstance(row, dict) else None
+        rows = row.get("rows") if isinstance(row, dict) else None
+        self._record("materialized", snap, rows, None)
+        return row
+
+    def _record(self, status, snapshot_id, row_count, error):
+        try:
+            self._client.post(
+                f"/w/{self._client.workspace}/assets/record_materialization",
+                json={
+                    "asset_kind": "ducklake",
+                    "asset_path": self._asset_path,
+                    "partition": self._partition,
+                    "status": status,
+                    "snapshot_id": snapshot_id,
+                    "row_count": row_count,
+                    "job_id": os.environ.get("WM_JOB_ID"),
+                    "error": error,
+                },
+            )
+        except Exception:
+            pass  # best-effort; never fail the user's materialization
+
+
 def infer_sql_type(value) -> str:
     """
     DuckDB executor requires explicit argument types at declaration