shahar1 commented on code in PR #69362:
URL: https://github.com/apache/airflow/pull/69362#discussion_r3650231948
##########
providers/common/sql/tests/unit/common/sql/operators/test_sql.py:
##########
@@ -2957,3 +2958,92 @@ def test_insert_rows_operator_with_rows_processor(self,
mock_get_db_hook):
(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
+
+ hook.run.assert_has_calls(
+ [
+ mock.call("CREATE TABLE users (id INT);"),
+ mock.call("DROP TABLE users;"),
+ ]
+ )
+
+ hook.bulk_load.assert_called_once_with(
+ table="users",
+ tmp_file="/tmp/users.tsv",
+ )
Review Comment:
It might be better to assert on the exact order of the calls: Create table
->>bulk load->Drop table
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -1916,3 +1916,75 @@ 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.
+
+ :param table: Name of the target table.
Review Comment:
nit: `table` is also templated
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -1916,3 +1916,75 @@ 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.
+
+ :param table: Name of the target table.
+ :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``.
+ """
Review Comment:
Maybe worth adding a refernce to the docs:
```rst
.. seealso::
For more information on how to use this operator, take a look at the
guide:
:ref:`howto/operator:SQLBulkLoadOperator`
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
