From dbc45645365cc9a278e444b933b9038e0abeb909 Mon Sep 17 00:00:00 2001 From: pratiksha rajendrabhai badheka Date: Fri, 6 Dec 2024 13:20:30 +0530 Subject: [PATCH 1/5] remove provider deprecations in trino --- .../src/airflow/providers/trino/CHANGELOG.rst | 9 ++ .../providers/trino/operators/trino.py | 60 ------------ providers/tests/trino/operators/test_trino.py | 95 ------------------- 3 files changed, 9 insertions(+), 155 deletions(-) diff --git a/providers/src/airflow/providers/trino/CHANGELOG.rst b/providers/src/airflow/providers/trino/CHANGELOG.rst index 42bf0211e1c93..c7bb3d14ef166 100644 --- a/providers/src/airflow/providers/trino/CHANGELOG.rst +++ b/providers/src/airflow/providers/trino/CHANGELOG.rst @@ -27,6 +27,15 @@ Changelog --------- +main +.... + +.. warning:: + All deprecated classes, parameters and features have been removed from the {provider_name} provider package. + The following breaking changes were introduced: + + * Remove ``TrinoOperator`` class from airflow.providers.trino.operators.trino. Please use ``airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator``` + 5.9.0 ..... diff --git a/providers/src/airflow/providers/trino/operators/trino.py b/providers/src/airflow/providers/trino/operators/trino.py index 5a3858342a523..92f81acbf2c73 100644 --- a/providers/src/airflow/providers/trino/operators/trino.py +++ b/providers/src/airflow/providers/trino/operators/trino.py @@ -18,63 +18,3 @@ """This module contains the Trino operator.""" from __future__ import annotations - -from collections.abc import Sequence -from typing import Any, ClassVar - -from deprecated import deprecated -from trino.exceptions import TrinoQueryError - -from airflow.exceptions import AirflowProviderDeprecationWarning -from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator -from airflow.providers.trino.hooks.trino import TrinoHook - - -@deprecated( - reason="Please use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator`.", - category=AirflowProviderDeprecationWarning, -) -class TrinoOperator(SQLExecuteQueryOperator): - """ - Executes sql code using a specific Trino query Engine. - - This class is deprecated. - - Please use :class:`airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator`. - - .. seealso:: - For more information on how to use this operator, take a look at the guide: - :ref:`howto/operator:TrinoOperator` - - :param sql: the SQL code to be executed as a single string, or - a list of str (sql statements), or a reference to a template file. - :param trino_conn_id: id of the connection config for the target Trino - environment - :param autocommit: What to set the connection's autocommit setting to - before executing the query - :param handler: The result handler which is called with the result of each statement. - :param parameters: (optional) the parameters to render the SQL query with. - """ - - template_fields: Sequence[str] = ("sql",) - template_fields_renderers: ClassVar[dict] = {"sql": "sql"} - template_ext: Sequence[str] = (".sql",) - ui_color = "#ededed" - - def __init__(self, *, trino_conn_id: str = "trino_default", **kwargs: Any) -> None: - super().__init__(conn_id=trino_conn_id, **kwargs) - - def on_kill(self) -> None: - if self._hook is not None and isinstance(self._hook, TrinoHook): # type: ignore[attr-defined] - query_id = "'" + self._hook.query_id + "'" # type: ignore[attr-defined] - try: - self.log.info("Stopping query run with queryId - %s", self._hook.query_id) # type: ignore[attr-defined] - self._hook.run( # type: ignore[attr-defined] - sql=f"CALL system.runtime.kill_query(query_id => {query_id},message => 'Job " - f"killed by " - f"user');", - handler=list, - ) - except TrinoQueryError as e: - self.log.info(str(e)) - self.log.info("Trino query (%s) terminated", query_id) diff --git a/providers/tests/trino/operators/test_trino.py b/providers/tests/trino/operators/test_trino.py index 24d933bf0d360..03cb33c14c40e 100644 --- a/providers/tests/trino/operators/test_trino.py +++ b/providers/tests/trino/operators/test_trino.py @@ -16,98 +16,3 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - -from unittest import mock - -from airflow.models.connection import Connection -from airflow.providers.common.compat.openlineage.facet import ( - Dataset, - SchemaDatasetFacet, - SchemaDatasetFacetFields, - SQLJobFacet, -) -from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator -from airflow.providers.trino.hooks.trino import TrinoHook - -TRINO_CONN_ID = "test_trino" -TASK_ID = "test_trino_task" -TRINO_DEFAULT = "trino_default" - - -class TestTrinoOperator: - @mock.patch("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator.get_db_hook") - def test_execute(self, mock_get_db_hook): - """Asserts that the run method is called when a TrinoOperator task is executed""" - - op = SQLExecuteQueryOperator( - task_id=TASK_ID, - sql="SELECT 1;", - conn_id=TRINO_CONN_ID, - handler=list, - ) - op.execute(None) - - mock_get_db_hook.return_value.run.assert_called_once_with( - sql="SELECT 1;", - autocommit=False, - handler=list, - parameters=None, - return_last=True, - ) - - -def test_execute_openlineage_events(): - DB_NAME = "tpch" - DB_SCHEMA_NAME = "sf1" - - class TrinoHookForTests(TrinoHook): - get_conn = mock.MagicMock(name="conn") - get_connection = mock.MagicMock() - - def get_first(self, *_): - return [f"{DB_NAME}.{DB_SCHEMA_NAME}"] - - dbapi_hook = TrinoHookForTests() - - sql = "SELECT name FROM tpch.sf1.customer LIMIT 3" - op = SQLExecuteQueryOperator(task_id="trino-operator", sql=sql, conn_id=TRINO_DEFAULT) - op._hook = dbapi_hook - rows = [ - (DB_SCHEMA_NAME, "customer", "custkey", 1, "bigint", DB_NAME), - (DB_SCHEMA_NAME, "customer", "name", 2, "varchar(25)", DB_NAME), - (DB_SCHEMA_NAME, "customer", "address", 3, "varchar(40)", DB_NAME), - (DB_SCHEMA_NAME, "customer", "nationkey", 4, "bigint", DB_NAME), - (DB_SCHEMA_NAME, "customer", "phone", 5, "varchar(15)", DB_NAME), - (DB_SCHEMA_NAME, "customer", "acctbal", 6, "double", DB_NAME), - ] - dbapi_hook.get_connection.return_value = Connection( - conn_id=TRINO_DEFAULT, - conn_type="trino", - host="trino", - port=8080, - ) - dbapi_hook.get_conn.return_value.cursor.return_value.fetchall.side_effect = [rows, []] - - lineage = op.get_openlineage_facets_on_start() - assert lineage.inputs == [ - Dataset( - namespace="trino://trino:8080", - name=f"{DB_NAME}.{DB_SCHEMA_NAME}.customer", - facets={ - "schema": SchemaDatasetFacet( - fields=[ - SchemaDatasetFacetFields(name="custkey", type="bigint"), - SchemaDatasetFacetFields(name="name", type="varchar(25)"), - SchemaDatasetFacetFields(name="address", type="varchar(40)"), - SchemaDatasetFacetFields(name="nationkey", type="bigint"), - SchemaDatasetFacetFields(name="phone", type="varchar(15)"), - SchemaDatasetFacetFields(name="acctbal", type="double"), - ] - ) - }, - ) - ] - - assert len(lineage.outputs) == 0 - - assert lineage.job_facets == {"sql": SQLJobFacet(query=sql)} From e6193a5419d443f9d44e7e567bfd6064e3f099cd Mon Sep 17 00:00:00 2001 From: pratiksha rajendrabhai badheka Date: Mon, 9 Dec 2024 13:55:32 +0530 Subject: [PATCH 2/5] remove trino.py and test_trino.py and fix mypy test --- .../providers/trino/operators/trino.py | 20 ------------------- providers/tests/trino/operators/test_trino.py | 18 ----------------- 2 files changed, 38 deletions(-) delete mode 100644 providers/src/airflow/providers/trino/operators/trino.py delete mode 100644 providers/tests/trino/operators/test_trino.py diff --git a/providers/src/airflow/providers/trino/operators/trino.py b/providers/src/airflow/providers/trino/operators/trino.py deleted file mode 100644 index 92f81acbf2c73..0000000000000 --- a/providers/src/airflow/providers/trino/operators/trino.py +++ /dev/null @@ -1,20 +0,0 @@ -# -# 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 module contains the Trino operator.""" - -from __future__ import annotations diff --git a/providers/tests/trino/operators/test_trino.py b/providers/tests/trino/operators/test_trino.py deleted file mode 100644 index 03cb33c14c40e..0000000000000 --- a/providers/tests/trino/operators/test_trino.py +++ /dev/null @@ -1,18 +0,0 @@ -# -# 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 49cae7860b93c1daf53e765b88c21a2374acd643 Mon Sep 17 00:00:00 2001 From: pratiksha rajendrabhai badheka Date: Mon, 9 Dec 2024 17:07:45 +0530 Subject: [PATCH 3/5] add trino.py and test_trino.py to fix static checks --- .../airflow/providers/trino/operators/trino.py | 17 +++++++++++++++++ providers/tests/trino/operators/test_trino.py | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) create mode 100644 providers/src/airflow/providers/trino/operators/trino.py create mode 100644 providers/tests/trino/operators/test_trino.py diff --git a/providers/src/airflow/providers/trino/operators/trino.py b/providers/src/airflow/providers/trino/operators/trino.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/src/airflow/providers/trino/operators/trino.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/tests/trino/operators/test_trino.py b/providers/tests/trino/operators/test_trino.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/tests/trino/operators/test_trino.py @@ -0,0 +1,17 @@ +# +# 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 8b9810f21cbe7b0578ec71a1219fd1ad4cffabad Mon Sep 17 00:00:00 2001 From: pratiksha rajendrabhai badheka Date: Tue, 10 Dec 2024 08:00:39 +0530 Subject: [PATCH 4/5] fix static checks --- .../providers/trino/operators/__init__.py | 16 ---------------- .../airflow/providers/trino/operators/trino.py | 17 ----------------- .../src/airflow/providers/trino/provider.yaml | 5 ----- .../tests/integration/trino/hooks/test_trino.py | 4 ++-- providers/tests/trino/operators/__init__.py | 16 ---------------- providers/tests/trino/operators/test_trino.py | 17 ----------------- 6 files changed, 2 insertions(+), 73 deletions(-) delete mode 100644 providers/src/airflow/providers/trino/operators/__init__.py delete mode 100644 providers/src/airflow/providers/trino/operators/trino.py delete mode 100644 providers/tests/trino/operators/__init__.py delete mode 100644 providers/tests/trino/operators/test_trino.py diff --git a/providers/src/airflow/providers/trino/operators/__init__.py b/providers/src/airflow/providers/trino/operators/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/providers/src/airflow/providers/trino/operators/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/providers/src/airflow/providers/trino/operators/trino.py b/providers/src/airflow/providers/trino/operators/trino.py deleted file mode 100644 index 217e5db960782..0000000000000 --- a/providers/src/airflow/providers/trino/operators/trino.py +++ /dev/null @@ -1,17 +0,0 @@ -# -# 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. diff --git a/providers/src/airflow/providers/trino/provider.yaml b/providers/src/airflow/providers/trino/provider.yaml index d94001f6d2b2e..960ce1c6c7102 100644 --- a/providers/src/airflow/providers/trino/provider.yaml +++ b/providers/src/airflow/providers/trino/provider.yaml @@ -83,11 +83,6 @@ integrations: - /docs/apache-airflow-providers-trino/operators/trino.rst tags: [software] -operators: - - integration-name: Trino - python-modules: - - airflow.providers.trino.operators.trino - asset-uris: - schemes: [trino] handler: airflow.providers.trino.assets.trino.sanitize_uri diff --git a/providers/tests/integration/trino/hooks/test_trino.py b/providers/tests/integration/trino/hooks/test_trino.py index 023989605f0b1..1faf183117b76 100644 --- a/providers/tests/integration/trino/hooks/test_trino.py +++ b/providers/tests/integration/trino/hooks/test_trino.py @@ -21,8 +21,8 @@ import pytest +from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator from airflow.providers.trino.hooks.trino import TrinoHook -from airflow.providers.trino.operators.trino import TrinoOperator @pytest.mark.integration("trino") @@ -50,7 +50,7 @@ def test_should_record_records_with_kerberos_auth(self): @mock.patch.dict("os.environ", AIRFLOW_CONN_TRINO_DEFAULT="trino://airflow@trino:8080/") def test_openlineage_methods(self): - op = TrinoOperator(task_id="trino_test", sql="SELECT name FROM tpch.sf1.customer LIMIT 3") + op = SQLExecuteQueryOperator(task_id="trino_test", sql="SELECT name FROM tpch.sf1.customer LIMIT 3") op.execute({}) lineage = op.get_openlineage_facets_on_start() assert lineage.inputs[0].namespace == "trino://trino:8080" diff --git a/providers/tests/trino/operators/__init__.py b/providers/tests/trino/operators/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/providers/tests/trino/operators/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. diff --git a/providers/tests/trino/operators/test_trino.py b/providers/tests/trino/operators/test_trino.py deleted file mode 100644 index 217e5db960782..0000000000000 --- a/providers/tests/trino/operators/test_trino.py +++ /dev/null @@ -1,17 +0,0 @@ -# -# 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 033265cb5559b7aa30472689a54f8d56cb030b31 Mon Sep 17 00:00:00 2001 From: pratiksha rajendrabhai badheka Date: Tue, 10 Dec 2024 23:08:29 +0530 Subject: [PATCH 5/5] add test method and add conn_id in SQLExecuteQueryOperator --- .../integration/trino/hooks/test_trino.py | 4 +- providers/tests/trino/hooks/test_trino.py | 61 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/providers/tests/integration/trino/hooks/test_trino.py b/providers/tests/integration/trino/hooks/test_trino.py index 1faf183117b76..d5a8b6d938a65 100644 --- a/providers/tests/integration/trino/hooks/test_trino.py +++ b/providers/tests/integration/trino/hooks/test_trino.py @@ -50,7 +50,9 @@ def test_should_record_records_with_kerberos_auth(self): @mock.patch.dict("os.environ", AIRFLOW_CONN_TRINO_DEFAULT="trino://airflow@trino:8080/") def test_openlineage_methods(self): - op = SQLExecuteQueryOperator(task_id="trino_test", sql="SELECT name FROM tpch.sf1.customer LIMIT 3") + op = SQLExecuteQueryOperator( + task_id="trino_test", sql="SELECT name FROM tpch.sf1.customer LIMIT 3", conn_id="trino_default" + ) op.execute({}) lineage = op.get_openlineage_facets_on_start() assert lineage.inputs[0].namespace == "trino://trino:8080" diff --git a/providers/tests/trino/hooks/test_trino.py b/providers/tests/trino/hooks/test_trino.py index 1cebc4eaaa469..d166c13636a0a 100644 --- a/providers/tests/trino/hooks/test_trino.py +++ b/providers/tests/trino/hooks/test_trino.py @@ -27,6 +27,12 @@ from airflow.exceptions import AirflowException from airflow.models import Connection +from airflow.providers.common.compat.openlineage.facet import ( + Dataset, + SchemaDatasetFacet, + SchemaDatasetFacetFields, +) +from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator from airflow.providers.trino.hooks.trino import TrinoHook from tests_common.test_utils.compat import AIRFLOW_V_3_0_PLUS @@ -37,6 +43,8 @@ TRINO_DBAPI_CONNECT = "airflow.providers.trino.hooks.trino.trino.dbapi.connect" JWT_AUTHENTICATION = "airflow.providers.trino.hooks.trino.trino.auth.JWTAuthentication" CERT_AUTHENTICATION = "airflow.providers.trino.hooks.trino.trino.auth.CertificateAuthentication" +TRINO_CONN_ID = "test_trino" +TRINO_DEFAULT = "trino_default" @pytest.fixture @@ -383,3 +391,56 @@ def test_connection_failure(self, mock_conn): def test_serialize_cell(self): assert self.db_hook._serialize_cell("foo", None) == "foo" assert self.db_hook._serialize_cell(1, None) == 1 + + +def test_execute_openlineage_events(): + DB_NAME = "tpch" + DB_SCHEMA_NAME = "sf1" + + class TrinoHookForTests(TrinoHook): + get_conn = mock.MagicMock(name="conn") + get_connection = mock.MagicMock() + + def get_first(self, *_): + return [f"{DB_NAME}.{DB_SCHEMA_NAME}"] + + dbapi_hook = TrinoHookForTests() + + sql = "SELECT name FROM tpch.sf1.customer LIMIT 3" + op = SQLExecuteQueryOperator(task_id="trino_test", sql=sql, conn_id=TRINO_DEFAULT) + op._hook = dbapi_hook + rows = [ + (DB_SCHEMA_NAME, "customer", "custkey", 1, "bigint", DB_NAME), + (DB_SCHEMA_NAME, "customer", "name", 2, "varchar(25)", DB_NAME), + (DB_SCHEMA_NAME, "customer", "address", 3, "varchar(40)", DB_NAME), + (DB_SCHEMA_NAME, "customer", "nationkey", 4, "bigint", DB_NAME), + (DB_SCHEMA_NAME, "customer", "phone", 5, "varchar(15)", DB_NAME), + (DB_SCHEMA_NAME, "customer", "acctbal", 6, "double", DB_NAME), + ] + dbapi_hook.get_connection.return_value = Connection( + conn_id=TRINO_DEFAULT, + conn_type="trino", + host="trino", + port=8080, + ) + dbapi_hook.get_conn.return_value.cursor.return_value.fetchall.side_effect = [rows, []] + + lineage = op.get_openlineage_facets_on_start() + assert lineage.inputs == [ + Dataset( + namespace="trino://trino:8080", + name=f"{DB_NAME}.{DB_SCHEMA_NAME}.customer", + facets={ + "schema": SchemaDatasetFacet( + fields=[ + SchemaDatasetFacetFields(name="custkey", type="bigint"), + SchemaDatasetFacetFields(name="name", type="varchar(25)"), + SchemaDatasetFacetFields(name="address", type="varchar(40)"), + SchemaDatasetFacetFields(name="nationkey", type="bigint"), + SchemaDatasetFacetFields(name="phone", type="varchar(15)"), + SchemaDatasetFacetFields(name="acctbal", type="double"), + ] + ) + }, + ) + ]