This is an automated email from the ASF dual-hosted git repository.
shahar1 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 c4a5f19e67c Add SQLBulkLoadOperator to Common SQL Provider (#69362)
c4a5f19e67c is described below
commit c4a5f19e67cd25143cc000d6ae0255e13c142b95
Author: SameerMesiah97 <[email protected]>
AuthorDate: Sun Jul 26 21:33:04 2026 +0100
Add SQLBulkLoadOperator to Common SQL Provider (#69362)
Add SQLBulkLoadOperator to bulk load tab-delimited files into SQL tables
using the underlying hook's native bulk loading implementation. The operator
supports optional pre- and post-load SQL statements for common setup and
cleanup tasks. Includes unit tests, documentation, and an example DAG.
---
providers/common/sql/docs/operators.rst | 29 ++++++
.../airflow/providers/common/sql/operators/sql.py | 76 ++++++++++++++++
.../system/common/sql/example_sql_bulk_load.py | 101 +++++++++++++++++++++
.../tests/unit/common/sql/operators/test_sql.py | 84 +++++++++++++++++
4 files changed, 290 insertions(+)
diff --git a/providers/common/sql/docs/operators.rst
b/providers/common/sql/docs/operators.rst
index 88e326a25fe..265832301c0 100644
--- a/providers/common/sql/docs/operators.rst
+++ b/providers/common/sql/docs/operators.rst
@@ -239,6 +239,35 @@ The example below shows how to instantiate the
SQLInsertRowsOperator task.
:start-after: [START howto_operator_sql_insert_rows]
:end-before: [END howto_operator_sql_insert_rows]
+.. _howto/operator:SQLBulkLoadOperator:
+
+Bulk load data into a table
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use the
:class:`~airflow.providers.common.sql.operators.sql.SQLBulkLoadOperator`
+to bulk load a tab-delimited file into a database table using the database's
+native bulk loading mechanism. Parameters of the operator are:
+
+- ``table`` - name of the target table (templated).
+- ``tmp_file`` - path to the tab-delimited file to load (templated).
+- ``conn_id`` - the Airflow connection ID used to connect to the database.
+- ``database`` (optional) - name of the database which overrides the one
defined
+ in the connection.
+- ``preoperator`` (optional) - SQL statement or list of statements to execute
+ before bulk loading (templated).
+- ``postoperator`` (optional) - SQL statement or list of statements to execute
+ after bulk loading (templated).
+- ``hook_params`` (optional) - dictionary of additional parameters passed to
the
+ underlying hook.
+
+The example below shows how to instantiate the SQLBulkLoadOperator task.
+
+.. exampleinclude:: /../tests/system/common/sql/example_sql_bulk_load.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_sql_bulk_load]
+ :end-before: [END howto_operator_sql_bulk_load]
+
.. _howto/operator:GenericTransfer:
Generic Transfer
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
index 21c8b6502b0..3593c258214 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
@@ -1916,3 +1916,79 @@ def _initialize_partition_clause(clause: str | None) ->
str | None:
raise ValueError("Invalid partition_clause: semicolons (;) not
allowed.")
return clause
+
+
+class SQLBulkLoadOperator(BaseSQLOperator):
+ """
+ Bulk load a tab-delimited file into a database table.
+
+ This operator delegates to the underlying hook's ``bulk_load``
implementation,
+ allowing each provider to use its native bulk loading mechanism.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SQLBulkLoadOperator`
+
+ :param table: Name of the target table (templated).
+ :param tmp_file: Path to the tab-delimited file to load (templated).
+ :param conn_id: The connection ID used to connect to the database.
Defaults to ``None``.
+ :param database: Name of the database which overrides the one defined in
the connection.
+ Defaults to ``None``.
+ :param preoperator: SQL statement or list of statements to execute before
bulk loading
+ (templated). Defaults to ``None``.
+ :param postoperator: SQL statement or list of statements to execute after
bulk loading
+ (templated). Defaults to ``None``.
+ """
+
+ template_fields: Sequence[str] = (
+ "table",
+ "tmp_file",
+ "preoperator",
+ "postoperator",
+ *BaseSQLOperator.template_fields,
+ )
+
+ def __init__(
+ self,
+ *,
+ table: str,
+ tmp_file: str,
+ conn_id: str | None = None,
+ database: str | None = None,
+ preoperator: str | list[str] | None = None,
+ postoperator: str | list[str] | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(
+ conn_id=conn_id,
+ database=database,
+ **kwargs,
+ )
+ self.table = table
+ self.tmp_file = tmp_file
+ self.preoperator = preoperator
+ self.postoperator = postoperator
+
+ def execute(self, context: Context) -> None:
+ hook = self.get_db_hook()
+
+ if self.preoperator:
+ self.log.info("Executing preoperator.")
+ hook.run(self.preoperator)
+
+ self.log.info(
+ "Bulk loading '%s' into table '%s'.",
+ self.tmp_file,
+ self.table,
+ )
+
+ hook.bulk_load(
+ table=self.table,
+ tmp_file=self.tmp_file,
+ )
+
+ if self.postoperator:
+ self.log.info("Executing postoperator.")
+ hook.run(self.postoperator)
+
+ self.log.info("Bulk load completed.")
diff --git
a/providers/common/sql/tests/system/common/sql/example_sql_bulk_load.py
b/providers/common/sql/tests/system/common/sql/example_sql_bulk_load.py
new file mode 100644
index 00000000000..8cea0c93a32
--- /dev/null
+++ b/providers/common/sql/tests/system/common/sql/example_sql_bulk_load.py
@@ -0,0 +1,101 @@
+#
+# 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 pathlib import Path
+
+from airflow import DAG
+from airflow.providers.common.sql.operators.sql import (
+ SQLBulkLoadOperator,
+)
+from airflow.providers.standard.operators.python import PythonOperator
+from airflow.utils.timezone import datetime
+
+TMP_FILE = "/tmp/actors.tsv"
+
+connection_args = {
+ "conn_id": "airflow_db",
+ "conn_type": "Postgres",
+ "host": "postgres",
+ "schema": "postgres",
+ "login": "postgres",
+ "password": "postgres",
+ "port": 5432,
+}
+
+
+def create_bulk_load_file():
+ Path(TMP_FILE).write_text(
+ "\n".join(
+ [
+ "Stallone\tSylvester\t78",
+ "Statham\tJason\t57",
+ "Li\tJet\t61",
+ "Lundgren\tDolph\t66",
+ "Norris\tChuck\t84",
+ ]
+ ),
+ encoding="utf-8",
+ )
+
+
+with DAG(
+ "example_sql_bulk_load",
+ description="Example DAG for SQLBulkLoadOperator.",
+ default_args=connection_args,
+ start_date=datetime(2021, 1, 1),
+ schedule=None,
+ catchup=False,
+) as dag:
+ """
+ ### Example SQL bulk load DAG
+
+ Runs the SQLBulkLoadOperator against the Airflow metadata DB.
+ """
+
+ create_file = PythonOperator(
+ task_id="create_file",
+ python_callable=create_bulk_load_file,
+ )
+
+ # [START howto_operator_sql_bulk_load]
+ bulk_load = SQLBulkLoadOperator(
+ task_id="bulk_load",
+ table="actors",
+ tmp_file=TMP_FILE,
+ preoperator=[
+ """
+ CREATE TABLE IF NOT EXISTS actors (
+ name TEXT NOT NULL,
+ firstname TEXT NOT NULL,
+ age BIGINT NOT NULL
+ );
+ """,
+ "TRUNCATE TABLE actors;",
+ ],
+ postoperator="DROP TABLE IF EXISTS actors;",
+ )
+ # [END howto_operator_sql_bulk_load]
+
+ create_file >> bulk_load
+
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+# Needed to run the example DAG with pytest (see:
contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)
diff --git a/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
b/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
index 14380b71736..ef37630d440 100644
--- a/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
+++ b/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
@@ -34,6 +34,7 @@ from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.providers.common.sql.operators.sql import (
BaseSQLOperator,
BranchSQLOperator,
+ SQLBulkLoadOperator,
SQLCheckOperator,
SQLCheckResult,
SQLColumnCheckOperator,
@@ -2957,3 +2958,86 @@ class TestSQLInsertRowsOperator:
(4, "Lundgren", "Dolph", 66),
(5, "Norris", "Chuck", 84),
]
+
+
+class TestSQLBulkLoadOperator:
+ def _construct_operator(self, table, tmp_file, **kwargs):
+ dag = DAG("test_dag", schedule=None,
start_date=datetime.datetime(2017, 1, 1))
+ return SQLBulkLoadOperator(
+ task_id="test_task",
+ conn_id="default_conn",
+ table=table,
+ tmp_file=tmp_file,
+ dag=dag,
+ **kwargs,
+ )
+
+ @mock.patch.object(SQLBulkLoadOperator, "get_db_hook")
+ def test_execute(self, mock_get_db_hook):
+ operator = self._construct_operator(
+ table="users",
+ tmp_file="/tmp/users.tsv",
+ preoperator="CREATE TABLE users (id INT);",
+ postoperator="DROP TABLE users;",
+ )
+
+ operator.execute(context=MagicMock())
+
+ hook = mock_get_db_hook.return_value
+
+ assert hook.mock_calls == [
+ mock.call.run("CREATE TABLE users (id INT);"),
+ mock.call.bulk_load(
+ table="users",
+ tmp_file="/tmp/users.tsv",
+ ),
+ mock.call.run("DROP TABLE users;"),
+ ]
+
+ @mock.patch.object(SQLBulkLoadOperator, "get_db_hook")
+ def test_execute_templated_fields(self, mock_get_db_hook):
+ operator = self._construct_operator(
+ table="{{ params.table }}",
+ tmp_file="{{ params.file }}",
+ preoperator="TRUNCATE TABLE {{ params.table }};",
+ postoperator="DROP TABLE {{ params.table }};",
+ )
+
+ operator.render_template_fields(
+ {
+ "params": {
+ "table": "users",
+ "file": "/tmp/users.tsv",
+ }
+ }
+ )
+
+ operator.execute(context=MagicMock())
+
+ hook = mock_get_db_hook.return_value
+
+ assert hook.mock_calls == [
+ mock.call.run("TRUNCATE TABLE users;"),
+ mock.call.bulk_load(
+ table="users",
+ tmp_file="/tmp/users.tsv",
+ ),
+ mock.call.run("DROP TABLE users;"),
+ ]
+
+ @mock.patch.object(SQLBulkLoadOperator, "get_db_hook")
+ def test_execute_without_pre_or_post_operator(self, mock_get_db_hook):
+ operator = self._construct_operator(
+ table="users",
+ tmp_file="/tmp/users.tsv",
+ )
+
+ operator.execute(context=MagicMock())
+
+ hook = mock_get_db_hook.return_value
+
+ hook.run.assert_not_called()
+ hook.bulk_load.assert_called_once_with(
+ table="users",
+ tmp_file="/tmp/users.tsv",
+ )
