diff --git a/airflow/providers/common/sql/CHANGELOG.rst b/airflow/providers/common/sql/CHANGELOG.rst index 7f48cf557310a..b436b9bc4355f 100644 --- a/airflow/providers/common/sql/CHANGELOG.rst +++ b/airflow/providers/common/sql/CHANGELOG.rst @@ -32,7 +32,9 @@ This release fixes a few errors that were introduced in common.sql operator whil * ``_process_output`` method in ``SQLExecuteQueryOperator`` has now consistent semantics and typing, it can also modify the returned (and stored in XCom) values in the operators that derive from the ``SQLExecuteQueryOperator``). -* last description of the cursor whether to return scalar values are now stored in DBApiHook +* descriptions of all returned results are stored as descriptions property in the DBApiHook +* last description of the cursor whether to return single query results values are now exposed in + DBApiHook via last_description property. Lack of consistency in the operator caused ``1.3.0`` to be yanked - the ``1.3.0`` should not be used - if you have ``1.3.0`` installed, upgrade to ``1.3.1``. diff --git a/airflow/providers/common/sql/hooks/sql.py b/airflow/providers/common/sql/hooks/sql.py index df808430fd9aa..06dbdb2f57641 100644 --- a/airflow/providers/common/sql/hooks/sql.py +++ b/airflow/providers/common/sql/hooks/sql.py @@ -30,6 +30,33 @@ from airflow.version import version +def return_single_query_results(sql: str | Iterable[str], return_last: bool, split_statements: bool): + """ + Determines when results of single query only should be returned. + + For compatibility reasons, the behaviour of the DBAPIHook is somewhat confusing. + In cases, when multiple queries are run, the return values will be an iterable (list) of results + - one for each query. However, in certain cases, when single query is run - the results will be just + the results of that single query without wrapping the results in a list. + + The cases when single query results are returned without wrapping them in a list are when: + + a) sql is string and last_statement is True (regardless what split_statement value is) + b) sql is string and split_statement is False + + In all other cases, the results are wrapped in a list, even if there is only one statement to process: + + a) always when sql is an iterable of string statements (regardless what last_statement value is) + b) when sql is string, split_statement is True and last_statement is False + + :param sql: sql to run (either string or list of strings) + :param return_last: whether last statement output should only be returned + :param split_statements: whether to split string statements. + :return: True if the hook should return single query results + """ + return isinstance(sql, str) and (return_last or not split_statements) + + def fetch_all_handler(cursor) -> list[tuple] | None: """Handler for DbApiHook.run() to return results""" if cursor.description is not None: @@ -114,8 +141,7 @@ def __init__(self, *args, schema: str | None = None, log_sql: bool = True, **kwa # Hook deriving from the DBApiHook to still have access to the field in its constructor self.__schema = schema self.log_sql = log_sql - self.scalar_return_last = False - self.last_description: Sequence[Sequence] | None = None + self.descriptions: list[Sequence[Sequence] | None] = [] def get_conn(self): """Returns a connection object""" @@ -222,6 +248,12 @@ def split_sql_string(sql: str) -> list[str]: statements: list[str] = list(filter(None, splits)) return statements + @property + def last_description(self) -> Sequence[Sequence] | None: + if not self.descriptions: + return None + return self.descriptions[-1] + def run( self, sql: str | Iterable[str], @@ -234,7 +266,42 @@ def run( """ Runs a command or a list of commands. Pass a list of sql statements to the sql parameter to get them to execute - sequentially + sequentially. + + The method will return either single query results (typically list of rows) or list of those results + where each element in the list are results of one of the queries (typically list of list of rows :D) + + For compatibility reasons, the behaviour of the DBAPIHook is somewhat confusing. + In cases, when multiple queries are run, the return values will be an iterable (list) of results + - one for each query. However, in certain cases, when single query is run - the results will be just + the results of that query without wrapping the results in a list. + + The cases when single query results are returned without wrapping them in a list are when: + + a) sql is string and last_statement is True (regardless what split_statement value is) + b) sql is string and split_statement is False + + In all other cases, the results are wrapped in a list, even if there is only one statement to process: + + a) always when sql is an iterable of string statements (regardless what last_statement value is) + b) when sql is string, split_statement is True and last_statement is False + + + In any of those cases, however you can access the following properties of the Hook after running it: + + * descriptions - has an array of cursor descriptions - each statement executed contain the list + of descriptions executed. If ``return_last`` is used, this is always a one-element array + * last_description - description of the last statement executed + + Note that return value from the hook will ONLY be actually returned when handler is provided. Setting + the ``handler`` to None, results in this method returning None. + + Handler is a way to process the rows from cursor (Iterator) into a value that is suitable to be + returned to XCom and generally fit in memory. As an optimization, handler is usually not executed + by the SQLExecuteQuery operator if `do_xcom_push` is not specified. + + You can use pre-defined handles (`fetch_all_handler``, ''fetch_one_handler``) or implement your + own handler. :param sql: the sql statement to be executed (str) or a list of sql statements to execute @@ -246,31 +313,38 @@ def run( :param return_last: Whether to return result for only last statement or for all after split :return: return only result of the ALL SQL expressions if handler was provided. """ - self.scalar_return_last = isinstance(sql, str) and return_last + self.descriptions = [] + if isinstance(sql, str): if split_statements: - sql = self.split_sql_string(sql) + sql_list: Iterable[str] = self.split_sql_string(sql) else: - sql = [sql] + sql_list = [sql] if sql.strip() else [] + else: + sql_list = sql - if sql: - self.log.debug("Executing following statements against DB: %s", list(sql)) + if sql_list: + self.log.debug("Executing following statements against DB: %s", sql_list) else: raise ValueError("List of SQL statements is empty") - + _last_result = None with closing(self.get_conn()) as conn: if self.supports_autocommit: self.set_autocommit(conn, autocommit) with closing(conn.cursor()) as cur: results = [] - for sql_statement in sql: + for sql_statement in sql_list: self._run_command(cur, sql_statement, parameters) if handler is not None: result = handler(cur) - results.append(result) - self.last_description = cur.description + if return_single_query_results(sql, return_last, split_statements): + _last_result = result + _last_description = cur.description + else: + results.append(result) + self.descriptions.append(cur.description) # If autocommit was set to False or db does not support autocommit, we do a manual commit. if not self.get_autocommit(conn): @@ -278,8 +352,9 @@ def run( if handler is None: return None - elif self.scalar_return_last: - return results[-1] + if return_single_query_results(sql, return_last, split_statements): + self.descriptions = [_last_description] + return _last_result else: return results diff --git a/airflow/providers/common/sql/operators/sql.py b/airflow/providers/common/sql/operators/sql.py index 314af43003488..09034b104ca85 100644 --- a/airflow/providers/common/sql/operators/sql.py +++ b/airflow/providers/common/sql/operators/sql.py @@ -19,14 +19,13 @@ import ast import re -from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, NoReturn, Sequence, SupportsAbs, overload +from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, NoReturn, Sequence, SupportsAbs from airflow.compat.functools import cached_property from airflow.exceptions import AirflowException, AirflowFailException from airflow.hooks.base import BaseHook from airflow.models import BaseOperator, SkipMixin -from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler -from airflow.typing_compat import Literal +from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler, return_single_query_results if TYPE_CHECKING: from airflow.utils.context import Context @@ -190,11 +189,17 @@ class SQLExecuteQueryOperator(BaseSQLOperator): Executes SQL code in a specific database :param sql: the SQL code or string pointing to a template file to be executed (templated). File must have a '.sql' extensions. + + When implementing a specific Operator, you can also implement `_process_output` method in the + hook to perform additional processing of values returned by the DB Hook of yours. For example, you + can join description retrieved from the cursors of your statements with returned values, or save + the output of your operator to a file. + :param autocommit: (optional) if True, each command is automatically committed (default: False). :param parameters: (optional) the parameters to render the SQL query with. :param handler: (optional) the function that will be applied to the cursor (default: fetch_all_handler). :param split_statements: (optional) if split single SQL string into statements (default: False). - :param return_last: (optional) if return the result of only last statement (default: True). + :param return_last: (optional) return the result of only last statement (default: True). .. seealso:: For more information on how to use this operator, take a look at the guide: @@ -225,54 +230,42 @@ def __init__( self.split_statements = split_statements self.return_last = return_last - @overload - def _process_output( - self, results: Any, description: Sequence[Sequence] | None, scalar_results: Literal[True] - ) -> Any: - pass - - @overload - def _process_output( - self, results: list[Any], description: Sequence[Sequence] | None, scalar_results: Literal[False] - ) -> Any: - pass - - def _process_output( - self, results: Any | list[Any], description: Sequence[Sequence] | None, scalar_results: bool - ) -> Any: + def _process_output(self, results: list[Any], descriptions: list[Sequence[Sequence] | None]) -> list[Any]: """ - Can be overridden by the subclass in case some extra processing is needed. + Processes output before it is returned by the operator. + + It can be overridden by the subclass in case some extra processing is needed. Note that unlike + DBApiHook return values returned - the results passed and returned by ``_process_output`` should + always be lists of results - each element of the list is a result from a single SQL statement + (typically this will be list of Rows). You have to make sure that this is the same for returned + values = there should be one element in the list for each statement executed by the hook.. + The "process_output" method can override the returned output - augmenting or processing the output as needed - the output returned will be returned as execute return value and if - do_xcom_push is set to True, it will be set as XCom returned + do_xcom_push is set to True, it will be set as XCom returned. :param results: results in the form of list of rows. - :param description: as returned by ``cur.description`` in the Python DBAPI - :param scalar_results: True if result is single scalar value rather than list of rows + :param descriptions: list of descriptions returned by ``cur.description`` in the Python DBAPI """ return results def execute(self, context): self.log.info("Executing: %s", self.sql) hook = self.get_db_hook() - if self.do_xcom_push: - output = hook.run( - sql=self.sql, - autocommit=self.autocommit, - parameters=self.parameters, - handler=self.handler, - split_statements=self.split_statements, - return_last=self.return_last, - ) - else: - output = hook.run( - sql=self.sql, - autocommit=self.autocommit, - parameters=self.parameters, - split_statements=self.split_statements, - ) - - return self._process_output(output, hook.last_description, hook.scalar_return_last) + output = hook.run( + sql=self.sql, + autocommit=self.autocommit, + parameters=self.parameters, + handler=self.handler if self.do_xcom_push else None, + split_statements=self.split_statements, + return_last=self.return_last, + ) + if return_single_query_results(self.sql, self.return_last, self.split_statements): + # For simplicity, we pass always list as input to _process_output, regardless if + # single query results are going to be returned, and we return the first element + # of the list in this case from the (always) list returned by _process_output + return self._process_output([output], hook.descriptions)[-1] + return self._process_output(output, hook.descriptions) def prepare_template(self) -> None: """Parse template file for attribute parameters.""" @@ -284,6 +277,7 @@ class SQLColumnCheckOperator(BaseSQLOperator): """ Performs one or more of the templated checks in the column_checks dictionary. Checks are performed on a per-column basis specified by the column_mapping. + Each check can take one or more of the following options: - equal_to: an exact value to equal, cannot be used with other comparison options - greater_than: value that result should be strictly greater than diff --git a/airflow/providers/databricks/hooks/databricks_sql.py b/airflow/providers/databricks/hooks/databricks_sql.py index f042435943889..245128aca97b2 100644 --- a/airflow/providers/databricks/hooks/databricks_sql.py +++ b/airflow/providers/databricks/hooks/databricks_sql.py @@ -24,7 +24,7 @@ from databricks.sql.client import Connection # type: ignore[attr-defined] from airflow.exceptions import AirflowException -from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints") @@ -151,49 +151,57 @@ def run( """ Runs a command or a list of commands. Pass a list of sql statements to the sql parameter to get them to execute - sequentially + sequentially. :param sql: the sql statement to be executed (str) or a list of sql statements to execute :param autocommit: What to set the connection's autocommit setting to - before executing the query. + before executing the query. Note that currently there is no commit functionality + in Databricks SQL so this flag has no effect. + :param parameters: The parameters to render the SQL query with. :param handler: The result handler which is called with the result of each statement. :param split_statements: Whether to split a single SQL string into statements and run separately :param return_last: Whether to return result for only last statement or for all after split - :return: return only result of the LAST SQL expression if handler was provided. + :return: return only result of the LAST SQL expression if handler was provided unless return_last + is set to False. """ - self.scalar_return_last = isinstance(sql, str) and return_last + self.descriptions = [] if isinstance(sql, str): if split_statements: - sql = self.split_sql_string(sql) + sql_list = [self.strip_sql_string(s) for s in self.split_sql_string(sql)] else: - sql = [self.strip_sql_string(sql)] + sql_list = [self.strip_sql_string(sql)] + else: + sql_list = [self.strip_sql_string(s) for s in sql] - if sql: - self.log.debug("Executing following statements against Databricks DB: %s", list(sql)) + if sql_list: + self.log.debug("Executing following statements against Databricks DB: %s", sql_list) else: raise ValueError("List of SQL statements is empty") results = [] - for sql_statement in sql: + for sql_statement in sql_list: # when using AAD tokens, it could expire if previous query run longer than token lifetime with closing(self.get_conn()) as conn: self.set_autocommit(conn, autocommit) with closing(conn.cursor()) as cur: self._run_command(cur, sql_statement, parameters) - if handler is not None: result = handler(cur) - results.append(result) - self.last_description = cur.description + if return_single_query_results(sql, return_last, split_statements): + results = [result] + self.descriptions = [cur.description] + else: + results.append(result) + self.descriptions.append(cur.description) self._sql_conn = None if handler is None: return None - elif self.scalar_return_last: + if return_single_query_results(sql, return_last, split_statements): return results[-1] else: return results diff --git a/airflow/providers/databricks/operators/databricks_sql.py b/airflow/providers/databricks/operators/databricks_sql.py index 379b0fd2c9bbb..178afc8d9854a 100644 --- a/airflow/providers/databricks/operators/databricks_sql.py +++ b/airflow/providers/databricks/operators/databricks_sql.py @@ -120,21 +120,17 @@ def get_db_hook(self) -> DatabricksSqlHook: } return DatabricksSqlHook(self.databricks_conn_id, **hook_params) - def _process_output( - self, results: Any | list[Any], description: Sequence[Sequence] | None, scalar_results: bool - ) -> Any: + def _process_output(self, results: list[Any], descriptions: list[Sequence[Sequence] | None]) -> list[Any]: if not self._output_path: - return description, results + return list(zip(descriptions, results)) if not self._output_format: raise AirflowException("Output format should be specified!") - if description is None: - self.log.warning("Description of the cursor is missing. Will not process the output") - return description, results - field_names = [field[0] for field in description] - if scalar_results: - list_results: list[Any] = [results] - else: - list_results = results + # Output to a file only the result of last query + last_description = descriptions[-1] + last_results = results[-1] + if last_description is None: + raise AirflowException("There is missing description present for the output file. .") + field_names = [field[0] for field in last_description] if self._output_format.lower() == "csv": with open(self._output_path, "w", newline="") as file: if self._csv_params: @@ -147,19 +143,19 @@ def _process_output( writer = csv.DictWriter(file, fieldnames=field_names, **csv_params) if write_header: writer.writeheader() - for row in list_results: + for row in last_results: writer.writerow(row.asDict()) elif self._output_format.lower() == "json": with open(self._output_path, "w") as file: - file.write(json.dumps([row.asDict() for row in list_results])) + file.write(json.dumps([row.asDict() for row in last_results])) elif self._output_format.lower() == "jsonl": with open(self._output_path, "w") as file: - for row in list_results: + for row in last_results: file.write(json.dumps(row.asDict())) file.write("\n") else: raise AirflowException(f"Unsupported output format: '{self._output_format}'") - return description, results + return list(zip(descriptions, results)) COPY_INTO_APPROVED_FORMATS = ["CSV", "JSON", "AVRO", "ORC", "PARQUET", "TEXT", "BINARYFILE"] diff --git a/airflow/providers/exasol/hooks/exasol.py b/airflow/providers/exasol/hooks/exasol.py index 49289df37a314..3dc6b81973c78 100644 --- a/airflow/providers/exasol/hooks/exasol.py +++ b/airflow/providers/exasol/hooks/exasol.py @@ -24,7 +24,7 @@ import pyexasol from pyexasol import ExaConnection -from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results class ExasolHook(DbApiHook): @@ -157,7 +157,6 @@ def run( :param return_last: Whether to return result for only last statement or for all after split :return: return only result of the LAST SQL expression if handler was provided. """ - self.scalar_return_last = isinstance(sql, str) and return_last if isinstance(sql, str): if split_statements: sql = self.split_sql_string(sql) @@ -187,7 +186,7 @@ def run( if handler is None: return None - elif self.scalar_return_last: + elif return_single_query_results(sql, return_last, split_statements): return results[-1] else: return results diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index e525efe763439..c08247be21716 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -32,7 +32,7 @@ from sqlalchemy import create_engine from airflow import AirflowException -from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results from airflow.utils.strings import to_boolean @@ -350,7 +350,6 @@ def run( """ self.query_ids = [] - self.scalar_return_last = isinstance(sql, str) and return_last if isinstance(sql, str): if split_statements: split_statements_tuple = util_text.split_statements(StringIO(sql)) @@ -387,7 +386,7 @@ def run( if handler is None: return None - elif self.scalar_return_last: + elif return_single_query_results(sql, return_last, split_statements): return results[-1] else: return results diff --git a/tests/providers/common/sql/hooks/test_sql.py b/tests/providers/common/sql/hooks/test_sql.py new file mode 100644 index 0000000000000..72684e5e00b8e --- /dev/null +++ b/tests/providers/common/sql/hooks/test_sql.py @@ -0,0 +1,226 @@ +# 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 unittest.mock import MagicMock + +import pytest + +from airflow.models import Connection +from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler +from airflow.utils.session import provide_session + +TASK_ID = "sql-operator" +HOST = "host" +DEFAULT_CONN_ID = "sqlite_default" +PASSWORD = "password" + + +class DBApiHookForTests(DbApiHook): + conn_name_attr = "conn_id" + get_conn = MagicMock(name="conn") + + +@provide_session +@pytest.fixture(autouse=True) +def create_connection(session): + conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() + conn.host = HOST + conn.login = None + conn.password = PASSWORD + conn.extra = None + session.commit() + + +def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: + return [(field,) for field in fields] + + +index = 0 + + +@pytest.mark.parametrize( + "return_last, split_statements, sql, cursor_calls," + "cursor_descriptions, cursor_results, hook_descriptions, hook_results, ", + [ + pytest.param( + True, + False, + "select * from test.test", + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last set and no split statements set on single query in string", + ), + pytest.param( + False, + False, + "select * from test.test;", + ["select * from test.test;"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last not set and no split statements set on single query in string", + ), + pytest.param( + True, + True, + "select * from test.test;", + ["select * from test.test;"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last set and split statements set on single query in string", + ), + pytest.param( + False, + True, + "select * from test.test;", + ["select * from test.test;"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last not set and split statements set on single query in string", + ), + pytest.param( + True, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test;", "select * from test.test2;"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id2",), ("value2",)]], + [[3, 4], [13, 14]], + id="The return_last set and split statements set on multiple queries in string", + ), # Failing + pytest.param( + False, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test;", "select * from test.test2;"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id",), ("value",)], [("id2",), ("value2",)]], + [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + id="The return_last not set and split statements set on multiple queries in string", + ), + pytest.param( + True, + True, + ["select * from test.test;"], + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last set on single query in list", + ), + pytest.param( + False, + True, + ["select * from test.test;"], + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last not set on single query in list", + ), + pytest.param( + True, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test", "select * from test.test2"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id2",), ("value2",)]], + [[3, 4], [13, 14]], + id="The return_last set set on multiple queries in list", + ), + pytest.param( + False, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test", "select * from test.test2"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id",), ("value",)], [("id2",), ("value2",)]], + [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + id="The return_last not set on multiple queries not set", + ), + ], +) +def test_query( + return_last, + split_statements, + sql, + cursor_calls, + cursor_descriptions, + cursor_results, + hook_descriptions, + hook_results, +): + modified_descriptions = [ + get_cursor_descriptions(cursor_description) for cursor_description in cursor_descriptions + ] + dbapi_hook = DBApiHookForTests() + dbapi_hook.get_conn.return_value.cursor.return_value.rowcount = 2 + dbapi_hook.get_conn.return_value.cursor.return_value._description_index = 0 + + def mock_execute(*args, **kwargs): + # the run method accesses description property directly, and we need to modify it after + # every execute, to make sure that different descriptions are returned. I could not find easier + # method with mocking + dbapi_hook.get_conn.return_value.cursor.return_value.description = modified_descriptions[ + dbapi_hook.get_conn.return_value.cursor.return_value._description_index + ] + dbapi_hook.get_conn.return_value.cursor.return_value._description_index += 1 + + dbapi_hook.get_conn.return_value.cursor.return_value.execute = mock_execute + dbapi_hook.get_conn.return_value.cursor.return_value.fetchall.side_effect = cursor_results + results = dbapi_hook.run( + sql=sql, handler=fetch_all_handler, return_last=return_last, split_statements=split_statements + ) + + assert dbapi_hook.descriptions == hook_descriptions + assert dbapi_hook.last_description == hook_descriptions[-1] + assert results == hook_results + + dbapi_hook.get_conn.return_value.cursor.return_value.close.assert_called() + + +@pytest.mark.parametrize( + "empty_statement", + [ + pytest.param([], id="Empty list"), + pytest.param("", id="Empty string"), + pytest.param("\n", id="Only EOL"), + ], +) +def test_no_query(empty_statement): + dbapi_hook = DBApiHookForTests() + dbapi_hook.get_conn.return_value.cursor.rowcount = 0 + with pytest.raises(ValueError) as err: + dbapi_hook.run(sql=empty_statement) + assert err.value.args[0] == "List of SQL statements is empty" diff --git a/tests/providers/common/sql/operators/test_sql.py b/tests/providers/common/sql/operators/test_sql.py index 1770ed8f5e659..216ac79280ceb 100644 --- a/tests/providers/common/sql/operators/test_sql.py +++ b/tests/providers/common/sql/operators/test_sql.py @@ -88,6 +88,8 @@ def test_dont_xcom_push(self, mock_get_db_hook): autocommit=False, parameters=None, split_statements=False, + handler=None, + return_last=True, ) diff --git a/tests/providers/common/sql/operators/test_sql_execute.py b/tests/providers/common/sql/operators/test_sql_execute.py new file mode 100644 index 0000000000000..0459472d3abd8 --- /dev/null +++ b/tests/providers/common/sql/operators/test_sql_execute.py @@ -0,0 +1,276 @@ +# +# 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 typing import Any, NamedTuple, Sequence +from unittest.mock import MagicMock + +import pytest + +from airflow.providers.common.sql.hooks.sql import fetch_all_handler +from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator + +DATE = "2017-04-20" +TASK_ID = "sql-operator" + + +class Row(NamedTuple): + id: str + value: str + + +class Row2(NamedTuple): + id2: str + value2: str + + +@pytest.mark.parametrize( + "sql, return_last, split_statement, hook_results, hook_descriptions, expected_results", + [ + pytest.param( + "select * from dummy", + True, + True, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: Single SQL statement, return_last, split statement", + ), + pytest.param( + "select * from dummy;select * from dummy2", + True, + True, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: Multiple SQL statements, return_last, split statement", + ), + pytest.param( + "select * from dummy", + False, + False, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: Single SQL statements, no return_last (doesn't matter), no split statement", + ), + pytest.param( + "select * from dummy", + True, + False, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: Single SQL statements, return_last (doesn't matter), no split statement", + ), + pytest.param( + ["select * from dummy"], + False, + False, + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + [[("id",), ("value",)]], + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + id="Non-Scalar: Single SQL statements in list, no return_last, no split statement", + ), + pytest.param( + ["select * from dummy", "select * from dummy2"], + False, + False, + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")], + ], + [[("id",), ("value",)], [("id2",), ("value2",)]], + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")], + ], + id="Non-Scalar: Multiple SQL statements in list, no return_last (no matter), no split statement", + ), + pytest.param( + ["select * from dummy", "select * from dummy2"], + True, + False, + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")], + ], + [[("id",), ("value",)], [("id2",), ("value2",)]], + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")], + ], + id="Non-Scalar: Multiple SQL statements in list, return_last (no matter), no split statement", + ), + ], +) +def test_exec_success(sql, return_last, split_statement, hook_results, hook_descriptions, expected_results): + """ + Test the execute function in case where SQL query was successful. + """ + + class SQLExecuteQueryOperatorForTest(SQLExecuteQueryOperator): + _mock_db_api_hook = MagicMock() + + def get_db_hook(self): + return self._mock_db_api_hook + + op = SQLExecuteQueryOperatorForTest( + task_id=TASK_ID, + sql=sql, + do_xcom_push=True, + return_last=return_last, + split_statements=split_statement, + ) + + op._mock_db_api_hook.run.return_value = hook_results + op._mock_db_api_hook.descriptions = hook_descriptions + + execute_results = op.execute(None) + + assert execute_results == expected_results + op._mock_db_api_hook.run.assert_called_once_with( + sql=sql, + parameters=None, + handler=fetch_all_handler, + autocommit=False, + return_last=return_last, + split_statements=split_statement, + ) + + +@pytest.mark.parametrize( + "sql, return_last, split_statement, hook_results, hook_descriptions, expected_results", + [ + pytest.param( + "select * from dummy", + True, + True, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Single SQL statement, return_last, split statement", + ), + pytest.param( + "select * from dummy;select * from dummy2", + True, + True, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Multiple SQL statements, return_last, split statement", + ), + pytest.param( + "select * from dummy", + False, + False, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Single SQL statements, no return_last (doesn't matter), no split statement", + ), + pytest.param( + "select * from dummy", + True, + False, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Single SQL statements, return_last (doesn't matter), no split statement", + ), + pytest.param( + ["select * from dummy"], + False, + False, + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + [[("id",), ("value",)]], + [([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")])], + id="Non-Scalar: Single SQL statements in list, no return_last, no split statement", + ), + pytest.param( + ["select * from dummy", "select * from dummy2"], + False, + False, + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")], + ], + [[("id",), ("value",)], [("id2",), ("value2",)]], + [ + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + ([("id2",), ("value2",)], [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")]), + ], + id="Non-Scalar: Multiple SQL statements in list, no return_last (no matter), no split statement", + ), + pytest.param( + ["select * from dummy", "select * from dummy2"], + True, + False, + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")], + ], + [[("id",), ("value",)], [("id2",), ("value2",)]], + [ + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + ([("id2",), ("value2",)], [Row2(id2=1, value2="value1"), Row2(id2=2, value2="value2")]), + ], + id="Non-Scalar: Multiple SQL statements in list, return_last (no matter), no split statement", + ), + ], +) +def test_exec_success_with_process_output( + sql, return_last, split_statement, hook_results, hook_descriptions, expected_results +): + """ + Test the execute function in case where SQL query was successful. + """ + + class SQLExecuteQueryOperatorForTestWithProcessOutput(SQLExecuteQueryOperator): + _mock_db_api_hook = MagicMock() + + def get_db_hook(self): + return self._mock_db_api_hook + + def _process_output( + self, results: list[Any], descriptions: list[Sequence[Sequence] | None] + ) -> list[Any]: + return list(zip(descriptions, results)) + + op = SQLExecuteQueryOperatorForTestWithProcessOutput( + task_id=TASK_ID, + sql=sql, + do_xcom_push=True, + return_last=return_last, + split_statements=split_statement, + ) + + op._mock_db_api_hook.run.return_value = hook_results + op._mock_db_api_hook.descriptions = hook_descriptions + + execute_results = op.execute(None) + + assert execute_results == expected_results + op._mock_db_api_hook.run.assert_called_once_with( + sql=sql, + parameters=None, + handler=fetch_all_handler, + autocommit=False, + return_last=return_last, + split_statements=split_statement, + ) diff --git a/tests/providers/databricks/hooks/test_databricks_sql.py b/tests/providers/databricks/hooks/test_databricks_sql.py index 673ac0bdddb4f..ecc0385278557 100644 --- a/tests/providers/databricks/hooks/test_databricks_sql.py +++ b/tests/providers/databricks/hooks/test_databricks_sql.py @@ -19,6 +19,7 @@ from __future__ import annotations from unittest import mock +from unittest.mock import patch import pytest @@ -34,25 +35,156 @@ TOKEN = "token" -class TestDatabricksSqlHookQueryByName: - """ - Tests for DatabricksHook. - """ - - @provide_session - def setup_method(self, method, session=None): - conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() - conn.host = HOST - conn.login = None - conn.password = TOKEN - conn.extra = None - session.commit() - - self.hook = DatabricksSqlHook(sql_endpoint_name="Test") - - @mock.patch("airflow.providers.databricks.hooks.databricks_sql.DatabricksSqlHook.get_conn") - @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") - def test_query(self, mock_requests, mock_conn): +@provide_session +@pytest.fixture(autouse=True) +def create_connection(session): + conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() + conn.host = HOST + conn.login = None + conn.password = TOKEN + conn.extra = None + session.commit() + + +@pytest.fixture +def databricks_hook(): + return DatabricksSqlHook(sql_endpoint_name="Test") + + +def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: + return [(field,) for field in fields] + + +@pytest.mark.parametrize( + "return_last, split_statements, sql, cursor_calls," + "cursor_descriptions, cursor_results, hook_descriptions, hook_results, ", + [ + pytest.param( + True, + False, + "select * from test.test", + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last set and no split statements set on single query in string", + ), + pytest.param( + False, + False, + "select * from test.test;", + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last not set and no split statements set on single query in string", + ), + pytest.param( + True, + True, + "select * from test.test;", + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last set and split statements set on single query in string", + ), + pytest.param( + False, + True, + "select * from test.test;", + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last not set and split statements set on single query in string", + ), + pytest.param( + True, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test", "select * from test.test2"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id2",), ("value2",)]], + [[3, 4], [13, 14]], + id="The return_last set and split statements set on multiple queries in string", + ), + pytest.param( + False, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test", "select * from test.test2"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id",), ("value",)], [("id2",), ("value2",)]], + [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + id="The return_last not set and split statements set on multiple queries in string", + ), + pytest.param( + True, + True, + ["select * from test.test;"], + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last set on single query in list", + ), + pytest.param( + False, + True, + ["select * from test.test;"], + ["select * from test.test"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last not set on single query in list", + ), + pytest.param( + True, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test", "select * from test.test2"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id2",), ("value2",)]], + [[3, 4], [13, 14]], + id="The return_last set set on multiple queries in list", + ), + pytest.param( + False, + True, + "select * from test.test;select * from test.test2;", + ["select * from test.test", "select * from test.test2"], + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id",), ("value",)], [("id2",), ("value2",)]], + [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + id="The return_last not set on multiple queries not set", + ), + ], +) +def test_query( + databricks_hook, + return_last, + split_statements, + sql, + cursor_calls, + cursor_descriptions, + cursor_results, + hook_descriptions, + hook_results, +): + with patch( + "airflow.providers.databricks.hooks.databricks_sql.DatabricksSqlHook.get_conn" + ) as mock_conn, patch("airflow.providers.databricks.hooks.databricks_base.requests") as mock_requests: mock_requests.codes.ok = 200 mock_requests.get.return_value.json.return_value = { "endpoints": [ @@ -68,26 +200,41 @@ def test_query(self, mock_requests, mock_conn): } status_code_mock = mock.PropertyMock(return_value=200) type(mock_requests.get.return_value).status_code = status_code_mock - - test_fields = ["id", "value"] - test_description = [(field,) for field in test_fields] - - conn = mock_conn.return_value - cur = mock.MagicMock(rowcount=0, description=test_description) - cur.fetchall.return_value = [] - conn.cursor.return_value = cur - - query = "select * from test.test;" - results = self.hook.run(sql=query, handler=fetch_all_handler) - - assert self.hook.last_description == test_description - assert results == [] - - cur.execute.assert_has_calls([mock.call(q) for q in [query]]) + connections = [] + cursors = [] + for index in range(len(cursor_descriptions)): + conn = mock.MagicMock() + cur = mock.MagicMock( + rowcount=len(cursor_results[index]), + description=get_cursor_descriptions(cursor_descriptions[index]), + ) + cur.fetchall.return_value = cursor_results[index] + conn.cursor.return_value = cur + cursors.append(cur) + connections.append(conn) + mock_conn.side_effect = connections + results = databricks_hook.run( + sql=sql, handler=fetch_all_handler, return_last=return_last, split_statements=split_statements + ) + + assert databricks_hook.descriptions == hook_descriptions + assert databricks_hook.last_description == hook_descriptions[-1] + assert results == hook_results + + for index, cur in enumerate(cursors): + cur.execute.assert_has_calls([mock.call(cursor_calls[index])]) cur.close.assert_called() - def test_no_query(self): - for empty_statement in ([], "", "\n"): - with pytest.raises(ValueError) as err: - self.hook.run(sql=empty_statement) - assert err.value.args[0] == "List of SQL statements is empty" + +@pytest.mark.parametrize( + "empty_statement", + [ + pytest.param([], id="Empty list"), + pytest.param("", id="Empty string"), + pytest.param("\n", id="Only EOL"), + ], +) +def test_no_query(databricks_hook, empty_statement): + with pytest.raises(ValueError) as err: + databricks_hook.run(sql=empty_statement) + assert err.value.args[0] == "List of SQL statements is empty" diff --git a/tests/providers/databricks/operators/test_databricks_copy.py b/tests/providers/databricks/operators/test_databricks_copy.py new file mode 100644 index 0000000000000..510e0a6c2fae9 --- /dev/null +++ b/tests/providers/databricks/operators/test_databricks_copy.py @@ -0,0 +1,230 @@ +# +# 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 + +import pytest + +from airflow import AirflowException +from airflow.providers.databricks.operators.databricks_sql import DatabricksCopyIntoOperator + +DATE = "2017-04-20" +TASK_ID = "databricks-sql-operator" +DEFAULT_CONN_ID = "databricks_default" +COPY_FILE_LOCATION = "s3://my-bucket/jsonData" + + +def test_copy_with_files(): + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + files=["file1", "file2", "file3"], + format_options={"dateFormat": "yyyy-MM-dd"}, + task_id=TASK_ID, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM '{COPY_FILE_LOCATION}' +FILEFORMAT = JSON +FILES = ('file1','file2','file3') +FORMAT_OPTIONS ('dateFormat' = 'yyyy-MM-dd') +""".strip() + ) + + +def test_copy_with_expression(): + expression = "col1, col2" + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="CSV", + table_name="test", + task_id=TASK_ID, + pattern="folder1/file_[a-g].csv", + expression_list=expression, + format_options={"header": "true"}, + force_copy=True, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM (SELECT {expression} FROM '{COPY_FILE_LOCATION}') +FILEFORMAT = CSV +PATTERN = 'folder1/file_[a-g].csv' +FORMAT_OPTIONS ('header' = 'true') +COPY_OPTIONS ('force' = 'true') +""".strip() + ) + + +def test_copy_with_credential(): + expression = "col1, col2" + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="CSV", + table_name="test", + task_id=TASK_ID, + expression_list=expression, + credential={"AZURE_SAS_TOKEN": "abc"}, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM (SELECT {expression} FROM '{COPY_FILE_LOCATION}' WITH (CREDENTIAL (AZURE_SAS_TOKEN = 'abc') )) +FILEFORMAT = CSV +""".strip() + ) + + +def test_copy_with_target_credential(): + expression = "col1, col2" + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="CSV", + table_name="test", + task_id=TASK_ID, + expression_list=expression, + storage_credential="abc", + credential={"AZURE_SAS_TOKEN": "abc"}, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test WITH (CREDENTIAL abc) +FROM (SELECT {expression} FROM '{COPY_FILE_LOCATION}' WITH (CREDENTIAL (AZURE_SAS_TOKEN = 'abc') )) +FILEFORMAT = CSV +""".strip() + ) + + +def test_copy_with_encryption(): + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="CSV", + table_name="test", + task_id=TASK_ID, + encryption={"TYPE": "AWS_SSE_C", "MASTER_KEY": "abc"}, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM '{COPY_FILE_LOCATION}' WITH ( ENCRYPTION (TYPE = 'AWS_SSE_C', MASTER_KEY = 'abc')) +FILEFORMAT = CSV +""".strip() + ) + + +def test_copy_with_encryption_and_credential(): + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="CSV", + table_name="test", + task_id=TASK_ID, + encryption={"TYPE": "AWS_SSE_C", "MASTER_KEY": "abc"}, + credential={"AZURE_SAS_TOKEN": "abc"}, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM '{COPY_FILE_LOCATION}' WITH (CREDENTIAL (AZURE_SAS_TOKEN = 'abc') """ + """ENCRYPTION (TYPE = 'AWS_SSE_C', MASTER_KEY = 'abc')) +FILEFORMAT = CSV +""".strip() + ) + + +def test_copy_with_validate_all(): + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + task_id=TASK_ID, + validate=True, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM '{COPY_FILE_LOCATION}' +FILEFORMAT = JSON +VALIDATE ALL +""".strip() + ) + + +def test_copy_with_validate_N_rows(): + op = DatabricksCopyIntoOperator( + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + task_id=TASK_ID, + validate=10, + ) + assert ( + op._create_sql_query() + == f"""COPY INTO test +FROM '{COPY_FILE_LOCATION}' +FILEFORMAT = JSON +VALIDATE 10 ROWS +""".strip() + ) + + +def test_incorrect_params_files_patterns(): + exception_message = "Only one of 'pattern' or 'files' should be specified" + with pytest.raises(AirflowException, match=exception_message): + DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="test", + files=["file1", "file2", "file3"], + pattern="abc", + ) + + +def test_incorrect_params_emtpy_table(): + exception_message = "table_name shouldn't be empty" + with pytest.raises(AirflowException, match=exception_message): + DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format="JSON", + table_name="", + ) + + +def test_incorrect_params_emtpy_location(): + exception_message = "file_location shouldn't be empty" + with pytest.raises(AirflowException, match=exception_message): + DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location="", + file_format="JSON", + table_name="abc", + ) + + +def test_incorrect_params_wrong_format(): + file_format = "JSONL" + exception_message = f"file_format '{file_format}' isn't supported" + with pytest.raises(AirflowException, match=exception_message): + DatabricksCopyIntoOperator( + task_id=TASK_ID, + file_location=COPY_FILE_LOCATION, + file_format=file_format, + table_name="abc", + ) diff --git a/tests/providers/databricks/operators/test_databricks_sql.py b/tests/providers/databricks/operators/test_databricks_sql.py index 2663deeec20fb..8489f45095d25 100644 --- a/tests/providers/databricks/operators/test_databricks_sql.py +++ b/tests/providers/databricks/operators/test_databricks_sql.py @@ -19,42 +19,118 @@ import os import tempfile -from unittest import mock +from unittest.mock import patch import pytest from databricks.sql.types import Row -from airflow import AirflowException from airflow.providers.common.sql.hooks.sql import fetch_all_handler -from airflow.providers.databricks.operators.databricks_sql import ( - DatabricksCopyIntoOperator, - DatabricksSqlOperator, -) +from airflow.providers.databricks.operators.databricks_sql import DatabricksSqlOperator DATE = "2017-04-20" TASK_ID = "databricks-sql-operator" DEFAULT_CONN_ID = "databricks_default" -COPY_FILE_LOCATION = "s3://my-bucket/jsonData" -class TestDatabricksSqlOperator: - @mock.patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") - def test_exec_success(self, db_mock_class): - """ - Test the execute function in case where SQL query was successful. - """ - sql = "select * from dummy" - op = DatabricksSqlOperator(task_id=TASK_ID, sql=sql, do_xcom_push=True) +@pytest.mark.parametrize( + "sql, return_last, split_statement, hook_results, hook_descriptions, expected_results", + [ + pytest.param( + "select * from dummy", + True, + True, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Single SQL statement, return_last, split statement", + ), + pytest.param( + "select * from dummy;select * from dummy2", + True, + True, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Multiple SQL statements, return_last, split statement", + ), + pytest.param( + "select * from dummy", + False, + False, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Single SQL statements, no return_last (doesn't matter), no split statement", + ), + pytest.param( + "select * from dummy", + True, + False, + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [[("id",), ("value",)]], + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + id="Scalar: Single SQL statements, return_last (doesn't matter), no split statement", + ), + pytest.param( + ["select * from dummy"], + False, + False, + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + [[("id",), ("value",)]], + [([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")])], + id="Non-Scalar: Single SQL statements in list, no return_last, no split statement", + ), + pytest.param( + ["select * from dummy", "select * from dummy2"], + False, + False, + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")], + ], + [[("id",), ("value",)], [("id2",), ("value2",)]], + [ + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + ([("id2",), ("value2",)], [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")]), + ], + id="Non-Scalar: Multiple SQL statements in list, no return_last (no matter), no split statement", + ), + pytest.param( + ["select * from dummy", "select * from dummy2"], + True, + False, + [ + [Row(id=1, value="value1"), Row(id=2, value="value2")], + [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")], + ], + [[("id",), ("value",)], [("id2",), ("value2",)]], + [ + ([("id",), ("value",)], [Row(id=1, value="value1"), Row(id=2, value="value2")]), + ([("id2",), ("value2",)], [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")]), + ], + id="Non-Scalar: Multiple SQL statements in list, return_last (no matter), no split statement", + ), + ], +) +def test_exec_success(sql, return_last, split_statement, hook_results, hook_descriptions, expected_results): + """ + Test the execute function in case where SQL query was successful. + """ + with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql=sql, + do_xcom_push=True, + return_last=return_last, + split_statements=split_statement, + ) db_mock = db_mock_class.return_value - mock_description = [("id",), ("value",)] - mock_results = [Row(id=1, value="value1")] - db_mock.run.return_value = mock_results - db_mock.last_description = mock_description - db_mock.scalar_return_last = False + db_mock.run.return_value = hook_results + db_mock.descriptions = hook_descriptions execute_results = op.execute(None) - assert execute_results == (mock_description, mock_results) + assert execute_results == expected_results db_mock_class.assert_called_once_with( DEFAULT_CONN_ID, http_path=None, @@ -70,32 +146,116 @@ def test_exec_success(self, db_mock_class): parameters=None, handler=fetch_all_handler, autocommit=False, - return_last=True, - split_statements=False, + return_last=return_last, + split_statements=split_statement, ) - @mock.patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") - def test_exec_write_file(self, db_mock_class): - """ - Test the execute function in case where SQL query was successful and data is written as CSV - """ - sql = "select * from dummy" + +@pytest.mark.parametrize( + "return_last, split_statements, sql, descriptions, hook_results", + [ + pytest.param( + True, + False, + "select * from dummy", + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: return_last True and split_statement False", + ), + pytest.param( + False, + True, + "select * from dummy", + [[("id",), ("value",)]], + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + id="Non-Scalar: return_last False and split_statement True", + ), + pytest.param( + True, + True, + "select * from dummy", + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: return_last True and no split_statement True", + ), + pytest.param( + False, + False, + "select * from dummy", + [[("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: return_last False and split_statement is False", + ), + pytest.param( + False, + True, + "select * from dummy2; select * from dummy", + [[("id2",), ("value2",)], [("id",), ("value",)]], + [ + [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + ], + id="Non-Scalar: return_last False and split_statement is True", + ), + pytest.param( + True, + True, + "select * from dummy2; select * from dummy", + [[("id2",), ("value2",)], [("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: return_last True and split_statement is True", + ), + pytest.param( + True, + True, + "select * from dummy2; select * from dummy", + [[("id2",), ("value2",)], [("id",), ("value",)]], + [Row(id=1, value="value1"), Row(id=2, value="value2")], + id="Scalar: return_last True and split_statement is True", + ), + pytest.param( + True, + True, + ["select * from dummy2", "select * from dummy"], + [[("id2",), ("value2",)], [("id",), ("value",)]], + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + id="Non-Scalar: sql is list and return_last is True", + ), + pytest.param( + False, + True, + ["select * from dummy2", "select * from dummy"], + [[("id2",), ("value2",)], [("id",), ("value",)]], + [[Row(id=1, value="value1"), Row(id=2, value="value2")]], + id="Non-Scalar: sql is list and return_last is False", + ), + ], +) +def test_exec_write_file(return_last, split_statements, sql, descriptions, hook_results): + """ + Test the execute function in case where SQL query was successful and data is written as CSV + """ + with patch("airflow.providers.databricks.operators.databricks_sql.DatabricksSqlHook") as db_mock_class: tempfile_path = tempfile.mkstemp()[1] - op = DatabricksSqlOperator(task_id=TASK_ID, sql=sql, output_path=tempfile_path) + op = DatabricksSqlOperator( + task_id=TASK_ID, + sql=sql, + output_path=tempfile_path, + return_last=return_last, + split_statements=split_statements, + ) db_mock = db_mock_class.return_value - mock_description = [("id",), ("value",)] - mock_results = [Row(id=1, value="value1")] + mock_results = hook_results db_mock.run.return_value = mock_results - db_mock.last_description = mock_description - db_mock.scalar_return_last = False + db_mock.descriptions = descriptions try: op.execute(None) results = [line.strip() for line in open(tempfile_path)] finally: os.remove(tempfile_path) - - assert results == ["id,value", "1,value1"] + # In all cases only result of last query i output as file + assert results == ["id,value", "1,value1", "2,value2"] db_mock_class.assert_called_once_with( DEFAULT_CONN_ID, http_path=None, @@ -111,198 +271,6 @@ def test_exec_write_file(self, db_mock_class): parameters=None, handler=fetch_all_handler, autocommit=False, - return_last=True, - split_statements=False, - ) - - -class TestDatabricksSqlCopyIntoOperator: - def test_copy_with_files(self): - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="JSON", - table_name="test", - files=["file1", "file2", "file3"], - format_options={"dateFormat": "yyyy-MM-dd"}, - task_id=TASK_ID, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM '{COPY_FILE_LOCATION}' -FILEFORMAT = JSON -FILES = ('file1','file2','file3') -FORMAT_OPTIONS ('dateFormat' = 'yyyy-MM-dd') -""".strip() - ) - - def test_copy_with_expression(self): - expression = "col1, col2" - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="CSV", - table_name="test", - task_id=TASK_ID, - pattern="folder1/file_[a-g].csv", - expression_list=expression, - format_options={"header": "true"}, - force_copy=True, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM (SELECT {expression} FROM '{COPY_FILE_LOCATION}') -FILEFORMAT = CSV -PATTERN = 'folder1/file_[a-g].csv' -FORMAT_OPTIONS ('header' = 'true') -COPY_OPTIONS ('force' = 'true') -""".strip() - ) - - def test_copy_with_credential(self): - expression = "col1, col2" - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="CSV", - table_name="test", - task_id=TASK_ID, - expression_list=expression, - credential={"AZURE_SAS_TOKEN": "abc"}, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM (SELECT {expression} FROM '{COPY_FILE_LOCATION}' WITH (CREDENTIAL (AZURE_SAS_TOKEN = 'abc') )) -FILEFORMAT = CSV -""".strip() + return_last=return_last, + split_statements=split_statements, ) - - def test_copy_with_target_credential(self): - expression = "col1, col2" - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="CSV", - table_name="test", - task_id=TASK_ID, - expression_list=expression, - storage_credential="abc", - credential={"AZURE_SAS_TOKEN": "abc"}, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test WITH (CREDENTIAL abc) -FROM (SELECT {expression} FROM '{COPY_FILE_LOCATION}' WITH (CREDENTIAL (AZURE_SAS_TOKEN = 'abc') )) -FILEFORMAT = CSV -""".strip() - ) - - def test_copy_with_encryption(self): - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="CSV", - table_name="test", - task_id=TASK_ID, - encryption={"TYPE": "AWS_SSE_C", "MASTER_KEY": "abc"}, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM '{COPY_FILE_LOCATION}' WITH ( ENCRYPTION (TYPE = 'AWS_SSE_C', MASTER_KEY = 'abc')) -FILEFORMAT = CSV -""".strip() - ) - - def test_copy_with_encryption_and_credential(self): - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="CSV", - table_name="test", - task_id=TASK_ID, - encryption={"TYPE": "AWS_SSE_C", "MASTER_KEY": "abc"}, - credential={"AZURE_SAS_TOKEN": "abc"}, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM '{COPY_FILE_LOCATION}' WITH (CREDENTIAL (AZURE_SAS_TOKEN = 'abc') """ - """ENCRYPTION (TYPE = 'AWS_SSE_C', MASTER_KEY = 'abc')) -FILEFORMAT = CSV -""".strip() - ) - - def test_copy_with_validate_all(self): - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="JSON", - table_name="test", - task_id=TASK_ID, - validate=True, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM '{COPY_FILE_LOCATION}' -FILEFORMAT = JSON -VALIDATE ALL -""".strip() - ) - - def test_copy_with_validate_N_rows(self): - op = DatabricksCopyIntoOperator( - file_location=COPY_FILE_LOCATION, - file_format="JSON", - table_name="test", - task_id=TASK_ID, - validate=10, - ) - assert ( - op._create_sql_query() - == f"""COPY INTO test -FROM '{COPY_FILE_LOCATION}' -FILEFORMAT = JSON -VALIDATE 10 ROWS -""".strip() - ) - - def test_incorrect_params_files_patterns(self): - exception_message = "Only one of 'pattern' or 'files' should be specified" - with pytest.raises(AirflowException, match=exception_message): - DatabricksCopyIntoOperator( - task_id=TASK_ID, - file_location=COPY_FILE_LOCATION, - file_format="JSON", - table_name="test", - files=["file1", "file2", "file3"], - pattern="abc", - ) - - def test_incorrect_params_emtpy_table(self): - exception_message = "table_name shouldn't be empty" - with pytest.raises(AirflowException, match=exception_message): - DatabricksCopyIntoOperator( - task_id=TASK_ID, - file_location=COPY_FILE_LOCATION, - file_format="JSON", - table_name="", - ) - - def test_incorrect_params_emtpy_location(self): - exception_message = "file_location shouldn't be empty" - with pytest.raises(AirflowException, match=exception_message): - DatabricksCopyIntoOperator( - task_id=TASK_ID, - file_location="", - file_format="JSON", - table_name="abc", - ) - - def test_incorrect_params_wrong_format(self): - file_format = "JSONL" - exception_message = f"file_format '{file_format}' isn't supported" - with pytest.raises(AirflowException, match=exception_message): - DatabricksCopyIntoOperator( - task_id=TASK_ID, - file_location=COPY_FILE_LOCATION, - file_format=file_format, - table_name="abc", - ) diff --git a/tests/providers/jdbc/operators/test_jdbc.py b/tests/providers/jdbc/operators/test_jdbc.py index e027bdb96df01..1b149bbb1fc4b 100644 --- a/tests/providers/jdbc/operators/test_jdbc.py +++ b/tests/providers/jdbc/operators/test_jdbc.py @@ -50,5 +50,7 @@ def test_execute_dont_push(self, mock_get_db_hook): sql=jdbc_operator.sql, autocommit=jdbc_operator.autocommit, parameters=jdbc_operator.parameters, + handler=None, + return_last=True, split_statements=False, )