This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 532e7d0911c Allow Common AI SQL imports without DataFusion (#69990)
532e7d0911c is described below
commit 532e7d0911cf109d14581c3f42281b40e08af6aa
Author: Aaron Chen <[email protected]>
AuthorDate: Wed Jul 22 18:13:48 2026 +0800
Allow Common AI SQL imports without DataFusion (#69990)
---
providers/common/ai/README.rst | 4 +-
providers/common/ai/docs/index.rst | 4 +-
providers/common/ai/docs/operators/llm_sql.rst | 6 ++
providers/common/ai/pyproject.toml | 4 +-
.../providers/common/ai/operators/llm_sql.py | 12 ++-
.../tests/unit/common/ai/operators/test_llm_sql.py | 87 +++++++++++++++++++---
6 files changed, 101 insertions(+), 16 deletions(-)
diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst
index f23ac091182..97f24220828 100644
--- a/providers/common/ai/README.rst
+++ b/providers/common/ai/README.rst
@@ -95,8 +95,8 @@ Extra Dependencies
``skills`` ``apache-airflow-providers-git>=0.4.0``,
``pydantic-ai-skills>=0.11.0``
``avro`` ``fastavro>=1.10.0; python_version < "3.14"``,
``fastavro>=1.12.1; python_version >= "3.14"``
``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``,
``pyarrow>=22.0.0; python_version >= '3.14'``
-``sql`` ``apache-airflow-providers-common-sql``, ``sqlglot>=30.0.0``
-``common.sql`` ``apache-airflow-providers-common-sql``
+``sql`` ``apache-airflow-providers-common-sql>=1.33.0``,
``sqlglot>=30.0.0``
+``common.sql`` ``apache-airflow-providers-common-sql>=1.33.0``
``langchain`` ``langchain>=1.0.0``
``llamaindex`` ``llama-index-core>=0.13.0``,
``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.0``
``pdf`` ``pypdf>=4.0.0``
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 51081846490..d1bff2ce31c 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -252,9 +252,9 @@ Extra Dependencies
``skills`` ``apache-airflow-providers-git>=0.4.0``,
``pydantic-ai-skills>=1.2.0``
``avro`` ``fastavro>=1.10.0; python_version < "3.14"``,
``fastavro>=1.12.1; python_version >= "3.14"``
``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``,
``pyarrow>=22.0.0; python_version >= '3.14'``
-``sql`` ``apache-airflow-providers-common-sql``, ``sqlglot>=30.0.0``
+``sql`` ``apache-airflow-providers-common-sql>=1.33.0``,
``sqlglot>=30.0.0``
``aws`` ``apache-airflow-providers-amazon>=9.0.0``
-``common.sql`` ``apache-airflow-providers-common-sql``
+``common.sql`` ``apache-airflow-providers-common-sql>=1.33.0``
``langchain`` ``langchain>=1.0.0``
``llamaindex`` ``dataclasses-json>=0.6.7``, ``llama-index-core>=0.13.0``,
``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.0``
``pdf`` ``pypdf>=4.0.0``
diff --git a/providers/common/ai/docs/operators/llm_sql.rst
b/providers/common/ai/docs/operators/llm_sql.rst
index 807fe5197e8..84d5a2226e1 100644
--- a/providers/common/ai/docs/operators/llm_sql.rst
+++ b/providers/common/ai/docs/operators/llm_sql.rst
@@ -60,6 +60,12 @@ The operator uses
:class:`~airflow.providers.common.sql.config.DataSourceConfig`
to register the object storage source as a table so the LLM can include it in
the schema context.
+.. note::
+
+ Object-storage schema introspection requires the ``datafusion`` extra of
+ ``apache-airflow-providers-common-sql``. Install it with
+ ``pip install "apache-airflow-providers-common-sql[datafusion]"``.
+
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
:language: python
:start-after: [START howto_operator_llm_sql_with_object_storage]
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index 2c1f13eec2a..cbaa9e82ee4 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -105,7 +105,7 @@ dependencies = [
"pyarrow>=22.0.0; python_version >= '3.14'",
]
"sql" = [
- "apache-airflow-providers-common-sql",
+ "apache-airflow-providers-common-sql>=1.33.0",
"sqlglot>=30.0.0",
]
# AWSToolset: allow-listed AWS API access for agents. The amazon provider
@@ -114,7 +114,7 @@ dependencies = [
"apache-airflow-providers-amazon>=9.0.0",
]
"common.sql" = [
- "apache-airflow-providers-common-sql"
+ "apache-airflow-providers-common-sql>=1.33.0"
]
"langchain" = [
"langchain>=1.0.0",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
index d850bf11372..e8a3b327b42 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
@@ -28,7 +28,6 @@ try:
resolve_sqlglot_dialect,
validate_sql as _validate_sql,
)
- from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
except ImportError as e:
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
@@ -226,6 +225,17 @@ class LLMSQLQueryOperator(LLMOperator):
def _introspect_object_storage_schema(self):
"""Use DataFusion Engine to get the schema of object stores."""
+ try:
+ from airflow.providers.common.sql.datafusion.engine import
DataFusionEngine
+ except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException(
+ "Object-storage schema introspection requires the `datafusion`
extra of "
+ "apache-airflow-providers-common-sql. Install it with: "
+ 'pip install "apache-airflow-providers-common-sql[datafusion]"'
+ ) from e
+
engine = DataFusionEngine()
engine.register_datasource(self.datasource_config)
return engine.get_schema(self.datasource_config.table_name)
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
index 812c7f58a36..a6fe35277ad 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
@@ -16,6 +16,8 @@
# under the License.
from __future__ import annotations
+import subprocess
+import sys
from datetime import timedelta
from unittest.mock import MagicMock, PropertyMock, patch
from uuid import uuid4
@@ -58,6 +60,77 @@ def _make_mock_agent(output: str):
return mock_agent
+def _run_python_without_datafusion(code: str) ->
subprocess.CompletedProcess[str]:
+ return subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ f"""
+import builtins
+
+real_import = builtins.__import__
+
+
+def blocked_import(name, *args, **kwargs):
+ if name == "datafusion" or name.startswith("datafusion."):
+ raise ModuleNotFoundError("No module named 'datafusion'")
+ return real_import(name, *args, **kwargs)
+
+
+builtins.__import__ = blocked_import
+{code}
+""",
+ ],
+ capture_output=True,
+ text=True,
+ check=False,
+ )
+
+
+class TestDataFusionOptionalDependency:
+ def test_relational_sql_imports_do_not_require_datafusion(self):
+ result = _run_python_without_datafusion(
+ """
+from airflow.providers.common.ai.decorators.llm_sql import llm_sql_task
+from airflow.providers.common.ai.operators.llm_sql import LLMSQLQueryOperator
+"""
+ )
+
+ assert result.returncode == 0, result.stderr
+
+ def test_object_storage_error_explains_how_to_install_datafusion(self):
+ result = _run_python_without_datafusion(
+ """
+from airflow.providers.common.ai.operators.llm_sql import LLMSQLQueryOperator
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+from airflow.providers.common.sql.config import DataSourceConfig
+
+operator = LLMSQLQueryOperator(
+ task_id="test",
+ prompt="test",
+ llm_conn_id="test",
+ datasource_config=DataSourceConfig(
+ conn_id="aws_default",
+ table_name="sales",
+ uri="s3://bucket/sales/",
+ format="parquet",
+ ),
+)
+
+try:
+ operator._introspect_object_storage_schema()
+except AirflowOptionalProviderFeatureException as error:
+ expected = 'pip install "apache-airflow-providers-common-sql[datafusion]"'
+ if expected not in str(error):
+ raise AssertionError(f"Missing installation guidance in: {{error}}")
from error
+else:
+ raise AssertionError("Expected object-storage introspection to require
DataFusion")
+"""
+ )
+
+ assert result.returncode == 0, result.stderr
+
+
class TestStripLLMOutput:
@pytest.mark.parametrize(
("raw", "expected"),
@@ -277,7 +350,7 @@ class TestLLMSQLQueryOperatorSchemaIntrospection:
assert op._get_schema_context() == "My custom schema info"
@patch(
- "airflow.providers.common.ai.operators.llm_sql.DataFusionEngine",
+ "airflow.providers.common.sql.datafusion.engine.DataFusionEngine",
autospec=True,
)
def test_introspect_object_storage_schema(self, mock_engine_cls):
@@ -305,7 +378,7 @@ class TestLLMSQLQueryOperatorSchemaIntrospection:
assert result == schema_text
@patch(
- "airflow.providers.common.ai.operators.llm_sql.DataFusionEngine",
+ "airflow.providers.common.sql.datafusion.engine.DataFusionEngine",
autospec=True,
)
def test_introspect_schemas_with_db_and_datasource_config(self,
mock_engine_cls):
@@ -343,7 +416,7 @@ class TestLLMSQLQueryOperatorSchemaIntrospection:
assert object_schema in result
@patch(
- "airflow.providers.common.ai.operators.llm_sql.DataFusionEngine",
+ "airflow.providers.common.sql.datafusion.engine.DataFusionEngine",
autospec=True,
)
def test_introspect_schemas_datasource_config_without_db_tables(self,
mock_engine_cls):
@@ -374,11 +447,7 @@ class TestLLMSQLQueryOperatorSchemaIntrospection:
assert "Table: s3_data" in result
assert "ts: TIMESTAMP\nvalue: DOUBLE" in result
- @patch(
- "airflow.providers.common.ai.operators.llm_sql.DataFusionEngine",
- autospec=True,
- )
- def test_introspect_schemas_raises_when_no_tables_and_no_datasource(self,
mock_engine_cls):
+ def test_introspect_schemas_raises_when_no_tables_and_no_datasource(self):
"""ValueError is raised when no db tables return schema and no
datasource_config is set."""
mock_db_hook = MagicMock(spec=["get_table_schema", "dialect_name"])
mock_db_hook.get_table_schema.return_value = []
@@ -397,7 +466,7 @@ class TestLLMSQLQueryOperatorSchemaIntrospection:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@patch(
- "airflow.providers.common.ai.operators.llm_sql.DataFusionEngine",
+ "airflow.providers.common.sql.datafusion.engine.DataFusionEngine",
autospec=True,
)
def test_execute_with_datasource_config_and_db_tables(self,
mock_engine_cls, mock_hook_cls):
