This is an automated email from the ASF dual-hosted git repository.
pankajastro pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3c19778c2ec Allow DataSourceConfig to represent a plain database table
(#73273)
3c19778c2ec is described below
commit 3c19778c2ec4ad8e6ec7155ab56315c423090a77
Author: Pankaj Singh <[email protected]>
AuthorDate: Sat Sep 19 09:02:24 2026 +0530
Allow DataSourceConfig to represent a plain database table (#73273)
LLMSchemaCompareOperator accepts a DataSourceConfig with only conn_id
and table_name, introspected via DbApiHook instead of DataFusion when
the connection resolves to one, but construction failed for that
shape. Also hoist the table_name check so it can no longer be skipped
by an explicit storage_type, and require a uri whenever format is set,
since format alone can't be registered without one.
Co-authored-by: Claude Sonnet 5 <[email protected]>
---
.../ai/docs/operators/llm_schema_compare.rst | 15 ++++++++----
.../ai/example_dags/example_llm_schema_compare.py | 7 ++++--
.../sql/src/airflow/providers/common/sql/config.py | 16 ++++++++++---
.../common/sql/datafusion/test_format_handlers.py | 2 +-
.../sql/tests/unit/common/sql/test_config.py | 28 ++++++++++++++++++++++
5 files changed, 57 insertions(+), 11 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index e7ae94e2c9b..2efed391b01 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -57,13 +57,17 @@ in the schema context sent to the LLM.
:start-after: [START howto_operator_llm_schema_compare_full]
:end-before: [END howto_operator_llm_schema_compare_full]
-With Object Storage
--------------------
+With Object Storage or a Database Table
+---------------------------------------
Use ``data_sources`` with
:class:`~airflow.providers.common.sql.config.DataSourceConfig` to include
object-storage sources (S3 Parquet, CSV, Iceberg, etc.) in the comparison.
-These can be freely combined with ``db_conn_ids``:
+These can be freely combined with ``db_conn_ids``. Whether an entry is
+introspected via ``DbApiHook`` or DataFusion depends on what its ``conn_id``
+resolves to, not on its ``uri``/``format`` fields — a ``DataSourceConfig``
+with neither ``uri`` nor ``format`` set only works when ``conn_id`` resolves
+to a ``DbApiHook``:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
:language: python
@@ -188,8 +192,9 @@ Parameters
- ``db_conn_ids``: List of database connection IDs to compare. Each must
resolve
to a ``DbApiHook``.
- ``table_names``: Tables to introspect from each ``db_conn_id``.
-- ``data_sources``: List of ``DataSourceConfig`` objects for object-storage or
- catalog-managed sources.
+- ``data_sources``: List of ``DataSourceConfig`` objects for object-storage
+ or catalog-managed sources. An entry with neither ``uri`` nor ``format``
+ set works only if its ``conn_id`` resolves to a ``DbApiHook``.
- ``context_strategy``: To fetch primary keys, foreign keys, and
indexes.``full`` or ``basic``,
strongly recommended for cross-system comparisons. default is ``full``
- ``require_approval``: If ``True``, the task pauses after the comparison and
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
index bf8020e684e..06caf979bbd 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
@@ -72,14 +72,17 @@ def example_llm_schema_compare_with_object_storage():
uri="s3://data-lake/customers/",
format="parquet",
)
+ # snowflake_default resolves to a DbApiHook, so no uri/format is needed;
+ # a plain database table can be listed alongside object-storage sources.
+ snowflake_source = DataSourceConfig(conn_id="snowflake_default",
table_name="customers")
LLMSchemaCompareOperator(
task_id="compare_s3_vs_db",
- prompt="Compare S3 Parquet schema against the Postgres table and flag
breaking changes",
+ prompt="Compare S3 Parquet schema against the Postgres and Snowflake
tables and flag breaking changes",
llm_conn_id="pydanticai_default",
db_conn_ids=["postgres_default"],
table_names=["customers"],
- data_sources=[s3_source],
+ data_sources=[s3_source, snowflake_source],
)
diff --git a/providers/common/sql/src/airflow/providers/common/sql/config.py
b/providers/common/sql/src/airflow/providers/common/sql/config.py
index 75cc15efedc..975bc58c9d4 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/config.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/config.py
@@ -62,6 +62,9 @@ class DataSourceConfig:
require ``uri`` or ``storage_type``; they use ``conn_id`` and
format-specific
keys in ``options`` (e.g. ``catalog_table_name`` for Iceberg).
+ **Plain database tables** (neither ``uri`` nor ``format`` set) are not
+ object-store backed either; ``storage_type`` stays ``None`` and is not
inferred.
+
:param conn_id: The connection ID to use for accessing the data source.
:param uri: The URI of the data source (e.g., file path, S3 bucket, etc.).
Not required for catalog-managed formats.
@@ -88,17 +91,24 @@ class DataSourceConfig:
return bool(self.format and self.format.lower() in TABLE_PROVIDERS)
def __post_init__(self):
+ if not self.table_name or not self.table_name.strip():
+ raise ValueError("Table name must be provided for storage type")
+
if self.is_table_provider:
if self.db_name is None:
raise ValueError(f"Database name must be provided for table
providers {TABLE_PROVIDERS}")
return
+ if not self.format and not self.uri:
+ # Plain database table: no object store involved, so storage_type
stays unset.
+ return
+
+ if not self.uri:
+ raise ValueError("URI must be provided when format is set")
+
if self.storage_type is None:
self.storage_type = self._extract_storage_type
- if self.storage_type is not None and (not self.table_name or not
self.table_name.strip()):
- raise ValueError("Table name must be provided for storage type")
-
@property
def _extract_storage_type(self) -> StorageType | None:
"""Extract storage type."""
diff --git
a/providers/common/sql/tests/unit/common/sql/datafusion/test_format_handlers.py
b/providers/common/sql/tests/unit/common/sql/datafusion/test_format_handlers.py
index c6221c62bf4..39ec10532d8 100644
---
a/providers/common/sql/tests/unit/common/sql/datafusion/test_format_handlers.py
+++
b/providers/common/sql/tests/unit/common/sql/datafusion/test_format_handlers.py
@@ -141,7 +141,7 @@ class TestFormatHandlers:
),
(
{"table_name": "t", "format": "parquet", "conn_id": "c"},
- "Unsupported storage type for URI:",
+ "URI must be provided when format is set",
),
],
)
diff --git a/providers/common/sql/tests/unit/common/sql/test_config.py
b/providers/common/sql/tests/unit/common/sql/test_config.py
index 0e90c873443..5478b5a2e44 100644
--- a/providers/common/sql/tests/unit/common/sql/test_config.py
+++ b/providers/common/sql/tests/unit/common/sql/test_config.py
@@ -44,10 +44,38 @@ class TestDataSourceConfig:
with pytest.raises(ValueError, match="Unsupported storage type for
URI"):
DataSourceConfig(conn_id="test", uri="unknown://bucket/path",
table_name="a_table")
+ def test_plain_db_table_without_uri_does_not_infer_storage_type(self):
+ config = DataSourceConfig(conn_id="postgres_default",
table_name="my_table")
+ assert config.storage_type is None
+ assert config.uri == ""
+
+ def test_plain_db_table_with_blank_table_name_raises_error(self):
+ with pytest.raises(ValueError, match="Table name must be provided for
storage type"):
+ DataSourceConfig(conn_id="postgres_default", table_name=" ")
+
def test_missing_table_name_raises_error(self):
with pytest.raises(ValueError, match="Table name must be provided for
storage type"):
DataSourceConfig(conn_id="test", uri="s3://bucket/path",
table_name="")
+ def test_missing_table_name_raises_error_with_explicit_storage_type(self):
+ with pytest.raises(ValueError, match="Table name must be provided for
storage type"):
+ DataSourceConfig(conn_id="pg", table_name="",
storage_type=StorageType.S3)
+
+ def test_iceberg_with_blank_table_name_raises_error(self):
+ with pytest.raises(ValueError, match="Table name must be provided for
storage type"):
+ DataSourceConfig(conn_id="iceberg_conn", table_name="",
format="iceberg", db_name="default")
+
+ def test_explicit_storage_type_without_uri_or_format_is_accepted(self):
+ config = DataSourceConfig(conn_id="pg", table_name="my_table",
storage_type=StorageType.S3)
+ assert config.storage_type == StorageType.S3
+ assert config.uri == ""
+
+ def test_format_without_uri_raises_error(self):
+ with pytest.raises(ValueError, match="URI must be provided when format
is set"):
+ DataSourceConfig(
+ conn_id="test", table_name="my_table", format="parquet",
storage_type=StorageType.LOCAL
+ )
+
def test_parquet_with_partition_cols(self):
config = DataSourceConfig(
conn_id="test_conn",
