apache-airflow-providers-google 17.1.0__py3-none-any.whl → 17.2.0rc1__py3-none-any.whl

airflow/providers/google/cloud/transfers/bigquery_to_mssql.py

@@ -21,14 +21,17 @@ from __future__ import annotations
 
 import warnings
 from collections.abc import Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
 from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
 from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
 
 if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -94,9 +97,13 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
         self.mssql_conn_id = mssql_conn_id
         self.source_project_dataset_table = source_project_dataset_table
 
-    def get_sql_hook(self) -> MsSqlHook:
+    @cached_property
+    def mssql_hook(self) -> MsSqlHook:
         return MsSqlHook(schema=self.database, mssql_conn_id=self.mssql_conn_id)
 
+    def get_sql_hook(self) -> MsSqlHook:
+        return self.mssql_hook
+
     def persist_links(self, context: Context) -> None:
         project_id, dataset_id, table_id = self.source_project_dataset_table.split(".")
         BigQueryTableLink.persist(
@@ -105,3 +112,67 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
             project_id=project_id,
             table_id=table_id,
         )
+
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import (
+            BIGQUERY_NAMESPACE,
+            get_facets_from_bq_table_for_given_fields,
+            get_identity_column_lineage_facet,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        if not self.bigquery_hook:
+            self.bigquery_hook = BigQueryHook(
+                gcp_conn_id=self.gcp_conn_id,
+                location=self.location,
+                impersonation_chain=self.impersonation_chain,
+            )
+
+        try:
+            table_obj = self.bigquery_hook.get_client().get_table(self.source_project_dataset_table)
+        except Exception:
+            self.log.debug(
+                "OpenLineage: could not fetch BigQuery table %s",
+                self.source_project_dataset_table,
+                exc_info=True,
+            )
+            return OperatorLineage()
+
+        if self.selected_fields:
+            if isinstance(self.selected_fields, str):
+                bigquery_field_names = list(self.selected_fields)
+            else:
+                bigquery_field_names = self.selected_fields
+        else:
+            bigquery_field_names = [f.name for f in getattr(table_obj, "schema", [])]
+
+        input_dataset = Dataset(
+            namespace=BIGQUERY_NAMESPACE,
+            name=self.source_project_dataset_table,
+            facets=get_facets_from_bq_table_for_given_fields(table_obj, bigquery_field_names),
+        )
+
+        db_info = self.mssql_hook.get_openlineage_database_info(self.mssql_hook.get_conn())
+        default_schema = self.mssql_hook.get_openlineage_default_schema()
+        namespace = f"{db_info.scheme}://{db_info.authority}"
+
+        if self.target_table_name and "." in self.target_table_name:
+            schema_name, table_name = self.target_table_name.split(".", 1)
+        else:
+            schema_name = default_schema or ""
+            table_name = self.target_table_name or ""
+
+        if self.database:
+            output_name = f"{self.database}.{schema_name}.{table_name}"
+        else:
+            output_name = f"{schema_name}.{table_name}"
+
+        column_lineage_facet = get_identity_column_lineage_facet(
+            bigquery_field_names, input_datasets=[input_dataset]
+        )
+
+        output_facets = column_lineage_facet or {}
+        output_dataset = Dataset(namespace=namespace, name=output_name, facets=output_facets)
+
+        return OperatorLineage(inputs=[input_dataset], outputs=[output_dataset])

airflow/providers/google/cloud/transfers/bigquery_to_sql.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import abc
 from collections.abc import Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
@@ -113,19 +114,22 @@ class BigQueryToSqlBaseOperator(BaseOperator):
     def persist_links(self, context: Context) -> None:
         """Persist the connection to the SQL provider."""
 
-    def execute(self, context: Context) -> None:
-        big_query_hook = BigQueryHook(
+    @cached_property
+    def bigquery_hook(self) -> BigQueryHook:
+        return BigQueryHook(
             gcp_conn_id=self.gcp_conn_id,
             location=self.location,
             impersonation_chain=self.impersonation_chain,
         )
+
+    def execute(self, context: Context) -> None:
         self.persist_links(context)
         sql_hook = self.get_sql_hook()
         for rows in bigquery_get_data(
             self.log,
             self.dataset_id,
             self.table_id,
-            big_query_hook,
+            self.bigquery_hook,
             self.batch_size,
             self.selected_fields,
         ):

airflow/providers/google/cloud/transfers/gcs_to_bigquery.py

@@ -635,12 +635,6 @@ class GCSToBigQueryOperator(BaseOperator):
             self.configuration["load"]["schema"] = {"fields": self.schema_fields}
 
         if self.schema_update_options:
-            if self.write_disposition not in ["WRITE_APPEND", "WRITE_TRUNCATE"]:
-                raise ValueError(
-                    "schema_update_options is only "
-                    "allowed if write_disposition is "
-                    "'WRITE_APPEND' or 'WRITE_TRUNCATE'."
-                )
             # To provide backward compatibility
             self.schema_update_options = list(self.schema_update_options or [])
             self.log.info("Adding experimental 'schemaUpdateOptions': %s", self.schema_update_options)

airflow/providers/google/cloud/transfers/gcs_to_sftp.py

@@ -80,7 +80,9 @@ class GCSToSFTPOperator(BaseOperator):
     :param destination_path: The sftp remote path. This is the specified directory path for
         uploading to the SFTP server.
     :param keep_directory_structure: (Optional) When set to False the path of the file
-         on the bucket is recreated within path passed in destination_path.
+        on the bucket is recreated within path passed in destination_path.
+    :param create_intermediate_dirs: (Optional) When set to True the intermediate directories
+        in the specified file path will be created.
     :param move_object: When move object is True, the object is moved instead
         of copied to the new location. This is the equivalent of a mv command
         as opposed to a cp command.
@@ -112,6 +114,7 @@ class GCSToSFTPOperator(BaseOperator):
         source_object: str,
         destination_path: str,
         keep_directory_structure: bool = True,
+        create_intermediate_dirs: bool = True,
         move_object: bool = False,
         gcp_conn_id: str = "google_cloud_default",
         sftp_conn_id: str = "ssh_default",
@@ -124,6 +127,7 @@ class GCSToSFTPOperator(BaseOperator):
         self.source_object = source_object
         self.destination_path = destination_path
         self.keep_directory_structure = keep_directory_structure
+        self.create_intermediate_dirs = create_intermediate_dirs
         self.move_object = move_object
         self.gcp_conn_id = gcp_conn_id
         self.sftp_conn_id = sftp_conn_id
@@ -190,7 +194,9 @@ class GCSToSFTPOperator(BaseOperator):
         )
 
         dir_path = os.path.dirname(destination_path)
-        sftp_hook.create_directory(dir_path)
+
+        if self.create_intermediate_dirs:
+            sftp_hook.create_directory(dir_path)
 
         with NamedTemporaryFile("w") as tmp:
             gcs_hook.download(

airflow/providers/google/cloud/transfers/oracle_to_gcs.py

@@ -21,12 +21,17 @@ import base64
 import calendar
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from functools import cached_property
+from typing import TYPE_CHECKING
 
 import oracledb
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
 from airflow.providers.oracle.hooks.oracle import OracleHook
 
+if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+
 
 class OracleToGCSOperator(BaseSQLToGCSOperator):
     """
@@ -62,10 +67,13 @@ class OracleToGCSOperator(BaseSQLToGCSOperator):
         self.ensure_utc = ensure_utc
         self.oracle_conn_id = oracle_conn_id
 
+    @cached_property
+    def db_hook(self) -> OracleHook:
+        return OracleHook(oracle_conn_id=self.oracle_conn_id)
+
     def query(self):
         """Query Oracle and returns a cursor to the results."""
-        oracle = OracleHook(oracle_conn_id=self.oracle_conn_id)
-        conn = oracle.get_conn()
+        conn = self.db_hook.get_conn()
         cursor = conn.cursor()
         if self.ensure_utc:
             # Ensure TIMESTAMP results are in UTC
@@ -121,3 +129,20 @@ class OracleToGCSOperator(BaseSQLToGCSOperator):
             else:
                 value = base64.standard_b64encode(value).decode("ascii")
         return value
+
+    def get_openlineage_facets_on_start(self) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        sql_parsing_result = get_openlineage_facets_with_sql(
+            hook=self.db_hook,
+            sql=self.sql,
+            conn_id=self.oracle_conn_id,
+            database=self.db_hook.service_name or self.db_hook.sid,
+        )
+        gcs_output_datasets = self._get_openlineage_output_datasets()
+        if sql_parsing_result:
+            sql_parsing_result.outputs = gcs_output_datasets
+            return sql_parsing_result
+        return OperatorLineage(outputs=gcs_output_datasets, job_facets={"sql": SQLJobFacet(self.sql)})

airflow/providers/google/cloud/transfers/postgres_to_gcs.py

@@ -31,7 +31,7 @@ import pendulum
 from slugify import slugify
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
-from airflow.providers.postgres.hooks.postgres import PostgresHook
+from airflow.providers.postgres.hooks.postgres import USE_PSYCOPG3, PostgresHook
 
 if TYPE_CHECKING:
     from airflow.providers.openlineage.extractors import OperatorLineage
@@ -52,9 +52,20 @@ class _PostgresServerSideCursorDecorator:
         self.initialized = False
 
     def __iter__(self):
+        """Make the cursor iterable."""
         return self
 
     def __next__(self):
+        """Fetch next row from the cursor."""
+        if USE_PSYCOPG3:
+            if self.rows:
+                return self.rows.pop()
+            self.initialized = True
+            row = self.cursor.fetchone()
+            if row is None:
+                raise StopIteration
+            return row
+        # psycopg2
         if self.rows:
             return self.rows.pop()
         self.initialized = True
@@ -141,13 +152,29 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
         return PostgresHook(postgres_conn_id=self.postgres_conn_id)
 
     def query(self):
-        """Query Postgres and returns a cursor to the results."""
+        """Execute the query and return a cursor."""
         conn = self.db_hook.get_conn()
-        cursor = conn.cursor(name=self._unique_name())
-        cursor.execute(self.sql, self.parameters)
-        if self.use_server_side_cursor:
-            cursor.itersize = self.cursor_itersize
-            return _PostgresServerSideCursorDecorator(cursor)
+
+        if USE_PSYCOPG3:
+            from psycopg.types.json import register_default_adapters
+
+            # Register JSON handlers for this connection if not already done
+            register_default_adapters(conn)
+
+            if self.use_server_side_cursor:
+                cursor_name = f"airflow_{self.task_id.replace('-', '_')}_{uuid.uuid4().hex}"[:63]
+                cursor = conn.cursor(name=cursor_name)
+                cursor.itersize = self.cursor_itersize
+                cursor.execute(self.sql, self.parameters)
+                return _PostgresServerSideCursorDecorator(cursor)
+            cursor = conn.cursor()
+            cursor.execute(self.sql, self.parameters)
+        else:
+            cursor = conn.cursor(name=self._unique_name())
+            cursor.execute(self.sql, self.parameters)
+            if self.use_server_side_cursor:
+                cursor.itersize = self.cursor_itersize
+                return _PostgresServerSideCursorDecorator(cursor)
         return cursor
 
     def field_to_bigquery(self, field) -> dict[str, str]:
@@ -182,8 +209,14 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
                 hours=formatted_time.tm_hour, minutes=formatted_time.tm_min, seconds=formatted_time.tm_sec
             )
             return str(time_delta)
-        if stringify_dict and isinstance(value, dict):
-            return json.dumps(value)
+        if stringify_dict:
+            if USE_PSYCOPG3:
+                from psycopg.types.json import Json
+
+                if isinstance(value, (dict, Json)):
+                    return json.dumps(value)
+            elif isinstance(value, dict):
+                return json.dumps(value)
         if isinstance(value, Decimal):
             return float(value)
         return value

airflow/providers/google/cloud/triggers/bigquery.py

@@ -25,16 +25,18 @@ from aiohttp.client_exceptions import ClientResponseError
 from asgiref.sync import sync_to_async
 
 from airflow.exceptions import AirflowException
-from airflow.models.taskinstance import TaskInstance
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryAsyncHook, BigQueryTableAsyncHook
 from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from airflow.utils.session import provide_session
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
 
+if not AIRFLOW_V_3_0_PLUS:
+    from airflow.models.taskinstance import TaskInstance
+    from airflow.utils.session import provide_session
+
 
 class BigQueryInsertJobTrigger(BaseTrigger):
     """
@@ -99,24 +101,26 @@ class BigQueryInsertJobTrigger(BaseTrigger):
             },
         )
 
-    @provide_session
-    def get_task_instance(self, session: Session) -> TaskInstance:
-        query = session.query(TaskInstance).filter(
-            TaskInstance.dag_id == self.task_instance.dag_id,
-            TaskInstance.task_id == self.task_instance.task_id,
-            TaskInstance.run_id == self.task_instance.run_id,
-            TaskInstance.map_index == self.task_instance.map_index,
-        )
-        task_instance = query.one_or_none()
-        if task_instance is None:
-            raise AirflowException(
-                "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
-                self.task_instance.dag_id,
-                self.task_instance.task_id,
-                self.task_instance.run_id,
-                self.task_instance.map_index,
+    if not AIRFLOW_V_3_0_PLUS:
+
+        @provide_session
+        def get_task_instance(self, session: Session) -> TaskInstance:
+            query = session.query(TaskInstance).filter(
+                TaskInstance.dag_id == self.task_instance.dag_id,
+                TaskInstance.task_id == self.task_instance.task_id,
+                TaskInstance.run_id == self.task_instance.run_id,
+                TaskInstance.map_index == self.task_instance.map_index,
             )
-        return task_instance
+            task_instance = query.one_or_none()
+            if task_instance is None:
+                raise AirflowException(
+                    "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
+                    self.task_instance.dag_id,
+                    self.task_instance.task_id,
+                    self.task_instance.run_id,
+                    self.task_instance.map_index,
+                )
+            return task_instance
 
     async def get_task_state(self):
         from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance

airflow/providers/google/cloud/triggers/cloud_composer.py

@@ -145,10 +145,23 @@ class CloudComposerAirflowCLICommandTrigger(BaseTrigger):
             )
             return
 
+        exit_code = result.get("exit_info", {}).get("exit_code")
+
+        if exit_code == 0:
+            yield TriggerEvent(
+                {
+                    "status": "success",
+                    "result": result,
+                }
+            )
+            return
+
+        error_output = "".join(line["content"] for line in result.get("error", []))
+        message = f"Airflow CLI command failed with exit code {exit_code}.\nError output:\n{error_output}"
         yield TriggerEvent(
             {
-                "status": "success",
-                "result": result,
+                "status": "error",
+                "message": message,
             }
         )
         return
@@ -166,6 +179,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         start_date: datetime,
         end_date: datetime,
         allowed_states: list[str],
+        composer_dag_run_id: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         poll_interval: int = 10,
@@ -179,6 +193,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         self.start_date = start_date
         self.end_date = end_date
         self.allowed_states = allowed_states
+        self.composer_dag_run_id = composer_dag_run_id
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
         self.poll_interval = poll_interval
@@ -200,6 +215,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
                 "start_date": self.start_date,
                 "end_date": self.end_date,
                 "allowed_states": self.allowed_states,
+                "composer_dag_run_id": self.composer_dag_run_id,
                 "gcp_conn_id": self.gcp_conn_id,
                 "impersonation_chain": self.impersonation_chain,
                 "poll_interval": self.poll_interval,
@@ -248,20 +264,42 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
                 return False
         return True
 
+    def _check_composer_dag_run_id_states(self, dag_runs: list[dict]) -> bool:
+        for dag_run in dag_runs:
+            if dag_run["run_id"] == self.composer_dag_run_id and dag_run["state"] in self.allowed_states:
+                return True
+        return False
+
     async def run(self):
         try:
             while True:
                 if datetime.now(self.end_date.tzinfo).timestamp() > self.end_date.timestamp():
                     dag_runs = await self._pull_dag_runs()
 
-                    self.log.info("Sensor waits for allowed states: %s", self.allowed_states)
-                    if self._check_dag_runs_states(
-                        dag_runs=dag_runs,
-                        start_date=self.start_date,
-                        end_date=self.end_date,
-                    ):
-                        yield TriggerEvent({"status": "success"})
-                        return
+                    if len(dag_runs) == 0:
+                        self.log.info("Dag runs are empty. Sensor waits for dag runs...")
+                        self.log.info("Sleeping for %s seconds.", self.poll_interval)
+                        await asyncio.sleep(self.poll_interval)
+                        continue
+
+                    if self.composer_dag_run_id:
+                        self.log.info(
+                            "Sensor waits for allowed states %s for specified RunID: %s",
+                            self.allowed_states,
+                            self.composer_dag_run_id,
+                        )
+                        if self._check_composer_dag_run_id_states(dag_runs=dag_runs):
+                            yield TriggerEvent({"status": "success"})
+                            return
+                    else:
+                        self.log.info("Sensor waits for allowed states: %s", self.allowed_states)
+                        if self._check_dag_runs_states(
+                            dag_runs=dag_runs,
+                            start_date=self.start_date,
+                            end_date=self.end_date,
+                        ):
+                            yield TriggerEvent({"status": "success"})
+                            return
                 self.log.info("Sleeping for %s seconds.", self.poll_interval)
                 await asyncio.sleep(self.poll_interval)
         except AirflowException as ex:

airflow/providers/google/cloud/triggers/dataplex.py

@@ -103,7 +103,13 @@ class DataplexDataQualityJobTrigger(BaseTrigger):
                 self.polling_interval_seconds,
             )
             await asyncio.sleep(self.polling_interval_seconds)
-        yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": self._convert_to_dict(job)})
+        yield TriggerEvent(
+            {
+                "job_id": self.job_id,
+                "job_state": DataScanJob.State(state).name,
+                "job": self._convert_to_dict(job),
+            }
+        )
 
     def _convert_to_dict(self, job: DataScanJob) -> dict:
         """Return a representation of a DataScanJob instance as a dict."""
@@ -185,7 +191,13 @@ class DataplexDataProfileJobTrigger(BaseTrigger):
                 self.polling_interval_seconds,
             )
             await asyncio.sleep(self.polling_interval_seconds)
-        yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": self._convert_to_dict(job)})
+        yield TriggerEvent(
+            {
+                "job_id": self.job_id,
+                "job_state": DataScanJob.State(state).name,
+                "job": self._convert_to_dict(job),
+            }
+        )
 
     def _convert_to_dict(self, job: DataScanJob) -> dict:
         """Return a representation of a DataScanJob instance as a dict."""