OpenLineage: record sql_metadata for Polars database I/O

docs/reference/lifecycle-hooks/OpenLineageAdapter.rst

@@ -18,16 +18,18 @@ beyond the client.
 SQL datasets
 ------------
 
-SQL loaders and savers (``@load_from.sql``, ``@save_to.sql`` and the pandas SQL materializers) record
-the datasource they used (see :ref:`sql-metadata-and-lineage`). How the adapter names their datasets
-is set by ``sql_dataset_identity``:
+SQL loaders and savers (``@load_from.sql``, ``@save_to.sql``, the pandas SQL materializers and the
+Polars database materializers) record the datasource they used (see :ref:`sql-metadata-and-lineage`).
+How the adapter names their datasets is set by ``sql_dataset_identity``:
 
 - ``"legacy"``, the default: datasets are named as in earlier Hamilton releases, under the adapter's
   *job* namespace with the bare ``table_name``. As before, a query read containing ``SELECT``
   produces a dataset with no name and the query in the job's ``sql`` facet, and other strings are
   used as the dataset name. Leaving the option unset warns once per adapter (a
   ``FutureWarning``) because the default will change; pass ``"legacy"`` explicitly to keep these names
-  without the warning.
+  without the warning. The Polars database materializers keep their earlier names too: the read is a
+  dataset named after the loader node, with the query as its ``dataSource`` URI, and the write is
+  named after the table. Unlike the pandas query read above, the read has no ``sql`` job facet.
 - ``"datasource"``: datasets are named after the datasource, following the `OpenLineage naming
   conventions <https://openlineage.io/docs/spec/naming/>`_, and every physical table a query reads is
   reported. A report written by one job and read by another then resolves to the same dataset, and
@@ -37,6 +39,9 @@ is set by ``sql_dataset_identity``:
 
     adapter = OpenLineageAdapter(client, "my_namespace", "my_job", sql_dataset_identity="datasource")
 
+Datasource mode for the Polars database materializers requires ``polars>=1.0``. With an older
+Polars they log a warning asking you to upgrade and emit no dataset.
+
 The rest of this section describes the ``"datasource"`` identity:
 
 .. list-table::
@@ -68,6 +73,11 @@ query. A read by another name that is not a plain identifier (``daily revenue``)
 that table when no SQL statement starts at it; one starting with ``table`` is left out, since
 Postgres' ``TABLE t`` is a statement ``openlineage-sql`` does not know.
 
+A write is handled alike, for the pandas and Polars writers: a table name that is not a plain identifier is
+parsed, and is the table written unless it parses as a statement that names tables, such as a table literally
+called ``INSERT INTO orders SELECT * FROM customers``. The tables that statement writes are emitted instead, or
+none when it only reads. Give such tables plain names.
+
 Whatever cannot be fully identified is left out and logged as a warning from the
 ``hamilton.plugins.h_openlineage`` logger, never guessed. The following cases are left out:
 

hamilton/plugins/polars_post_1_0_0_extensions.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import csv
 import dataclasses
 from collections.abc import Collection, Mapping, Sequence
 from io import BytesIO, IOBase, TextIOWrapper
@@ -807,7 +808,14 @@ def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]:
             connection=self.connection,
             **self._get_loading_kwargs(),
         )
-        metadata = utils.get_file_and_dataframe_metadata(self.query, df)
+        # with iter_batches=True, pl.read_database returns a generator, so there is no row count
+        rows = len(df) if isinstance(df, pl.DataFrame) else None
+        metadata = {
+            **utils.get_file_and_dataframe_metadata(self.query, df),
+            **utils.get_sql_metadata(
+                self.query, rows, db_connection=self.connection, operation="read"
+            ),
+        }
         return df, metadata
 
     @classmethod
@@ -842,12 +850,26 @@ def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]:
         if isinstance(data, pl.LazyFrame):
             data = data.collect()
 
-        data.write_database(
+        rows = data.write_database(
             table_name=self.table_name,
             connection=self.connection,
             **self._get_saving_kwargs(),
         )
-        return utils.get_file_and_dataframe_metadata(self.table_name, data)
+        # polars unpacks table_name as [[catalog.]schema.]table, honouring quotes; lineage wants
+        # the schema recorded apart, and the connection already identifies the catalog
+        components = next(csv.reader([self.table_name], delimiter="."))
+        table = components[-1]
+        schema = components[-2] if len(components) > 1 else None
+        return {
+            **utils.get_file_and_dataframe_metadata(self.table_name, data),
+            **utils.get_sql_metadata(
+                table,
+                rows,
+                db_connection=self.connection,
+                schema=schema or None,
+                operation="write",
+            ),
+        }
 
     @classmethod
     def name(cls) -> str:

hamilton/plugins/polars_pre_1_0_0_extension.py

@@ -47,6 +47,7 @@
     from polars.type_aliases import CsvEncoding, SchemaDefinition
 else:
     CsvEncoding = type
+    SchemaDefinition = type
 if has_alias and hasattr(pl.type_aliases, "CsvQuoteStyle"):
     from polars.type_aliases import CsvQuoteStyle
 else:
@@ -729,6 +730,17 @@ def name(cls) -> str:
         return "spreadsheet"
 
 
+def _sql_metadata_without_datasource(
+    query_or_table: str, rows: int | None, operation: Literal["read", "write"]
+) -> dict[str, Any]:
+    """``sql_metadata`` for database I/O, whose datasource polars<1.0 cannot be asked for."""
+    metadata = utils.get_sql_metadata(query_or_table, rows, operation=operation)
+    metadata["sql_metadata"]["notes"] = (
+        "Datasource lineage for Polars database I/O requires polars>=1.0; upgrade polars"
+    )
+    return metadata
+
+
 @dataclasses.dataclass
 class PolarsDatabaseReader(DataLoader):
     """
@@ -768,7 +780,12 @@ def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]:
             connection=self.connection,
             **self._get_loading_kwargs(),
         )
-        metadata = utils.get_file_and_dataframe_metadata(self.query, df)
+        # with iter_batches=True, pl.read_database returns a generator, so there is no row count
+        rows = len(df) if isinstance(df, pl.DataFrame) else None
+        metadata = {
+            **utils.get_file_and_dataframe_metadata(self.query, df),
+            **_sql_metadata_without_datasource(self.query, rows, "read"),
+        }
         return df, metadata
 
     @classmethod
@@ -803,12 +820,15 @@ def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]:
         if isinstance(data, pl.LazyFrame):
             data = data.collect()
 
-        data.write_database(
+        rows = data.write_database(
             table_name=self.table_name,
             connection=self.connection,
             **self._get_saving_kwargs(),
         )
-        return utils.get_file_and_dataframe_metadata(self.table_name, data)
+        return {
+            **utils.get_file_and_dataframe_metadata(self.table_name, data),
+            **_sql_metadata_without_datasource(self.table_name, rows, "write"),
+        }
 
     @classmethod
     def name(cls) -> str:

tests/plugins/test_h_openlineage.py

@@ -18,6 +18,7 @@
 import json
 import logging
 import sqlite3
+import typing
 import warnings
 
 import pandas as pd
@@ -36,6 +37,13 @@
 from hamilton.plugins import h_openlineage  # noqa: E402
 from hamilton.plugins.pandas_extensions import PandasSqlReader  # noqa: E402
 
+try:
+    import polars as pl
+
+    from hamilton.plugins import polars_post_1_0_0_extensions as polars_post
+except ImportError:  # pragma: no cover - polars is optional
+    pl = None
+
 REVENUE_QUERY = """
 -- daily revenue per customer country
 WITH paid AS (SELECT * FROM orders WHERE status = 'paid')
@@ -845,3 +853,187 @@ def orders(df: pd.DataFrame) -> pd.DataFrame:
     assert dataset_ids(events, "inputs") == [
         (f"sqlite://{other_path.resolve().as_posix()}", "orders")
     ]
+
+
+requires_polars = pytest.mark.skipif(pl is None, reason="polars is not installed")
+
+
+class SubscriptableStandIn:
+    """Stands in for a name a type alias refers to when sqlalchemy has no class of that name."""
+
+    def __class_getitem__(cls, item):
+        return cls
+
+
+def sqlalchemy_class(name):
+    """The sqlalchemy class polars' alias means by ``name`` (or ``Alchemy<name>``), if any."""
+    import sqlalchemy
+    import sqlalchemy.orm
+
+    modules = [sqlalchemy, sqlalchemy.orm]
+    try:
+        import sqlalchemy.ext.asyncio
+
+        modules.append(sqlalchemy.ext.asyncio)
+    except ImportError:  # the asyncio extension needs greenlet
+        pass
+    for candidate in (name, name.removeprefix("Alchemy")):
+        for module in modules:
+            if hasattr(module, candidate):
+                return getattr(module, candidate)
+    return SubscriptableStandIn
+
+
+@pytest.fixture
+def polars_connection_type_hints(monkeypatch):
+    """Resolves the names polars' ``ConnectionOrCursor`` alias only imports for type checking.
+
+    Hamilton reads the database classes' type hints to build ``@load_from.database`` and
+    ``@save_to.database``, which fails on those unresolved names. They differ between polars
+    versions, so each one that is missing is bound to the sqlalchemy class it means.
+    """
+    for _ in range(50):
+        try:
+            typing.get_type_hints(polars_post.PolarsDatabaseReader)
+            typing.get_type_hints(polars_post.PolarsDatabaseWriter)
+            return
+        except NameError as e:
+            monkeypatch.setattr(polars_post, e.name, sqlalchemy_class(e.name), raising=False)
+    raise AssertionError("polars database type hints still do not resolve")
+
+
+def polars_database_modules():
+    @save_to.database(
+        table_name=value("orders"),
+        connection=source("engine"),
+        if_table_exists=value("replace"),
+        output_name_="saved_orders",
+    )
+    def orders_to_save(raw_orders: pl.DataFrame) -> pl.DataFrame:
+        return raw_orders
+
+    @load_from.database(query=value("SELECT * FROM orders"), connection=source("engine"))
+    def orders_read(orders_df: pl.DataFrame) -> pl.DataFrame:
+        return orders_df
+
+    return (
+        ad_hoc_utils.create_temporary_module(orders_to_save),
+        ad_hoc_utils.create_temporary_module(orders_read),
+    )
+
+
+def run_polars_database_flow(tmp_path, namespace, **adapter_kwargs):
+    """Writes ``orders`` then reads it back through Polars' database materializers."""
+    db_path = tmp_path / "shop.db"
+    inputs = {
+        "engine": create_engine(f"sqlite:///{db_path}"),
+        "raw_orders": pl.DataFrame({"id": [1, 2, 3], "amount": [10.0, 5.0, 7.0]}),
+    }
+    writer, reader = polars_database_modules()
+    run_with_lineage(
+        tmp_path, namespace, inputs, module=writer, final_vars=("saved_orders",), **adapter_kwargs
+    )
+    result, events = run_with_lineage(
+        tmp_path, namespace, inputs, module=reader, final_vars=("orders_read",), **adapter_kwargs
+    )
+    return db_path, result, [strip_private(e) for e in events if e["eventType"] == "RUNNING"]
+
+
+@requires_polars
+def test_polars_database_datasource_identity_names_the_sqlite_table(
+    tmp_path, polars_connection_type_hints
+):
+    db_path, _, (write, read) = run_polars_database_flow(
+        tmp_path, "polars_ns", sql_dataset_identity="datasource"
+    )
+    namespace = f"sqlite://{db_path.resolve().as_posix()}"
+    assert [(d["namespace"], d["name"]) for d in write["outputs"]] == [(namespace, "orders")]
+    assert [(d["namespace"], d["name"]) for d in read["inputs"]] == [(namespace, "orders")]
+    assert read["job"]["facets"]["sql"]["query"] == "SELECT * FROM orders"
+    for dataset in write["outputs"] + read["inputs"]:
+        facets = dataset.get("facets", {})
+        assert "storage" not in facets
+
+
+@requires_polars
+def test_polars_database_default_identity_keeps_legacy_datasets(
+    tmp_path, polars_connection_type_hints
+):
+    """Pinned to ``main``'s datasets: file metadata still wins over the SQL metadata.
+
+    Only the ``FutureWarning`` is new, since the metadata now carries ``sql_metadata``.
+    """
+    with pytest.warns(FutureWarning, match="sql_dataset_identity='datasource'"):
+        _, _, (write, read) = run_polars_database_flow(tmp_path, "polars_legacy_ns")
+    assert (read["inputs"][0]["namespace"], read["inputs"][0]["name"]) == (
+        "polars_legacy_ns",
+        "orders_read.load_data.orders_df",
+    )
+    assert set(read["inputs"][0]["facets"]) == {"storage", "dataSource", "schema"}
+    assert read["inputs"][0]["facets"]["storage"]["storageLayer"] == "FileSystem"
+    assert read["inputs"][0]["facets"]["dataSource"]["uri"] == "SELECT * FROM orders"
+    assert "sql" not in read["job"]["facets"]
+    assert (write["outputs"][0]["namespace"], write["outputs"][0]["name"]) == (
+        "polars_legacy_ns",
+        "orders",
+    )
+    assert set(write["outputs"][0]["facets"]) == {"storage", "dataSource", "schema"}
+    assert write["outputs"][0]["facets"]["storage"]["storageLayer"] == "FileSystem"
+    assert write["outputs"][0]["facets"]["dataSource"]["uri"] == "orders"
+
+
+@requires_polars
+def test_pre_1_0_polars_database_asks_datasource_users_to_upgrade(tmp_path, caplog):
+    """The pre-1.0 classes keep file metadata and add a connection-less ``sql_metadata``."""
+    from hamilton.plugins import polars_pre_1_0_0_extension as pre_1_0
+
+    connection = sqlite3.connect(tmp_path / "old.db")
+    frame = pl.DataFrame({"id": [1, 2]})
+    write_metadata = pre_1_0.PolarsDatabaseWriter(
+        table_name="orders",
+        connection=f"sqlite:///{tmp_path / 'old.db'}",
+        if_table_exists="replace",
+    ).save_data(frame)
+    _, read_metadata = pre_1_0.PolarsDatabaseReader(
+        query="SELECT * FROM orders", connection=connection
+    ).load_data(pl.DataFrame)
+
+    for metadata, operation in ((write_metadata, "write"), (read_metadata, "read")):
+        assert metadata["file_metadata"]["path"] == (
+            "orders" if operation == "write" else "SELECT * FROM orders"
+        )
+        assert metadata["sql_metadata"]["source"] is None
+        assert metadata["sql_metadata"]["operation"] == operation
+        node_ = type("FakeNode", (), {"name": f"{operation}_node"})()
+        caplog.clear()
+        with caplog.at_level(logging.WARNING, logger="hamilton.plugins.h_openlineage"):
+            datasets, _ = h_openlineage._sql_lineage(metadata, operation, node_)
+        assert datasets == []
+        assert "requires polars>=1.0" in caplog.text
+
+
+@requires_polars
+@pytest.mark.parametrize(
+    ("table_name", "query", "name"),
+    [
+        ("main.orders", "SELECT * FROM main.orders", "orders"),
+        ('"my-tbl"', 'SELECT * FROM "my-tbl"', "my-tbl"),
+        ('main."my-tbl"', 'SELECT * FROM main."my-tbl"', "my-tbl"),
+    ],
+)
+def test_polars_database_write_resolves_to_the_table_a_query_reads(
+    tmp_path, table_name, query, name
+):
+    path = tmp_path / "qualified.db"
+    engine = create_engine(f"sqlite:///{path}")
+    written = polars_post.PolarsDatabaseWriter(
+        table_name=table_name, connection=engine, if_table_exists="replace"
+    ).save_data(pl.DataFrame({"id": [1, 2]}))
+    _, read = polars_post.PolarsDatabaseReader(query=query, connection=engine).load_data(
+        pl.DataFrame
+    )
+
+    output = h_openlineage.sql_datasets(written).outputs
+    input_ = h_openlineage.sql_datasets(read).inputs
+    assert identities(output) == identities(input_)
+    assert identities(output) == [(f"sqlite://{path.resolve().as_posix()}", name)]