Add deferred pagination mode to GenericTransfer

airflow/operators/__init__.py

@@ -43,7 +43,7 @@
         "BranchDateTimeOperator": "airflow.providers.standard.operators.datetime.BranchDateTimeOperator",
     },
     "generic_transfer": {
-        "GenericTransfer": "airflow.providers.standard.operators.generic_transfer.GenericTransfer",
+        "GenericTransfer": "airflow.providers.common.sql.operators.generic_transfer.GenericTransfer",
     },
     "weekday": {
         "BranchDayOfWeekOperator": "airflow.providers.standard.operators.weekday.BranchDayOfWeekOperator",

generated/provider_dependencies.json

@@ -1281,7 +1281,6 @@
   },
   "standard": {
     "deps": [
-      "apache-airflow-providers-common-sql>=1.20.0",
       "apache-airflow>=2.9.0"
     ],
     "devel-deps": [],

providers/common/sql/provider.yaml

@@ -76,6 +76,7 @@ operators:
   - integration-name: Common SQL
     python-modules:
       - airflow.providers.common.sql.operators.sql
+      - airflow.providers.common.sql.operators.generic_transfer
 
 dialects:
   - dialect-type: default
@@ -87,6 +88,11 @@ hooks:
       - airflow.providers.common.sql.hooks.handlers
       - airflow.providers.common.sql.hooks.sql
 
+triggers:
+  - integration-name: Common SQL
+    python-modules:
+      - airflow.providers.common.sql.triggers.sql
+
 sensors:
   - integration-name: Common SQL
     python-modules:

providers/common/sql/src/airflow/providers/common/sql/get_provider_info.py

@@ -80,7 +80,10 @@ def get_provider_info():
         "operators": [
             {
                 "integration-name": "Common SQL",
-                "python-modules": ["airflow.providers.common.sql.operators.sql"],
+                "python-modules": [
+                    "airflow.providers.common.sql.operators.sql",
+                    "airflow.providers.common.sql.operators.generic_transfer",
+                ],
             }
         ],
         "dialects": [
@@ -98,6 +101,12 @@ def get_provider_info():
                 ],
             }
         ],
+        "triggers": [
+            {
+                "integration-name": "Common SQL",
+                "python-modules": ["airflow.providers.common.sql.triggers.sql"],
+            }
+        ],
         "sensors": [
             {"integration-name": "Common SQL", "python-modules": ["airflow.providers.common.sql.sensors.sql"]}
         ],

providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py

@@ -0,0 +1,219 @@
+#
+# 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.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+from airflow.models import BaseOperator
+from airflow.providers.common.sql.hooks.sql import DbApiHook
+from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger
+
+if TYPE_CHECKING:
+    import jinja2
+
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
+
+
+class GenericTransfer(BaseOperator):
+    """
+    Moves data from a connection to another.
+
+    Assuming that they both provide the required methods in their respective hooks.
+    The source hook needs to expose a `get_records` method, and the destination a
+    `insert_rows` method.
+
+    This is meant to be used on small-ish datasets that fit in memory.
+
+    :param sql: SQL query to execute against the source database. (templated)
+    :param destination_table: target table. (templated)
+    :param source_conn_id: source connection. (templated)
+    :param source_hook_params: source hook parameters.
+    :param destination_conn_id: destination connection. (templated)
+    :param destination_hook_params: destination hook parameters.
+    :param preoperator: sql statement or list of statements to be
+        executed prior to loading the data. (templated)
+    :param insert_args: extra params for `insert_rows` method.
+    :param page_size: number of records to be read in paginated mode (optional).
+    """
+
+    template_fields: Sequence[str] = (
+        "source_conn_id",
+        "destination_conn_id",
+        "sql",
+        "destination_table",
+        "preoperator",
+        "insert_args",
+    )
+    template_ext: Sequence[str] = (
+        ".sql",
+        ".hql",
+    )
+    template_fields_renderers = {"preoperator": "sql"}
+    ui_color = "#b0f07c"
+
+    def __init__(
+        self,
+        *,
+        sql: str,
+        destination_table: str,
+        source_conn_id: str,
+        source_hook_params: dict | None = None,
+        destination_conn_id: str,
+        destination_hook_params: dict | None = None,
+        preoperator: str | list[str] | None = None,
+        insert_args: dict | None = None,
+        page_size: int | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.sql = sql
+        self.destination_table = destination_table
+        self.source_conn_id = source_conn_id
+        self.source_hook_params = source_hook_params
+        self.destination_conn_id = destination_conn_id
+        self.destination_hook_params = destination_hook_params
+        self.preoperator = preoperator
+        self.insert_args = insert_args or {}
+        self.page_size = page_size
+        self._paginated_sql_statement_format = kwargs.get(
+            "paginated_sql_statement_format", "{} LIMIT {} OFFSET {}"
+        )
+
+    @classmethod
+    def get_hook(cls, conn_id: str, hook_params: dict | None = None) -> DbApiHook:
+        """
+        Return DbApiHook for this connection id.
+
+        :param conn_id: connection id
+        :param hook_params: hook parameters
+        :return: DbApiHook for this connection
+        """
+        connection = BaseHook.get_connection(conn_id)
+        hook = connection.get_hook(hook_params=hook_params)
+        if not isinstance(hook, DbApiHook):
+            raise RuntimeError(f"Hook for connection {conn_id!r} must be of type {DbApiHook.__name__}")
+        return hook
+
+    @cached_property
+    def source_hook(self) -> DbApiHook:
+        return self.get_hook(conn_id=self.source_conn_id, hook_params=self.source_hook_params)
+
+    @cached_property
+    def destination_hook(self) -> DbApiHook:
+        return self.get_hook(conn_id=self.destination_conn_id, hook_params=self.destination_hook_params)
+
+    def get_paginated_sql(self, offset: int) -> str:
+        """Format the paginated SQL statement using the current format."""
+        return self._paginated_sql_statement_format.format(self.sql, self.page_size, offset)
+
+    def render_template_fields(
+        self,
+        context: Context,
+        jinja_env: jinja2.Environment | None = None,
+    ) -> None:
+        super().render_template_fields(context=context, jinja_env=jinja_env)
+
+        # Make sure string are converted to integers
+        if isinstance(self.page_size, str):
+            self.page_size = int(self.page_size)
+        commit_every = self.insert_args.get("commit_every")
+        if isinstance(commit_every, str):
+            self.insert_args["commit_every"] = int(commit_every)
+
+    def execute(self, context: Context):
+        if self.preoperator:
+            self.log.info("Running preoperator")
+            self.log.info(self.preoperator)
+            self.destination_hook.run(self.preoperator)
+
+        if self.page_size and isinstance(self.sql, str):
+            self.defer(
+                trigger=SQLExecuteQueryTrigger(
+                    conn_id=self.source_conn_id,
+                    hook_params=self.source_hook_params,
+                    sql=self.get_paginated_sql(0),
+                ),
+                method_name=self.execute_complete.__name__,
+            )
+        else:
+            self.log.info("Extracting data from %s", self.source_conn_id)
+            self.log.info("Executing: \n %s", self.sql)
+
+            results = self.destination_hook.get_records(self.sql)
+
+            self.log.info("Inserting rows into %s", self.destination_conn_id)
+            self.destination_hook.insert_rows(table=self.destination_table, rows=results, **self.insert_args)
+
+    def execute_complete(
+        self,
+        context: Context,
+        event: dict[Any, Any] | None = None,
+    ) -> Any:
+        if event:
+            if event.get("status") == "failure":
+                raise AirflowException(event.get("message"))
+
+            results = event.get("results")
+
+            if results:
+                map_index = context["ti"].map_index
+                offset = (
+                    context["ti"].xcom_pull(
+                        key="offset",
+                        task_ids=self.task_id,
+                        dag_id=self.dag_id,
+                        map_indexes=map_index,
+                        default=0,
+                    )
+                    + self.page_size
+                )
+
+                self.log.info("Offset increased to %d", offset)
+                self.xcom_push(context=context, key="offset", value=offset)
+
+                self.log.info("Inserting %d rows into %s", len(results), self.destination_conn_id)
+                self.destination_hook.insert_rows(
+                    table=self.destination_table, rows=results, **self.insert_args
+                )
+                self.log.info(
+                    "Inserting %d rows into %s done!",
+                    len(results),
+                    self.destination_conn_id,
+                )
+
+                self.defer(
+                    trigger=SQLExecuteQueryTrigger(
+                        conn_id=self.source_conn_id,
+                        hook_params=self.source_hook_params,
+                        sql=self.get_paginated_sql(offset),
+                    ),
+                    method_name=self.execute_complete.__name__,
+                )
+            else:
+                self.log.info(
+                    "No more rows to fetch into %s; ending transfer.",
+                    self.destination_table,
+                )

providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.pyi

@@ -0,0 +1,85 @@
+# 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.
+#
+# This is automatically generated stub for the `common.sql` provider
+#
+# This file is generated automatically by the `update-common-sql-api stubs` pre-commit
+# and the .pyi file represents part of the "public" API that the
+# `common.sql` provider exposes to other providers.
+#
+# Any, potentially breaking change in the stubs will require deliberate manual action from the contributor
+# making a change to the `common.sql` provider. Those stubs are also used by MyPy automatically when checking
+# if only public API of the common.sql provider is used by all the other providers.
+#
+# You can read more in the README_API.md file
+#
+"""
+Definition of the public interface for airflow.providers.common.sql.operators.generic_transfer
+isort:skip_file
+"""
+
+from collections.abc import Sequence
+from functools import cached_property as cached_property
+from typing import Any, ClassVar
+
+import jinja2
+from _typeshed import Incomplete as Incomplete
+
+from airflow.models import BaseOperator
+from airflow.providers.common.sql.hooks.sql import DbApiHook as DbApiHook
+from airflow.utils.context import Context as Context
+
+class GenericTransfer(BaseOperator):
+    template_fields: Sequence[str]
+    template_ext: Sequence[str]
+    template_fields_renderers: ClassVar[dict]
+    ui_color: str
+    sql: Incomplete
+    destination_table: Incomplete
+    source_conn_id: Incomplete
+    source_hook_params: Incomplete
+    destination_conn_id: Incomplete
+    destination_hook_params: Incomplete
+    preoperator: Incomplete
+    insert_args: Incomplete
+    page_size: Incomplete
+    def __init__(
+        self,
+        *,
+        sql: str,
+        destination_table: str,
+        source_conn_id: str,
+        source_hook_params: dict | None = None,
+        destination_conn_id: str,
+        destination_hook_params: dict | None = None,
+        preoperator: str | list[str] | None = None,
+        insert_args: dict | None = None,
+        page_size: int | None = None,
+        **kwargs,
+    ) -> None: ...
+    @classmethod
+    def get_hook(cls, conn_id: str, hook_params: dict | None = None) -> DbApiHook: ...
+    @cached_property
+    def source_hook(self) -> DbApiHook: ...
+    @cached_property
+    def destination_hook(self) -> DbApiHook: ...
+    def get_paginated_sql(self, offset: int) -> str: ...
+    def render_template_fields(
+        self, context: Context, jinja_env: jinja2.Environment | None = None
+    ) -> None: ...
+    def execute(self, context: Context): ...
+    def execute_complete(self, context: Context, event: dict[Any, Any] | None = None) -> Any: ...