Bring back accidentally removed test 'test_execute_openlineage_events()'

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
 .....
 

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

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,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 = 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", conn_id="trino_default"
+        )
         op.execute({})
         lineage = op.get_openlineage_facets_on_start()
         assert lineage.inputs[0].namespace == "trino://trino:8080"

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"),
+                    ]
+                )
+            },
+        )
+    ]