apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -22,21 +22,23 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import Conflict
+from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
 from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import BaseOperator
 from airflow.utils.helpers import merge_dicts
-from google.api_core.exceptions import Conflict
-from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class BigQueryToGCSOperator(BaseOperator):
     """
@@ -213,8 +215,9 @@ class BigQueryToGCSOperator(BaseOperator):
             job_id=self.job_id,
             dag_id=self.dag_id,
             task_id=self.task_id,
-            logical_date=context["logical_date"],
+            logical_date=None,
             configuration=configuration,
+            run_after=hook.get_run_after_or_logical_date(context),
             force_rerun=self.force_rerun,
         )
 
@@ -239,21 +242,19 @@ class BigQueryToGCSOperator(BaseOperator):
                     f"want to force rerun it consider setting `force_rerun=True`."
                     f"Or, if you want to reattach in this scenario add {job.state} to `reattach_states`"
                 )
-            else:
-                # Job already reached state DONE
-                if job.state == "DONE":
-                    raise AirflowException("Job is already in state DONE. Can not reattach to this job.")
+            # Job already reached state DONE
+            if job.state == "DONE":
+                raise AirflowException("Job is already in state DONE. Can not reattach to this job.")
 
-                # We are reattaching to a job
-                self.log.info("Reattaching to existing Job in state %s", job.state)
-                self._handle_job_error(job)
+            # We are reattaching to a job
+            self.log.info("Reattaching to existing Job in state %s", job.state)
+            self._handle_job_error(job)
 
         self.job_id = job.job_id
         conf = job.to_api_repr()["configuration"]["extract"]["sourceTable"]
         dataset_id, project_id, table_id = conf["datasetId"], conf["projectId"], conf["tableId"]
         BigQueryTableLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_id,
             project_id=project_id,
             table_id=table_id,

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

@@ -21,6 +21,7 @@ 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
@@ -29,7 +30,7 @@ from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToS
 from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
@@ -94,14 +95,17 @@ 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(
             context=context,
-            task_instance=self,
             dataset_id=dataset_id,
             project_id=project_id,
             table_id=table_id,

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

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import warnings
 from collections.abc import Sequence
+from functools import cached_property
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
@@ -76,5 +77,15 @@ class BigQueryToMySqlOperator(BigQueryToSqlBaseOperator):
         )
         self.mysql_conn_id = mysql_conn_id
 
-    def get_sql_hook(self) -> MySqlHook:
+    @cached_property
+    def mysql_hook(self) -> MySqlHook:
         return MySqlHook(schema=self.database, mysql_conn_id=self.mysql_conn_id)
+
+    def get_sql_hook(self) -> MySqlHook:
+        return self.mysql_hook
+
+    def execute(self, context):
+        # Set source_project_dataset_table here, after hooks are initialized and project_id is available
+        project_id = self.bigquery_hook.project_id
+        self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
+        return super().execute(context)

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

@@ -19,15 +19,19 @@
 
 from __future__ import annotations
 
+from functools import cached_property
 from typing import TYPE_CHECKING
 
+from psycopg2.extensions import register_adapter
+from psycopg2.extras import Json
+
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
 from airflow.providers.google.cloud.utils.bigquery_get_data import bigquery_get_data
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
@@ -75,26 +79,36 @@ class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
         self.postgres_conn_id = postgres_conn_id
         self.replace_index = replace_index
 
-    def get_sql_hook(self) -> PostgresHook:
+    @cached_property
+    def postgres_hook(self) -> PostgresHook:
+        register_adapter(list, Json)
+        register_adapter(dict, Json)
         return PostgresHook(database=self.database, postgres_conn_id=self.postgres_conn_id)
 
+    def get_sql_hook(self) -> PostgresHook:
+        return self.postgres_hook
+
     def execute(self, context: Context) -> None:
-        big_query_hook = BigQueryHook(
-            gcp_conn_id=self.gcp_conn_id,
-            location=self.location,
-            impersonation_chain=self.impersonation_chain,
-        )
+        if not self.bigquery_hook:
+            self.bigquery_hook = BigQueryHook(
+                gcp_conn_id=self.gcp_conn_id,
+                location=self.location,
+                impersonation_chain=self.impersonation_chain,
+            )
+        # Set source_project_dataset_table here, after hooks are initialized and project_id is available
+        project_id = self.bigquery_hook.project_id
+        self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
+
         self.persist_links(context)
-        sql_hook: PostgresHook = 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,
         ):
-            sql_hook.insert_rows(
+            self.postgres_hook.insert_rows(
                 table=self.target_table_name,
                 rows=rows,
                 target_fields=self.selected_fields,

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

@@ -21,15 +21,17 @@ from __future__ import annotations
 
 import abc
 from collections.abc import Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.utils.bigquery_get_data import bigquery_get_data
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
+    from airflow.providers.common.compat.sdk import Context
     from airflow.providers.common.sql.hooks.sql import DbApiHook
-    from airflow.utils.context import Context
+    from airflow.providers.openlineage.extractors import OperatorLineage
 
 
 class BigQueryToSqlBaseOperator(BaseOperator):
@@ -113,19 +115,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,
         ):
@@ -136,3 +141,97 @@ class BigQueryToSqlBaseOperator(BaseOperator):
                 replace=self.replace,
                 commit_every=self.batch_size,
             )
+
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """
+        Build a generic OpenLineage facet for BigQuery -> SQL transfers.
+
+        This consolidates nearly identical implementations from child
+        operators. Children still provide a concrete SQL hook via
+        ``get_sql_hook()`` and may override behavior if needed.
+        """
+        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:
+            if not getattr(self, "source_project_dataset_table", None):
+                project_id = self.bigquery_hook.project_id
+                self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
+
+            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",
+                getattr(self, "source_project_dataset_table", None),
+                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),
+        )
+
+        sql_hook = self.get_sql_hook()
+        db_info = sql_hook.get_openlineage_database_info(sql_hook.get_conn())
+        if db_info is None:
+            self.log.debug("OpenLineage: could not get database info from SQL hook %s", type(sql_hook))
+            return OperatorLineage()
+        namespace = f"{db_info.scheme}://{db_info.authority}"
+
+        schema_name = None
+        if hasattr(sql_hook, "get_openlineage_default_schema"):
+            try:
+                schema_name = sql_hook.get_openlineage_default_schema()
+            except Exception:
+                schema_name = None
+
+        if self.target_table_name and "." in self.target_table_name:
+            schema_part, table_part = self.target_table_name.split(".", 1)
+        else:
+            schema_part = schema_name or ""
+            table_part = self.target_table_name or ""
+
+        if db_info and db_info.scheme == "mysql":
+            output_name = f"{self.database}.{table_part}" if self.database else f"{table_part}"
+        else:
+            if self.database:
+                if schema_part:
+                    output_name = f"{self.database}.{schema_part}.{table_part}"
+                else:
+                    output_name = f"{self.database}.{table_part}"
+            else:
+                if schema_part:
+                    output_name = f"{schema_part}.{table_part}"
+                else:
+                    output_name = f"{table_part}"
+
+        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/calendar_to_gcs.py

@@ -21,9 +21,9 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Any
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.suite.hooks.calendar import GoogleCalendarHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from datetime import datetime

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

@@ -31,12 +31,12 @@ from uuid import UUID
 from cassandra.util import Date, OrderedMapSerializedKey, SortedSet, Time
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 NotSetType = NewType("NotSetType", object)
 NOT_SET = NotSetType(object())
@@ -261,29 +261,27 @@ class CassandraToGCSOperator(BaseOperator):
         """Convert value to BQ type."""
         if not value or isinstance(value, (str, int, float, bool, dict)):
             return value
-        elif isinstance(value, bytes):
+        if isinstance(value, bytes):
             return b64encode(value).decode("ascii")
-        elif isinstance(value, UUID):
+        if isinstance(value, UUID):
             if self.encode_uuid:
                 return b64encode(value.bytes).decode("ascii")
-            else:
-                return str(value)
-        elif isinstance(value, (datetime, Date)):
             return str(value)
-        elif isinstance(value, Decimal):
+        if isinstance(value, (datetime, Date)):
+            return str(value)
+        if isinstance(value, Decimal):
             return float(value)
-        elif isinstance(value, Time):
+        if isinstance(value, Time):
             return str(value).split(".")[0]
-        elif isinstance(value, (list, SortedSet)):
+        if isinstance(value, (list, SortedSet)):
             return self.convert_array_types(value)
-        elif hasattr(value, "_fields"):
+        if hasattr(value, "_fields"):
             return self.convert_user_type(value)
-        elif isinstance(value, tuple):
+        if isinstance(value, tuple):
             return self.convert_tuple_type(value)
-        elif isinstance(value, OrderedMapSerializedKey):
+        if isinstance(value, OrderedMapSerializedKey):
             return self.convert_map_type(value)
-        else:
-            raise AirflowException(f"Unexpected value: {value}")
+        raise AirflowException(f"Unexpected value: {value}")
 
     def convert_array_types(self, value: list[Any] | SortedSet) -> list[Any]:
         """Map convert_value over array."""
@@ -376,19 +374,17 @@ class CassandraToGCSOperator(BaseOperator):
         """Convert type to equivalent BQ type."""
         if cls.is_simple_type(type_):
             return CassandraToGCSOperator.CQL_TYPE_MAP[type_.cassname]
-        elif cls.is_record_type(type_):
+        if cls.is_record_type(type_):
             return "RECORD"
-        elif cls.is_array_type(type_):
+        if cls.is_array_type(type_):
             return cls.get_bq_type(type_.subtypes[0])
-        else:
-            raise AirflowException("Not a supported type_: " + type_.cassname)
+        raise AirflowException("Not a supported type_: " + type_.cassname)
 
     @classmethod
     def get_bq_mode(cls, type_: Any) -> str:
         """Convert type to equivalent BQ mode."""
         if cls.is_array_type(type_) or type_.cassname == "MapType":
             return "REPEATED"
-        elif cls.is_record_type(type_) or cls.is_simple_type(type_):
+        if cls.is_record_type(type_) or cls.is_simple_type(type_):
             return "NULLABLE"
-        else:
-            raise AirflowException("Not a supported type_: " + type_.cassname)
+        raise AirflowException("Not a supported type_: " + type_.cassname)

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

@@ -26,14 +26,14 @@ from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.facebook.ads.hooks.ads import FacebookAdsReportingHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from facebook_business.adobjects.adsinsights import AdsInsights
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class FlushAction(Enum):
@@ -157,8 +157,7 @@ class FacebookAdsReportToGcsOperator(BaseOperator):
     def _generate_rows_with_action(self, type_check: bool):
         if type_check and self.upload_as_account:
             return {FlushAction.EXPORT_EVERY_ACCOUNT: []}
-        else:
-            return {FlushAction.EXPORT_ONCE: []}
+        return {FlushAction.EXPORT_ONCE: []}
 
     def _prepare_rows_for_upload(
         self,
@@ -209,7 +208,7 @@ class FacebookAdsReportToGcsOperator(BaseOperator):
 
     def _flush_rows(self, converted_rows: list[Any] | None, object_name: str):
         if converted_rows:
-            headers = converted_rows[0].keys()
+            headers = self.fields
             with tempfile.NamedTemporaryFile("w", suffix=".csv") as csvfile:
                 writer = csv.DictWriter(csvfile, fieldnames=headers)
                 writer.writeheader()