Add LLMSQLQueryOperator and @task.llm_sql

dev/breeze/tests/test_selective_checks.py

@@ -2249,7 +2249,7 @@ def test_upgrade_to_newer_dependencies(
             ("providers/common/sql/src/airflow/providers/common/sql/common_sql_python.py",),
             {
                 "docs-list-as-string": "amazon apache.drill apache.druid apache.hive "
-                "apache.impala apache.pinot common.compat common.sql databricks elasticsearch "
+                "apache.impala apache.pinot common.ai common.compat common.sql databricks elasticsearch "
                 "exasol google jdbc microsoft.mssql mysql odbc openlineage "
                 "oracle pgvector postgres presto slack snowflake sqlite teradata trino vertica ydb",
             },

providers/common/ai/docs/index.rst

@@ -36,6 +36,7 @@
 
     Connection types <connections/pydantic_ai>
     Hooks <hooks/pydantic_ai>
+    Operators <operators/llm_sql>
 
 .. toctree::
     :hidden:
@@ -120,6 +121,7 @@ You can install such cross-provider dependencies when installing from PyPI. For
 Dependent package                                                                                                   Extra
 ==================================================================================================================  =================
 `apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_  ``common.compat``
+`apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_        ``common.sql``
 ==================================================================================================================  =================
 
 Downloading official packages

providers/common/ai/docs/operators/llm_sql.rst

@@ -0,0 +1,87 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _howto/operator:llm_sql_query:
+
+``LLMSQLQueryOperator``
+======================
+
+Use :class:`~airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator` to generate
+SQL queries from natural language using an LLM.
+
+The operator generates SQL but does not execute it. The generated query is returned
+as XCom and can be passed to ``SQLExecuteQueryOperator`` or used in downstream tasks.
+
+.. seealso::
+    :ref:`Connection configuration <howto/connection:pydantic_ai>`
+
+Basic Usage
+-----------
+
+Provide a natural language ``prompt`` and the operator generates a SQL query:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
+    :language: python
+    :start-after: [START howto_operator_llm_sql_basic]
+    :end-before: [END howto_operator_llm_sql_basic]
+
+With Schema Introspection
+-------------------------
+
+Use ``db_conn_id`` and ``table_names`` to automatically include database schema
+in the LLM's context. This produces more accurate queries because the LLM knows
+the actual column names and types:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
+    :language: python
+    :start-after: [START howto_operator_llm_sql_schema]
+    :end-before: [END howto_operator_llm_sql_schema]
+
+TaskFlow Decorator
+------------------
+
+The ``@task.llm_sql`` decorator lets you write a function that returns the
+prompt. The decorator handles LLM connection, schema introspection, SQL generation,
+and safety validation:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
+    :language: python
+    :start-after: [START howto_decorator_llm_sql]
+    :end-before: [END howto_decorator_llm_sql]
+
+Dynamic Task Mapping
+--------------------
+
+Generate SQL for multiple prompts in parallel using ``expand()``:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
+    :language: python
+    :start-after: [START howto_operator_llm_sql_expand]
+    :end-before: [END howto_operator_llm_sql_expand]
+
+SQL Safety Validation
+---------------------
+
+By default, the operator validates generated SQL using an allowlist approach:
+
+- Only ``SELECT``, ``UNION``, ``INTERSECT``, and ``EXCEPT`` statements are allowed.
+- Multi-statement SQL (semicolon-separated) is rejected.
+- Disallowed statements (``INSERT``, ``UPDATE``, ``DELETE``, ``DROP``, etc.) raise
+  :class:`~airflow.providers.common.ai.utils.sql_validation.SQLSafetyError`.
+
+You can disable validation with ``validate_sql=False`` or customize the allowed
+statement types with ``allowed_sql_types``.

providers/common/ai/provider.yaml

@@ -51,3 +51,12 @@ connection-types:
       placeholders:
         host: "https://api.openai.com/v1 (optional, for custom endpoints)"
         extra: '{"model": "openai:gpt-5"}'
+
+operators:
+  - integration-name: Pydantic AI
+    python-modules:
+      - airflow.providers.common.ai.operators.llm_sql
+
+task-decorators:
+  - class-name: airflow.providers.common.ai.decorators.llm_sql.llm_sql_task
+    name: llm_sql

providers/common/ai/pyproject.toml

@@ -72,13 +72,21 @@ dependencies = [
 "common.compat" = [
     "apache-airflow-providers-common-compat"
 ]
+"sql" = [
+    "apache-airflow-providers-common-sql",
+    "sqlglot>=26.0.0",
+]
+"common.sql" = [
+    "apache-airflow-providers-common-sql"
+]
 
 [dependency-groups]
 dev = [
     "apache-airflow",
     "apache-airflow-task-sdk",
     "apache-airflow-devel-common",
     "apache-airflow-providers-common-compat",
+    "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
 ]
 

providers/common/ai/src/airflow/providers/common/ai/decorators/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

providers/common/ai/src/airflow/providers/common/ai/decorators/llm_sql.py

@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+TaskFlow decorator for LLM SQL generation.
+
+The user writes a function that **returns the prompt**. The decorator handles
+the LLM call, schema introspection, and safety validation. The decorated task's
+XCom output is the generated SQL string.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Collection, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from airflow.providers.common.ai.operators.llm_sql import LLMSQLQueryOperator
+from airflow.providers.common.compat.sdk import (
+    DecoratedOperator,
+    TaskDecorator,
+    context_merge,
+    task_decorator_factory,
+)
+from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION
+from airflow.utils.operator_helpers import determine_kwargs
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+
+
+class _LLMSQLDecoratedOperator(DecoratedOperator, LLMSQLQueryOperator):
+    """
+    Wraps a callable that returns a prompt for LLM SQL generation.
+
+    The user function is called at execution time to produce the prompt string.
+    All other parameters (``llm_conn_id``, ``db_conn_id``, ``table_names``, etc.)
+    are passed through to :class:`~airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator`.
+
+    :param python_callable: A reference to a callable that returns the prompt string.
+    :param op_args: Positional arguments for the callable.
+    :param op_kwargs: Keyword arguments for the callable.
+    """
+
+    template_fields: Sequence[str] = (
+        *DecoratedOperator.template_fields,
+        *LLMSQLQueryOperator.template_fields,
+    )
+    template_fields_renderers: ClassVar[dict[str, str]] = {
+        **DecoratedOperator.template_fields_renderers,
+    }
+
+    custom_operator_name: str = "@task.llm_sql"
+
+    def __init__(
+        self,
+        *,
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            python_callable=python_callable,
+            op_args=op_args,
+            op_kwargs=op_kwargs,
+            prompt=SET_DURING_EXECUTION,
+            **kwargs,
+        )
+
+    def execute(self, context: Context) -> Any:
+        context_merge(context, self.op_kwargs)
+        kwargs = determine_kwargs(self.python_callable, self.op_args, context)
+
+        self.prompt = self.python_callable(*self.op_args, **kwargs)
+
+        if not isinstance(self.prompt, str) or not self.prompt.strip():
+            raise TypeError("The returned value from the @task.llm_sql callable must be a non-empty string.")
+
+        self.render_template_fields(context)
+        # Call LLMSQLQueryOperator.execute directly, not super().execute(),
+        # because we need to skip DecoratedOperator.execute — the callable
+        # invocation is already handled above.
+        return LLMSQLQueryOperator.execute(self, context)
+
+
+def llm_sql_task(
+    python_callable: Callable | None = None,
+    **kwargs,
+) -> TaskDecorator:
+    """
+    Wrap a function that returns a natural language prompt into an LLM SQL task.
+
+    The function body constructs the prompt (can use Airflow context, XCom, etc.).
+    The decorator handles: LLM connection, schema introspection, SQL generation,
+    and safety validation.
+
+    Usage::
+
+        @task.llm_sql(
+            llm_conn_id="openai_default",
+            db_conn_id="postgres_default",
+            table_names=["customers", "orders"],
+        )
+        def build_query(ds=None):
+            return f"Find top 10 customers by revenue in {ds}"
+
+    :param python_callable: Function to decorate.
+    """
+    return task_decorator_factory(
+        python_callable=python_callable,
+        decorated_operator_class=_LLMSQLDecoratedOperator,
+        **kwargs,
+    )

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py

@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Example DAGs demonstrating LLMSQLQueryOperator usage."""
+
+from __future__ import annotations
+
+from airflow.providers.common.ai.operators.llm_sql import LLMSQLQueryOperator
+from airflow.providers.common.compat.sdk import dag, task
+
+
+# [START howto_operator_llm_sql_basic]
+@dag(schedule=None)
+def example_llm_sql_basic():
+    LLMSQLQueryOperator(
+        task_id="generate_sql",
+        prompt="Find the top 10 customers by total revenue",
+        llm_conn_id="pydantic_ai_default",
+        schema_context=(
+            "Table: customers\n"
+            "Columns: id INT, name TEXT, email TEXT\n\n"
+            "Table: orders\n"
+            "Columns: id INT, customer_id INT, total DECIMAL, created_at TIMESTAMP"
+        ),
+    )
+
+
+# [END howto_operator_llm_sql_basic]
+
+example_llm_sql_basic()
+
+
+# [START howto_operator_llm_sql_schema]
+@dag(schedule=None)
+def example_llm_sql_schema_introspection():
+    LLMSQLQueryOperator(
+        task_id="generate_sql",
+        prompt="Calculate monthly revenue for 2024",
+        llm_conn_id="pydantic_ai_default",
+        db_conn_id="postgres_default",
+        table_names=["orders", "customers"],
+        dialect="postgres",
+    )
+
+
+# [END howto_operator_llm_sql_schema]
+
+example_llm_sql_schema_introspection()
+
+
+# [START howto_decorator_llm_sql]
+@dag(schedule=None)
+def example_llm_sql_decorator():
+    @task.llm_sql(
+        llm_conn_id="pydantic_ai_default",
+        schema_context="Table: users\nColumns: id INT, name TEXT, signup_date DATE",
+    )
+    def build_churn_query(ds=None):
+        return f"Find users who signed up before {ds} and have no orders"
+
+    build_churn_query()
+
+
+# [END howto_decorator_llm_sql]
+
+example_llm_sql_decorator()
+
+
+# [START howto_operator_llm_sql_expand]
+@dag(schedule=None)
+def example_llm_sql_expand():
+    LLMSQLQueryOperator.partial(
+        task_id="generate_sql",
+        llm_conn_id="pydantic_ai_default",
+        schema_context=(
+            "Table: orders\nColumns: id INT, customer_id INT, total DECIMAL, created_at TIMESTAMP"
+        ),
+    ).expand(
+        prompt=[
+            "Total revenue by month",
+            "Top 10 customers by order count",
+            "Average order value by day of week",
+        ]
+    )
+
+
+# [END howto_operator_llm_sql_expand]
+
+example_llm_sql_expand()