feat: Add OpenLineage support for CloudSQLExecuteQueryOperator (apach…

…e#45182)

Signed-off-by: Kacper Muda <[email protected]>

providers/src/airflow/providers/google/cloud/hooks/cloud_sql.py

@@ -67,6 +67,8 @@
     from google.cloud.secretmanager_v1 import AccessSecretVersionResponse
     from requests import Session
 
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+
 UNIX_PATH_MAX = 108
 
 # Time to sleep between active checks of the operation results
@@ -1148,7 +1150,7 @@ def get_sqlproxy_runner(self) -> CloudSqlProxyRunner:
             gcp_conn_id=self.gcp_conn_id,
         )
 
-    def get_database_hook(self, connection: Connection) -> BaseHook:
+    def get_database_hook(self, connection: Connection) -> DbApiHook:
         """
         Retrieve database hook.
 
@@ -1158,7 +1160,7 @@ def get_database_hook(self, connection: Connection) -> BaseHook:
         if self.database_type == "postgres":
             from airflow.providers.postgres.hooks.postgres import PostgresHook
 
-            db_hook: BaseHook = PostgresHook(connection=connection, database=self.database)
+            db_hook: DbApiHook = PostgresHook(connection=connection, database=self.database)
         else:
             from airflow.providers.mysql.hooks.mysql import MySqlHook
 

providers/src/airflow/providers/google/cloud/operators/cloud_sql.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Mapping, Sequence
+from contextlib import contextmanager
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
@@ -38,8 +39,7 @@
 
 if TYPE_CHECKING:
     from airflow.models import Connection
-    from airflow.providers.mysql.hooks.mysql import MySqlHook
-    from airflow.providers.postgres.hooks.postgres import PostgresHook
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -1256,7 +1256,8 @@ def __init__(
         self.ssl_client_key = ssl_client_key
         self.ssl_secret_id = ssl_secret_id
 
-    def _execute_query(self, hook: CloudSQLDatabaseHook, database_hook: PostgresHook | MySqlHook) -> None:
+    @contextmanager
+    def cloud_sql_proxy_context(self, hook: CloudSQLDatabaseHook):
         cloud_sql_proxy_runner = None
         try:
             if hook.use_proxy:
@@ -1266,27 +1267,27 @@ def _execute_query(self, hook: CloudSQLDatabaseHook, database_hook: PostgresHook
                 # be taken over here by another bind(0).
                 # It's quite unlikely to happen though!
                 cloud_sql_proxy_runner.start_proxy()
-            self.log.info('Executing: "%s"', self.sql)
-            database_hook.run(self.sql, self.autocommit, parameters=self.parameters)
+            yield
         finally:
             if cloud_sql_proxy_runner:
                 cloud_sql_proxy_runner.stop_proxy()
 
     def execute(self, context: Context):
-        self.gcp_connection = BaseHook.get_connection(self.gcp_conn_id)
-
         hook = self.hook
         hook.validate_ssl_certs()
         connection = hook.create_connection()
         hook.validate_socket_path_length()
         database_hook = hook.get_database_hook(connection=connection)
         try:
-            self._execute_query(hook, database_hook)
+            with self.cloud_sql_proxy_context(hook):
+                self.log.info('Executing: "%s"', self.sql)
+                database_hook.run(self.sql, self.autocommit, parameters=self.parameters)
         finally:
             hook.cleanup_database_hook()
 
     @cached_property
     def hook(self):
+        self.gcp_connection = BaseHook.get_connection(self.gcp_conn_id)
         return CloudSQLDatabaseHook(
             gcp_cloudsql_conn_id=self.gcp_cloudsql_conn_id,
             gcp_conn_id=self.gcp_conn_id,
@@ -1297,3 +1298,14 @@ def hook(self):
             ssl_key=self.ssl_client_key,
             ssl_secret_id=self.ssl_secret_id,
         )
+
+    def get_openlineage_facets_on_complete(self, _) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+
+        with self.cloud_sql_proxy_context(self.hook):
+            return get_openlineage_facets_with_sql(
+                hook=self.hook.db_hook,
+                sql=self.sql,  # type:ignore[arg-type]  # Iterable[str] instead of list[str]
+                conn_id=self.gcp_cloudsql_conn_id,
+                database=self.hook.database,
+            )

providers/tests/common/compat/openlineage/utils/test_sql.py

@@ -0,0 +1,23 @@
+# 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
+
+
+def test_import():
+    from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+
+    assert get_openlineage_facets_with_sql is not None

providers/tests/google/cloud/operators/test_cloud_sql.py

@@ -19,11 +19,19 @@
 
 import os
 from unittest import mock
+from unittest.mock import MagicMock
 
 import pytest
 
 from airflow.exceptions import AirflowException, TaskDeferred
 from airflow.models import Connection
+from airflow.providers.common.compat.openlineage.facet import (
+    Dataset,
+    SchemaDatasetFacet,
+    SchemaDatasetFacetFields,
+    SQLJobFacet,
+)
+from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.cloud.operators.cloud_sql import (
     CloudSQLCloneInstanceOperator,
     CloudSQLCreateInstanceDatabaseOperator,
@@ -822,3 +830,66 @@ def test_create_operator_with_too_long_unix_socket_path(self, get_connection):
             operator.execute(None)
         err = ctx.value
         assert "The UNIX socket path length cannot exceed" in str(err)
+
+    @pytest.mark.parametrize(
+        "connection_port, default_port, expected_port",
+        [(None, 4321, 4321), (1234, None, 1234), (1234, 4321, 1234)],
+    )
+    def test_execute_openlineage_events(self, connection_port, default_port, expected_port):
+        class DBApiHookForTests(DbApiHook):
+            conn_name_attr = "sql_default"
+            get_conn = MagicMock(name="conn")
+            get_connection = MagicMock()
+
+            def get_openlineage_database_info(self, connection):
+                from airflow.providers.openlineage.sqlparser import DatabaseInfo
+
+                return DatabaseInfo(
+                    scheme="sqlscheme",
+                    authority=DbApiHook.get_openlineage_authority_part(connection, default_port=default_port),
+                )
+
+        dbapi_hook = DBApiHookForTests()
+
+        class CloudSQLExecuteQueryOperatorForTest(CloudSQLExecuteQueryOperator):
+            @property
+            def hook(self):
+                return MagicMock(db_hook=dbapi_hook, database="")
+
+        sql = """CREATE TABLE IF NOT EXISTS popular_orders_day_of_week (
+            order_day_of_week VARCHAR(64) NOT NULL,
+            order_placed_on   TIMESTAMP NOT NULL,
+            orders_placed     INTEGER NOT NULL
+        );
+FORGOT TO COMMENT"""
+        op = CloudSQLExecuteQueryOperatorForTest(task_id="task_id", sql=sql)
+        DB_SCHEMA_NAME = "PUBLIC"
+        rows = [
+            (DB_SCHEMA_NAME, "popular_orders_day_of_week", "order_day_of_week", 1, "varchar"),
+            (DB_SCHEMA_NAME, "popular_orders_day_of_week", "order_placed_on", 2, "timestamp"),
+            (DB_SCHEMA_NAME, "popular_orders_day_of_week", "orders_placed", 3, "int4"),
+        ]
+        dbapi_hook.get_connection.return_value = Connection(
+            conn_id="sql_default", conn_type="postgresql", host="host", port=connection_port
+        )
+        dbapi_hook.get_conn.return_value.cursor.return_value.fetchall.side_effect = [rows, []]
+
+        lineage = op.get_openlineage_facets_on_complete(None)
+        assert len(lineage.inputs) == 0
+        assert lineage.job_facets == {"sql": SQLJobFacet(query=sql)}
+        assert lineage.run_facets["extractionError"].failedTasks == 1
+        assert lineage.outputs == [
+            Dataset(
+                namespace=f"sqlscheme://host:{expected_port}",
+                name="PUBLIC.popular_orders_day_of_week",
+                facets={
+                    "schema": SchemaDatasetFacet(
+                        fields=[
+                            SchemaDatasetFacetFields(name="order_day_of_week", type="varchar"),
+                            SchemaDatasetFacetFields(name="order_placed_on", type="timestamp"),
+                            SchemaDatasetFacetFields(name="orders_placed", type="int4"),
+                        ]
+                    )
+                },
+            )
+        ]