apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.3.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (257) hide show
  1. airflow/providers/google/3rd-party-licenses/NOTICE +2 -12
  2. airflow/providers/google/__init__.py +3 -3
  3. airflow/providers/google/ads/hooks/ads.py +39 -6
  4. airflow/providers/google/ads/operators/ads.py +2 -2
  5. airflow/providers/google/ads/transfers/ads_to_gcs.py +2 -2
  6. airflow/providers/google/assets/gcs.py +1 -11
  7. airflow/providers/google/cloud/bundles/__init__.py +16 -0
  8. airflow/providers/google/cloud/bundles/gcs.py +161 -0
  9. airflow/providers/google/cloud/hooks/alloy_db.py +1 -1
  10. airflow/providers/google/cloud/hooks/bigquery.py +176 -293
  11. airflow/providers/google/cloud/hooks/cloud_batch.py +1 -1
  12. airflow/providers/google/cloud/hooks/cloud_build.py +1 -1
  13. airflow/providers/google/cloud/hooks/cloud_composer.py +288 -15
  14. airflow/providers/google/cloud/hooks/cloud_logging.py +109 -0
  15. airflow/providers/google/cloud/hooks/cloud_memorystore.py +1 -1
  16. airflow/providers/google/cloud/hooks/cloud_run.py +18 -10
  17. airflow/providers/google/cloud/hooks/cloud_sql.py +102 -23
  18. airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +29 -7
  19. airflow/providers/google/cloud/hooks/compute.py +1 -1
  20. airflow/providers/google/cloud/hooks/compute_ssh.py +6 -2
  21. airflow/providers/google/cloud/hooks/datacatalog.py +10 -1
  22. airflow/providers/google/cloud/hooks/dataflow.py +72 -95
  23. airflow/providers/google/cloud/hooks/dataform.py +1 -1
  24. airflow/providers/google/cloud/hooks/datafusion.py +21 -19
  25. airflow/providers/google/cloud/hooks/dataplex.py +2 -2
  26. airflow/providers/google/cloud/hooks/dataprep.py +1 -1
  27. airflow/providers/google/cloud/hooks/dataproc.py +73 -72
  28. airflow/providers/google/cloud/hooks/dataproc_metastore.py +1 -1
  29. airflow/providers/google/cloud/hooks/dlp.py +1 -1
  30. airflow/providers/google/cloud/hooks/functions.py +1 -1
  31. airflow/providers/google/cloud/hooks/gcs.py +112 -15
  32. airflow/providers/google/cloud/hooks/gdm.py +1 -1
  33. airflow/providers/google/cloud/hooks/gen_ai.py +196 -0
  34. airflow/providers/google/cloud/hooks/kubernetes_engine.py +3 -3
  35. airflow/providers/google/cloud/hooks/looker.py +6 -2
  36. airflow/providers/google/cloud/hooks/managed_kafka.py +1 -1
  37. airflow/providers/google/cloud/hooks/mlengine.py +4 -3
  38. airflow/providers/google/cloud/hooks/pubsub.py +3 -0
  39. airflow/providers/google/cloud/hooks/secret_manager.py +102 -10
  40. airflow/providers/google/cloud/hooks/spanner.py +74 -9
  41. airflow/providers/google/cloud/hooks/stackdriver.py +11 -9
  42. airflow/providers/google/cloud/hooks/tasks.py +1 -1
  43. airflow/providers/google/cloud/hooks/translate.py +2 -2
  44. airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +2 -210
  45. airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +3 -3
  46. airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +28 -2
  47. airflow/providers/google/cloud/hooks/vertex_ai/experiment_service.py +202 -0
  48. airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +308 -8
  49. airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +79 -75
  50. airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +1 -1
  51. airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +1 -1
  52. airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +1 -1
  53. airflow/providers/google/cloud/hooks/vertex_ai/ray.py +223 -0
  54. airflow/providers/google/cloud/hooks/vision.py +3 -3
  55. airflow/providers/google/cloud/hooks/workflows.py +1 -1
  56. airflow/providers/google/cloud/links/alloy_db.py +0 -46
  57. airflow/providers/google/cloud/links/base.py +77 -13
  58. airflow/providers/google/cloud/links/bigquery.py +0 -47
  59. airflow/providers/google/cloud/links/bigquery_dts.py +0 -20
  60. airflow/providers/google/cloud/links/bigtable.py +0 -48
  61. airflow/providers/google/cloud/links/cloud_build.py +0 -73
  62. airflow/providers/google/cloud/links/cloud_functions.py +0 -33
  63. airflow/providers/google/cloud/links/cloud_memorystore.py +0 -58
  64. airflow/providers/google/cloud/links/{life_sciences.py → cloud_run.py} +5 -27
  65. airflow/providers/google/cloud/links/cloud_sql.py +0 -33
  66. airflow/providers/google/cloud/links/cloud_storage_transfer.py +17 -44
  67. airflow/providers/google/cloud/links/cloud_tasks.py +7 -26
  68. airflow/providers/google/cloud/links/compute.py +0 -58
  69. airflow/providers/google/cloud/links/data_loss_prevention.py +0 -169
  70. airflow/providers/google/cloud/links/datacatalog.py +23 -54
  71. airflow/providers/google/cloud/links/dataflow.py +0 -34
  72. airflow/providers/google/cloud/links/dataform.py +0 -64
  73. airflow/providers/google/cloud/links/datafusion.py +1 -96
  74. airflow/providers/google/cloud/links/dataplex.py +0 -154
  75. airflow/providers/google/cloud/links/dataprep.py +0 -24
  76. airflow/providers/google/cloud/links/dataproc.py +11 -95
  77. airflow/providers/google/cloud/links/datastore.py +0 -31
  78. airflow/providers/google/cloud/links/kubernetes_engine.py +9 -60
  79. airflow/providers/google/cloud/links/managed_kafka.py +0 -70
  80. airflow/providers/google/cloud/links/mlengine.py +0 -70
  81. airflow/providers/google/cloud/links/pubsub.py +0 -32
  82. airflow/providers/google/cloud/links/spanner.py +0 -33
  83. airflow/providers/google/cloud/links/stackdriver.py +0 -30
  84. airflow/providers/google/cloud/links/translate.py +17 -187
  85. airflow/providers/google/cloud/links/vertex_ai.py +28 -195
  86. airflow/providers/google/cloud/links/workflows.py +0 -52
  87. airflow/providers/google/cloud/log/gcs_task_handler.py +58 -22
  88. airflow/providers/google/cloud/log/stackdriver_task_handler.py +9 -6
  89. airflow/providers/google/cloud/openlineage/CloudStorageTransferJobFacet.json +68 -0
  90. airflow/providers/google/cloud/openlineage/CloudStorageTransferRunFacet.json +60 -0
  91. airflow/providers/google/cloud/openlineage/DataFusionRunFacet.json +32 -0
  92. airflow/providers/google/cloud/openlineage/facets.py +102 -1
  93. airflow/providers/google/cloud/openlineage/mixins.py +10 -8
  94. airflow/providers/google/cloud/openlineage/utils.py +15 -1
  95. airflow/providers/google/cloud/operators/alloy_db.py +71 -56
  96. airflow/providers/google/cloud/operators/bigquery.py +73 -636
  97. airflow/providers/google/cloud/operators/bigquery_dts.py +4 -6
  98. airflow/providers/google/cloud/operators/bigtable.py +37 -8
  99. airflow/providers/google/cloud/operators/cloud_base.py +21 -1
  100. airflow/providers/google/cloud/operators/cloud_batch.py +3 -3
  101. airflow/providers/google/cloud/operators/cloud_build.py +76 -33
  102. airflow/providers/google/cloud/operators/cloud_composer.py +129 -41
  103. airflow/providers/google/cloud/operators/cloud_logging_sink.py +341 -0
  104. airflow/providers/google/cloud/operators/cloud_memorystore.py +69 -43
  105. airflow/providers/google/cloud/operators/cloud_run.py +24 -6
  106. airflow/providers/google/cloud/operators/cloud_sql.py +8 -17
  107. airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +93 -12
  108. airflow/providers/google/cloud/operators/compute.py +9 -41
  109. airflow/providers/google/cloud/operators/datacatalog.py +157 -21
  110. airflow/providers/google/cloud/operators/dataflow.py +40 -16
  111. airflow/providers/google/cloud/operators/dataform.py +15 -5
  112. airflow/providers/google/cloud/operators/datafusion.py +42 -21
  113. airflow/providers/google/cloud/operators/dataplex.py +194 -110
  114. airflow/providers/google/cloud/operators/dataprep.py +1 -5
  115. airflow/providers/google/cloud/operators/dataproc.py +80 -36
  116. airflow/providers/google/cloud/operators/dataproc_metastore.py +97 -89
  117. airflow/providers/google/cloud/operators/datastore.py +23 -7
  118. airflow/providers/google/cloud/operators/dlp.py +6 -29
  119. airflow/providers/google/cloud/operators/functions.py +17 -8
  120. airflow/providers/google/cloud/operators/gcs.py +12 -9
  121. airflow/providers/google/cloud/operators/gen_ai.py +389 -0
  122. airflow/providers/google/cloud/operators/kubernetes_engine.py +62 -100
  123. airflow/providers/google/cloud/operators/looker.py +2 -2
  124. airflow/providers/google/cloud/operators/managed_kafka.py +108 -53
  125. airflow/providers/google/cloud/operators/natural_language.py +1 -1
  126. airflow/providers/google/cloud/operators/pubsub.py +68 -15
  127. airflow/providers/google/cloud/operators/spanner.py +26 -13
  128. airflow/providers/google/cloud/operators/speech_to_text.py +2 -3
  129. airflow/providers/google/cloud/operators/stackdriver.py +1 -9
  130. airflow/providers/google/cloud/operators/tasks.py +1 -12
  131. airflow/providers/google/cloud/operators/text_to_speech.py +2 -3
  132. airflow/providers/google/cloud/operators/translate.py +41 -17
  133. airflow/providers/google/cloud/operators/translate_speech.py +2 -3
  134. airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +39 -19
  135. airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +30 -10
  136. airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +55 -27
  137. airflow/providers/google/cloud/operators/vertex_ai/dataset.py +70 -8
  138. airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +43 -9
  139. airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py +435 -0
  140. airflow/providers/google/cloud/operators/vertex_ai/feature_store.py +532 -1
  141. airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +135 -115
  142. airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +12 -10
  143. airflow/providers/google/cloud/operators/vertex_ai/model_service.py +57 -11
  144. airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +31 -8
  145. airflow/providers/google/cloud/operators/vertex_ai/ray.py +393 -0
  146. airflow/providers/google/cloud/operators/video_intelligence.py +1 -1
  147. airflow/providers/google/cloud/operators/vision.py +2 -2
  148. airflow/providers/google/cloud/operators/workflows.py +18 -15
  149. airflow/providers/google/cloud/secrets/secret_manager.py +3 -2
  150. airflow/providers/google/cloud/sensors/bigquery.py +3 -3
  151. airflow/providers/google/cloud/sensors/bigquery_dts.py +2 -3
  152. airflow/providers/google/cloud/sensors/bigtable.py +11 -4
  153. airflow/providers/google/cloud/sensors/cloud_composer.py +533 -30
  154. airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +2 -3
  155. airflow/providers/google/cloud/sensors/dataflow.py +26 -10
  156. airflow/providers/google/cloud/sensors/dataform.py +2 -3
  157. airflow/providers/google/cloud/sensors/datafusion.py +4 -5
  158. airflow/providers/google/cloud/sensors/dataplex.py +2 -3
  159. airflow/providers/google/cloud/sensors/dataprep.py +2 -2
  160. airflow/providers/google/cloud/sensors/dataproc.py +2 -3
  161. airflow/providers/google/cloud/sensors/dataproc_metastore.py +2 -3
  162. airflow/providers/google/cloud/sensors/gcs.py +4 -5
  163. airflow/providers/google/cloud/sensors/looker.py +2 -3
  164. airflow/providers/google/cloud/sensors/pubsub.py +4 -5
  165. airflow/providers/google/cloud/sensors/tasks.py +2 -2
  166. airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +2 -3
  167. airflow/providers/google/cloud/sensors/workflows.py +2 -3
  168. airflow/providers/google/cloud/transfers/adls_to_gcs.py +1 -1
  169. airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +2 -2
  170. airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +4 -3
  171. airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py +11 -8
  172. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +10 -5
  173. airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +7 -3
  174. airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +12 -1
  175. airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +24 -10
  176. airflow/providers/google/cloud/transfers/bigquery_to_sql.py +104 -5
  177. airflow/providers/google/cloud/transfers/calendar_to_gcs.py +1 -1
  178. airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +3 -3
  179. airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +4 -4
  180. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +21 -13
  181. airflow/providers/google/cloud/transfers/gcs_to_gcs.py +4 -3
  182. airflow/providers/google/cloud/transfers/gcs_to_local.py +6 -4
  183. airflow/providers/google/cloud/transfers/gcs_to_sftp.py +11 -5
  184. airflow/providers/google/cloud/transfers/gdrive_to_gcs.py +6 -2
  185. airflow/providers/google/cloud/transfers/gdrive_to_local.py +2 -2
  186. airflow/providers/google/cloud/transfers/http_to_gcs.py +193 -0
  187. airflow/providers/google/cloud/transfers/local_to_gcs.py +2 -2
  188. airflow/providers/google/cloud/transfers/mssql_to_gcs.py +1 -1
  189. airflow/providers/google/cloud/transfers/oracle_to_gcs.py +36 -11
  190. airflow/providers/google/cloud/transfers/postgres_to_gcs.py +42 -9
  191. airflow/providers/google/cloud/transfers/s3_to_gcs.py +13 -7
  192. airflow/providers/google/cloud/transfers/salesforce_to_gcs.py +2 -2
  193. airflow/providers/google/cloud/transfers/sftp_to_gcs.py +14 -5
  194. airflow/providers/google/cloud/transfers/sheets_to_gcs.py +3 -3
  195. airflow/providers/google/cloud/transfers/sql_to_gcs.py +10 -10
  196. airflow/providers/google/cloud/triggers/bigquery.py +76 -35
  197. airflow/providers/google/cloud/triggers/cloud_build.py +1 -1
  198. airflow/providers/google/cloud/triggers/cloud_composer.py +303 -47
  199. airflow/providers/google/cloud/triggers/cloud_run.py +3 -3
  200. airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +92 -2
  201. airflow/providers/google/cloud/triggers/dataflow.py +122 -0
  202. airflow/providers/google/cloud/triggers/datafusion.py +1 -1
  203. airflow/providers/google/cloud/triggers/dataplex.py +14 -2
  204. airflow/providers/google/cloud/triggers/dataproc.py +123 -53
  205. airflow/providers/google/cloud/triggers/kubernetes_engine.py +47 -28
  206. airflow/providers/google/cloud/triggers/mlengine.py +1 -1
  207. airflow/providers/google/cloud/triggers/pubsub.py +15 -19
  208. airflow/providers/google/cloud/triggers/vertex_ai.py +1 -1
  209. airflow/providers/google/cloud/utils/bigquery_get_data.py +1 -1
  210. airflow/providers/google/cloud/utils/credentials_provider.py +2 -2
  211. airflow/providers/google/cloud/utils/field_sanitizer.py +1 -1
  212. airflow/providers/google/cloud/utils/field_validator.py +2 -3
  213. airflow/providers/google/common/auth_backend/google_openid.py +4 -4
  214. airflow/providers/google/common/deprecated.py +2 -1
  215. airflow/providers/google/common/hooks/base_google.py +27 -9
  216. airflow/providers/google/common/hooks/operation_helpers.py +1 -1
  217. airflow/providers/google/common/links/storage.py +0 -22
  218. airflow/providers/google/common/utils/get_secret.py +31 -0
  219. airflow/providers/google/common/utils/id_token_credentials.py +3 -4
  220. airflow/providers/google/firebase/hooks/firestore.py +1 -1
  221. airflow/providers/google/firebase/operators/firestore.py +3 -3
  222. airflow/providers/google/get_provider_info.py +56 -52
  223. airflow/providers/google/go_module_utils.py +35 -3
  224. airflow/providers/google/leveldb/hooks/leveldb.py +27 -2
  225. airflow/providers/google/leveldb/operators/leveldb.py +2 -2
  226. airflow/providers/google/marketing_platform/hooks/campaign_manager.py +1 -1
  227. airflow/providers/google/marketing_platform/hooks/display_video.py +3 -109
  228. airflow/providers/google/marketing_platform/hooks/search_ads.py +1 -1
  229. airflow/providers/google/marketing_platform/links/analytics_admin.py +5 -14
  230. airflow/providers/google/marketing_platform/operators/analytics_admin.py +2 -3
  231. airflow/providers/google/marketing_platform/operators/campaign_manager.py +6 -6
  232. airflow/providers/google/marketing_platform/operators/display_video.py +28 -489
  233. airflow/providers/google/marketing_platform/operators/search_ads.py +2 -2
  234. airflow/providers/google/marketing_platform/sensors/campaign_manager.py +2 -2
  235. airflow/providers/google/marketing_platform/sensors/display_video.py +3 -64
  236. airflow/providers/google/suite/hooks/calendar.py +2 -2
  237. airflow/providers/google/suite/hooks/sheets.py +16 -2
  238. airflow/providers/google/suite/operators/sheets.py +8 -3
  239. airflow/providers/google/suite/sensors/drive.py +2 -2
  240. airflow/providers/google/suite/transfers/gcs_to_gdrive.py +3 -3
  241. airflow/providers/google/suite/transfers/gcs_to_sheets.py +1 -1
  242. airflow/providers/google/suite/transfers/local_to_drive.py +3 -3
  243. airflow/providers/google/suite/transfers/sql_to_sheets.py +5 -4
  244. airflow/providers/google/version_compat.py +15 -1
  245. {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/METADATA +90 -46
  246. apache_airflow_providers_google-19.3.0.dist-info/RECORD +331 -0
  247. apache_airflow_providers_google-19.3.0.dist-info/licenses/NOTICE +5 -0
  248. airflow/providers/google/cloud/hooks/automl.py +0 -673
  249. airflow/providers/google/cloud/hooks/life_sciences.py +0 -159
  250. airflow/providers/google/cloud/links/automl.py +0 -193
  251. airflow/providers/google/cloud/operators/automl.py +0 -1362
  252. airflow/providers/google/cloud/operators/life_sciences.py +0 -119
  253. airflow/providers/google/cloud/operators/mlengine.py +0 -112
  254. apache_airflow_providers_google-15.1.0rc1.dist-info/RECORD +0 -321
  255. {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/WHEEL +0 -0
  256. {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/entry_points.txt +0 -0
  257. {airflow/providers/google → apache_airflow_providers_google-19.3.0.dist-info/licenses}/LICENSE +0 -0
@@ -19,12 +19,12 @@
19
19
 
20
20
  from __future__ import annotations
21
21
 
22
- from collections.abc import Sequence
22
+ from collections.abc import Callable, Sequence
23
23
  from functools import cached_property
24
- from typing import TYPE_CHECKING, Any, Callable
24
+ from typing import TYPE_CHECKING, Any
25
25
 
26
26
  from airflow.configuration import conf
27
- from airflow.exceptions import AirflowException
27
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, PokeReturnValue
28
28
  from airflow.providers.google.cloud.hooks.dataflow import (
29
29
  DEFAULT_DATAFLOW_LOCATION,
30
30
  DataflowHook,
@@ -37,10 +37,9 @@ from airflow.providers.google.cloud.triggers.dataflow import (
37
37
  DataflowJobStatusTrigger,
38
38
  )
39
39
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
40
- from airflow.sensors.base import BaseSensorOperator
41
40
 
42
41
  if TYPE_CHECKING:
43
- from airflow.utils.context import Context
42
+ from airflow.providers.common.compat.sdk import Context
44
43
 
45
44
 
46
45
  class DataflowJobStatusSensor(BaseSensorOperator):
@@ -342,7 +341,7 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
342
341
  self.deferrable = deferrable
343
342
  self.poll_interval = poll_interval
344
343
 
345
- def poke(self, context: Context) -> bool:
344
+ def poke(self, context: Context) -> PokeReturnValue | bool:
346
345
  if self.fail_on_terminal_state:
347
346
  job = self.hook.get_job(
348
347
  job_id=self.job_id,
@@ -359,8 +358,17 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
359
358
  project_id=self.project_id,
360
359
  location=self.location,
361
360
  )
361
+ result = result if self.callback is None else self.callback(result)
362
+
363
+ if isinstance(result, PokeReturnValue):
364
+ return result
362
365
 
363
- return result if self.callback is None else self.callback(result)
366
+ if bool(result):
367
+ return PokeReturnValue(
368
+ is_done=True,
369
+ xcom_value=result,
370
+ )
371
+ return False
364
372
 
365
373
  def execute(self, context: Context) -> Any:
366
374
  """Airflow runs this method on the worker and defers using the trigger."""
@@ -464,7 +472,7 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
464
472
  self.deferrable = deferrable
465
473
  self.poll_interval = poll_interval
466
474
 
467
- def poke(self, context: Context) -> bool:
475
+ def poke(self, context: Context) -> PokeReturnValue | bool:
468
476
  if self.fail_on_terminal_state:
469
477
  job = self.hook.get_job(
470
478
  job_id=self.job_id,
@@ -481,8 +489,16 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
481
489
  project_id=self.project_id,
482
490
  location=self.location,
483
491
  )
484
-
485
- return result if self.callback is None else self.callback(result)
492
+ result = result if self.callback is None else self.callback(result)
493
+ if isinstance(result, PokeReturnValue):
494
+ return result
495
+
496
+ if bool(result):
497
+ return PokeReturnValue(
498
+ is_done=True,
499
+ xcom_value=result,
500
+ )
501
+ return False
486
502
 
487
503
  def execute(self, context: Context) -> Any:
488
504
  """Airflow runs this method on the worker and defers using the trigger."""
@@ -22,12 +22,11 @@ from __future__ import annotations
22
22
  from collections.abc import Iterable, Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from airflow.exceptions import AirflowException
25
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
26
26
  from airflow.providers.google.cloud.hooks.dataform import DataformHook
27
- from airflow.sensors.base import BaseSensorOperator
28
27
 
29
28
  if TYPE_CHECKING:
30
- from airflow.utils.context import Context
29
+ from airflow.providers.common.compat.sdk import Context
31
30
 
32
31
 
33
32
  class DataformWorkflowInvocationStateSensor(BaseSensorOperator):
@@ -22,13 +22,12 @@ from __future__ import annotations
22
22
  from collections.abc import Iterable, Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from airflow.exceptions import AirflowException, AirflowNotFoundException
25
+ from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException, BaseSensorOperator
26
26
  from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
27
27
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
28
- from airflow.sensors.base import BaseSensorOperator
29
28
 
30
29
  if TYPE_CHECKING:
31
- from airflow.utils.context import Context
30
+ from airflow.providers.common.compat.sdk import Context
32
31
 
33
32
 
34
33
  class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
@@ -110,7 +109,7 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
110
109
  pipeline_id=self.pipeline_id,
111
110
  namespace=self.namespace,
112
111
  )
113
- pipeline_status = pipeline_workflow["status"]
112
+ pipeline_status = pipeline_workflow.get("status")
114
113
  except AirflowNotFoundException:
115
114
  message = "Specified Pipeline ID was not found."
116
115
  raise AirflowException(message)
@@ -127,4 +126,4 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
127
126
  self.log.debug(
128
127
  "Current status of the pipeline workflow for %s: %s.", self.pipeline_id, pipeline_status
129
128
  )
130
- return pipeline_status in self.expected_statuses
129
+ return pipeline_status is not None and pipeline_status in self.expected_statuses
@@ -25,19 +25,18 @@ from typing import TYPE_CHECKING
25
25
  if TYPE_CHECKING:
26
26
  from google.api_core.retry import Retry
27
27
 
28
- from airflow.utils.context import Context
28
+ from airflow.providers.common.compat.sdk import Context
29
29
 
30
30
  from google.api_core.exceptions import GoogleAPICallError
31
31
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
32
32
  from google.cloud.dataplex_v1.types import DataScanJob
33
33
 
34
- from airflow.exceptions import AirflowException
34
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
35
35
  from airflow.providers.google.cloud.hooks.dataplex import (
36
36
  AirflowDataQualityScanException,
37
37
  AirflowDataQualityScanResultTimeoutException,
38
38
  DataplexHook,
39
39
  )
40
- from airflow.sensors.base import BaseSensorOperator
41
40
 
42
41
 
43
42
  class TaskState:
@@ -22,11 +22,11 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
+ from airflow.providers.common.compat.sdk import BaseSensorOperator
25
26
  from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook, JobGroupStatuses
26
- from airflow.sensors.base import BaseSensorOperator
27
27
 
28
28
  if TYPE_CHECKING:
29
- from airflow.utils.context import Context
29
+ from airflow.providers.common.compat.sdk import Context
30
30
 
31
31
 
32
32
  class DataprepJobGroupIsFinishedSensor(BaseSensorOperator):
@@ -26,13 +26,12 @@ from typing import TYPE_CHECKING
26
26
  from google.api_core.exceptions import ServerError
27
27
  from google.cloud.dataproc_v1.types import Batch, JobStatus
28
28
 
29
- from airflow.exceptions import AirflowException
29
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
30
30
  from airflow.providers.google.cloud.hooks.dataproc import DataprocHook
31
31
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
32
- from airflow.sensors.base import BaseSensorOperator
33
32
 
34
33
  if TYPE_CHECKING:
35
- from airflow.utils.context import Context
34
+ from airflow.providers.common.compat.sdk import Context
36
35
 
37
36
 
38
37
  class DataprocJobSensor(BaseSensorOperator):
@@ -20,15 +20,14 @@ from __future__ import annotations
20
20
  from collections.abc import Sequence
21
21
  from typing import TYPE_CHECKING
22
22
 
23
- from airflow.exceptions import AirflowException
23
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
24
24
  from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
25
25
  from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
26
- from airflow.sensors.base import BaseSensorOperator
27
26
 
28
27
  if TYPE_CHECKING:
29
28
  from google.api_core.operation import Operation
30
29
 
31
- from airflow.utils.context import Context
30
+ from airflow.providers.common.compat.sdk import Context
32
31
 
33
32
 
34
33
  class MetastoreHivePartitionSensor(BaseSensorOperator):
@@ -21,14 +21,14 @@ from __future__ import annotations
21
21
 
22
22
  import os
23
23
  import textwrap
24
- from collections.abc import Sequence
24
+ from collections.abc import Callable, Sequence
25
25
  from datetime import datetime, timedelta
26
- from typing import TYPE_CHECKING, Any, Callable
26
+ from typing import TYPE_CHECKING, Any
27
27
 
28
28
  from google.cloud.storage.retry import DEFAULT_RETRY
29
29
 
30
30
  from airflow.configuration import conf
31
- from airflow.exceptions import AirflowException
31
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, poke_mode_only
32
32
  from airflow.providers.google.cloud.hooks.gcs import GCSHook
33
33
  from airflow.providers.google.cloud.triggers.gcs import (
34
34
  GCSBlobTrigger,
@@ -36,12 +36,11 @@ from airflow.providers.google.cloud.triggers.gcs import (
36
36
  GCSPrefixBlobTrigger,
37
37
  GCSUploadSessionTrigger,
38
38
  )
39
- from airflow.sensors.base import BaseSensorOperator, poke_mode_only
40
39
 
41
40
  if TYPE_CHECKING:
42
41
  from google.api_core.retry import Retry
43
42
 
44
- from airflow.utils.context import Context
43
+ from airflow.providers.common.compat.sdk import Context
45
44
 
46
45
 
47
46
  class GCSObjectExistenceSensor(BaseSensorOperator):
@@ -21,12 +21,11 @@ from __future__ import annotations
21
21
 
22
22
  from typing import TYPE_CHECKING
23
23
 
24
- from airflow.exceptions import AirflowException
24
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
25
25
  from airflow.providers.google.cloud.hooks.looker import JobStatus, LookerHook
26
- from airflow.sensors.base import BaseSensorOperator
27
26
 
28
27
  if TYPE_CHECKING:
29
- from airflow.utils.context import Context
28
+ from airflow.providers.common.compat.sdk import Context
30
29
 
31
30
 
32
31
  class LookerCheckPdtBuildSensor(BaseSensorOperator):
@@ -19,21 +19,20 @@
19
19
 
20
20
  from __future__ import annotations
21
21
 
22
- from collections.abc import Sequence
22
+ from collections.abc import Callable, Sequence
23
23
  from datetime import timedelta
24
- from typing import TYPE_CHECKING, Any, Callable
24
+ from typing import TYPE_CHECKING, Any
25
25
 
26
26
  from google.cloud import pubsub_v1
27
27
  from google.cloud.pubsub_v1.types import ReceivedMessage
28
28
 
29
29
  from airflow.configuration import conf
30
- from airflow.exceptions import AirflowException
30
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
31
31
  from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
32
32
  from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
33
- from airflow.sensors.base import BaseSensorOperator
34
33
 
35
34
  if TYPE_CHECKING:
36
- from airflow.utils.context import Context
35
+ from airflow.providers.common.compat.sdk import Context
37
36
 
38
37
 
39
38
  class PubSubMessageTransformException(AirflowException):
@@ -22,12 +22,12 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
+ from airflow.providers.common.compat.sdk import BaseSensorOperator
25
26
  from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
26
27
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
27
- from airflow.sensors.base import BaseSensorOperator
28
28
 
29
29
  if TYPE_CHECKING:
30
- from airflow.utils.context import Context
30
+ from airflow.providers.common.compat.sdk import Context
31
31
 
32
32
 
33
33
  class TaskQueueEmptySensor(BaseSensorOperator):
@@ -23,12 +23,11 @@ import time
23
23
  from collections.abc import Sequence
24
24
  from typing import TYPE_CHECKING
25
25
 
26
- from airflow.exceptions import AirflowException
26
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
27
27
  from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
28
- from airflow.sensors.base import BaseSensorOperator
29
28
 
30
29
  if TYPE_CHECKING:
31
- from airflow.utils.context import Context
30
+ from airflow.providers.common.compat.sdk import Context
32
31
 
33
32
 
34
33
  class FeatureViewSyncSensor(BaseSensorOperator):
@@ -22,15 +22,14 @@ from typing import TYPE_CHECKING
22
22
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
23
23
  from google.cloud.workflows.executions_v1beta import Execution
24
24
 
25
- from airflow.exceptions import AirflowException
25
+ from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
26
26
  from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook
27
27
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
28
- from airflow.sensors.base import BaseSensorOperator
29
28
 
30
29
  if TYPE_CHECKING:
31
30
  from google.api_core.retry import Retry
32
31
 
33
- from airflow.utils.context import Context
32
+ from airflow.providers.common.compat.sdk import Context
34
33
 
35
34
 
36
35
  class WorkflowExecutionSensor(BaseSensorOperator):
@@ -35,7 +35,7 @@ except ModuleNotFoundError as e:
35
35
  raise AirflowOptionalProviderFeatureException(e)
36
36
 
37
37
  if TYPE_CHECKING:
38
- from airflow.utils.context import Context
38
+ from airflow.providers.common.compat.sdk import Context
39
39
 
40
40
 
41
41
  class ADLSToGCSOperator(ADLSListOperator):
@@ -21,8 +21,8 @@ import tempfile
21
21
  from collections.abc import Sequence
22
22
  from typing import TYPE_CHECKING
23
23
 
24
- from airflow.models import BaseOperator
25
24
  from airflow.providers.google.cloud.hooks.gcs import GCSHook
25
+ from airflow.providers.google.version_compat import BaseOperator
26
26
 
27
27
  try:
28
28
  from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
@@ -32,7 +32,7 @@ except ModuleNotFoundError as e:
32
32
  raise AirflowOptionalProviderFeatureException(e)
33
33
 
34
34
  if TYPE_CHECKING:
35
- from airflow.utils.context import Context
35
+ from airflow.providers.common.compat.sdk import Context
36
36
 
37
37
 
38
38
  class AzureBlobStorageToGCSOperator(BaseOperator):
@@ -22,9 +22,10 @@ from collections.abc import Sequence
22
22
  from tempfile import NamedTemporaryFile
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
26
- from airflow.models import BaseOperator
25
+ from airflow.exceptions import AirflowProviderDeprecationWarning
26
+ from airflow.providers.common.compat.sdk import AirflowException
27
27
  from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url, gcs_object_is_directory
28
+ from airflow.providers.google.version_compat import BaseOperator
28
29
 
29
30
  try:
30
31
  from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook
@@ -34,7 +35,7 @@ except ModuleNotFoundError as e:
34
35
  raise AirflowOptionalProviderFeatureException(e)
35
36
 
36
37
  if TYPE_CHECKING:
37
- from airflow.utils.context import Context
38
+ from airflow.providers.common.compat.sdk import Context
38
39
 
39
40
 
40
41
  class AzureFileShareToGCSOperator(BaseOperator):
@@ -22,12 +22,13 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from airflow.models import BaseOperator
26
25
  from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
27
26
  from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
27
+ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
28
+ from airflow.providers.google.version_compat import BaseOperator
28
29
 
29
30
  if TYPE_CHECKING:
30
- from airflow.utils.context import Context
31
+ from airflow.providers.common.compat.sdk import Context
31
32
 
32
33
 
33
34
  class BigQueryToBigQueryOperator(BaseOperator):
@@ -73,6 +74,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
73
74
  If set as a sequence, the identities from the list must grant
74
75
  Service Account Token Creator IAM role to the directly preceding identity, with first
75
76
  account from the list granting this role to the originating account (templated).
77
+ :param project_id: Google Cloud Project where the job is running
76
78
  """
77
79
 
78
80
  template_fields: Sequence[str] = (
@@ -93,6 +95,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
93
95
  write_disposition: str = "WRITE_EMPTY",
94
96
  create_disposition: str = "CREATE_IF_NEEDED",
95
97
  gcp_conn_id: str = "google_cloud_default",
98
+ project_id: str = PROVIDE_PROJECT_ID,
96
99
  labels: dict | None = None,
97
100
  encryption_configuration: dict | None = None,
98
101
  location: str | None = None,
@@ -112,6 +115,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
112
115
  self.impersonation_chain = impersonation_chain
113
116
  self.hook: BigQueryHook | None = None
114
117
  self._job_conf: dict = {}
118
+ self.project_id = project_id
115
119
 
116
120
  def _prepare_job_configuration(self):
117
121
  self.source_project_dataset_tables = (
@@ -124,7 +128,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
124
128
  for source_project_dataset_table in self.source_project_dataset_tables:
125
129
  source_project, source_dataset, source_table = self.hook.split_tablename(
126
130
  table_input=source_project_dataset_table,
127
- default_project_id=self.hook.project_id,
131
+ default_project_id=self.project_id,
128
132
  var_name="source_project_dataset_table",
129
133
  )
130
134
  source_project_dataset_tables_fixup.append(
@@ -133,7 +137,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
133
137
 
134
138
  destination_project, destination_dataset, destination_table = self.hook.split_tablename(
135
139
  table_input=self.destination_project_dataset_table,
136
- default_project_id=self.hook.project_id,
140
+ default_project_id=self.project_id,
137
141
  )
138
142
  configuration = {
139
143
  "copy": {
@@ -168,18 +172,17 @@ class BigQueryToBigQueryOperator(BaseOperator):
168
172
  impersonation_chain=self.impersonation_chain,
169
173
  )
170
174
 
171
- if not self.hook.project_id:
172
- raise ValueError("The project_id should be set")
175
+ if not self.project_id:
176
+ self.project_id = self.hook.project_id
173
177
 
174
178
  configuration = self._prepare_job_configuration()
175
179
  self._job_conf = self.hook.insert_job(
176
- configuration=configuration, project_id=self.hook.project_id
180
+ configuration=configuration, project_id=self.project_id
177
181
  ).to_api_repr()
178
182
 
179
183
  dest_table_info = self._job_conf["configuration"]["copy"]["destinationTable"]
180
184
  BigQueryTableLink.persist(
181
185
  context=context,
182
- task_instance=self,
183
186
  dataset_id=dest_table_info["datasetId"],
184
187
  project_id=dest_table_info["projectId"],
185
188
  table_id=dest_table_info["tableId"],
@@ -26,18 +26,18 @@ from google.api_core.exceptions import Conflict
26
26
  from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
27
27
 
28
28
  from airflow.configuration import conf
29
- from airflow.exceptions import AirflowException
30
- from airflow.models import BaseOperator
29
+ from airflow.providers.common.compat.sdk import AirflowException
31
30
  from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
32
31
  from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
33
32
  from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
34
33
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
34
+ from airflow.providers.google.version_compat import BaseOperator
35
35
  from airflow.utils.helpers import merge_dicts
36
36
 
37
37
  if TYPE_CHECKING:
38
38
  from google.api_core.retry import Retry
39
39
 
40
- from airflow.utils.context import Context
40
+ from airflow.providers.common.compat.sdk import Context
41
41
 
42
42
 
43
43
  class BigQueryToGCSOperator(BaseOperator):
@@ -87,6 +87,7 @@ class BigQueryToGCSOperator(BaseOperator):
87
87
  :param reattach_states: Set of BigQuery job's states in case of which we should reattach
88
88
  to the job. Should be other than final states.
89
89
  :param deferrable: Run operator in the deferrable mode
90
+ :return: URIs for the objects created in Google Cloud Storage
90
91
  """
91
92
 
92
93
  template_fields: Sequence[str] = (
@@ -215,8 +216,9 @@ class BigQueryToGCSOperator(BaseOperator):
215
216
  job_id=self.job_id,
216
217
  dag_id=self.dag_id,
217
218
  task_id=self.task_id,
218
- logical_date=context["logical_date"],
219
+ logical_date=None,
219
220
  configuration=configuration,
221
+ run_after=hook.get_run_after_or_logical_date(context),
220
222
  force_rerun=self.force_rerun,
221
223
  )
222
224
 
@@ -254,7 +256,6 @@ class BigQueryToGCSOperator(BaseOperator):
254
256
  dataset_id, project_id, table_id = conf["datasetId"], conf["projectId"], conf["tableId"]
255
257
  BigQueryTableLink.persist(
256
258
  context=context,
257
- task_instance=self,
258
259
  dataset_id=dataset_id,
259
260
  project_id=project_id,
260
261
  table_id=table_id,
@@ -275,6 +276,8 @@ class BigQueryToGCSOperator(BaseOperator):
275
276
  else:
276
277
  job.result(timeout=self.result_timeout, retry=self.result_retry)
277
278
 
279
+ return self.destination_cloud_storage_uris
280
+
278
281
  def execute_complete(self, context: Context, event: dict[str, Any]):
279
282
  """
280
283
  Return immediately and relies on trigger to throw a success event. Callback for the trigger.
@@ -291,6 +294,8 @@ class BigQueryToGCSOperator(BaseOperator):
291
294
  # Save job_id as an attribute to be later used by listeners
292
295
  self.job_id = event.get("job_id")
293
296
 
297
+ return self.destination_cloud_storage_uris
298
+
294
299
  def get_openlineage_facets_on_complete(self, task_instance):
295
300
  """Implement on_complete as we will include final BQ job id."""
296
301
  from airflow.providers.common.compat.openlineage.facet import (
@@ -21,6 +21,7 @@ from __future__ import annotations
21
21
 
22
22
  import warnings
23
23
  from collections.abc import Sequence
24
+ from functools import cached_property
24
25
  from typing import TYPE_CHECKING
25
26
 
26
27
  from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -29,7 +30,7 @@ from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToS
29
30
  from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
30
31
 
31
32
  if TYPE_CHECKING:
32
- from airflow.utils.context import Context
33
+ from airflow.providers.common.compat.sdk import Context
33
34
 
34
35
 
35
36
  class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
@@ -94,14 +95,17 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
94
95
  self.mssql_conn_id = mssql_conn_id
95
96
  self.source_project_dataset_table = source_project_dataset_table
96
97
 
97
- def get_sql_hook(self) -> MsSqlHook:
98
+ @cached_property
99
+ def mssql_hook(self) -> MsSqlHook:
98
100
  return MsSqlHook(schema=self.database, mssql_conn_id=self.mssql_conn_id)
99
101
 
102
+ def get_sql_hook(self) -> MsSqlHook:
103
+ return self.mssql_hook
104
+
100
105
  def persist_links(self, context: Context) -> None:
101
106
  project_id, dataset_id, table_id = self.source_project_dataset_table.split(".")
102
107
  BigQueryTableLink.persist(
103
108
  context=context,
104
- task_instance=self,
105
109
  dataset_id=dataset_id,
106
110
  project_id=project_id,
107
111
  table_id=table_id,
@@ -21,6 +21,7 @@ from __future__ import annotations
21
21
 
22
22
  import warnings
23
23
  from collections.abc import Sequence
24
+ from functools import cached_property
24
25
 
25
26
  from airflow.exceptions import AirflowProviderDeprecationWarning
26
27
  from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
@@ -76,5 +77,15 @@ class BigQueryToMySqlOperator(BigQueryToSqlBaseOperator):
76
77
  )
77
78
  self.mysql_conn_id = mysql_conn_id
78
79
 
79
- def get_sql_hook(self) -> MySqlHook:
80
+ @cached_property
81
+ def mysql_hook(self) -> MySqlHook:
80
82
  return MySqlHook(schema=self.database, mysql_conn_id=self.mysql_conn_id)
83
+
84
+ def get_sql_hook(self) -> MySqlHook:
85
+ return self.mysql_hook
86
+
87
+ def execute(self, context):
88
+ # Set source_project_dataset_table here, after hooks are initialized and project_id is available
89
+ project_id = self.bigquery_hook.project_id
90
+ self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
91
+ return super().execute(context)
@@ -19,15 +19,19 @@
19
19
 
20
20
  from __future__ import annotations
21
21
 
22
+ from functools import cached_property
22
23
  from typing import TYPE_CHECKING
23
24
 
25
+ from psycopg2.extensions import register_adapter
26
+ from psycopg2.extras import Json
27
+
24
28
  from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
25
29
  from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
26
30
  from airflow.providers.google.cloud.utils.bigquery_get_data import bigquery_get_data
27
31
  from airflow.providers.postgres.hooks.postgres import PostgresHook
28
32
 
29
33
  if TYPE_CHECKING:
30
- from airflow.utils.context import Context
34
+ from airflow.providers.common.compat.sdk import Context
31
35
 
32
36
 
33
37
  class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
@@ -75,26 +79,36 @@ class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
75
79
  self.postgres_conn_id = postgres_conn_id
76
80
  self.replace_index = replace_index
77
81
 
78
- def get_sql_hook(self) -> PostgresHook:
82
+ @cached_property
83
+ def postgres_hook(self) -> PostgresHook:
84
+ register_adapter(list, Json)
85
+ register_adapter(dict, Json)
79
86
  return PostgresHook(database=self.database, postgres_conn_id=self.postgres_conn_id)
80
87
 
88
+ def get_sql_hook(self) -> PostgresHook:
89
+ return self.postgres_hook
90
+
81
91
  def execute(self, context: Context) -> None:
82
- big_query_hook = BigQueryHook(
83
- gcp_conn_id=self.gcp_conn_id,
84
- location=self.location,
85
- impersonation_chain=self.impersonation_chain,
86
- )
92
+ if not self.bigquery_hook:
93
+ self.bigquery_hook = BigQueryHook(
94
+ gcp_conn_id=self.gcp_conn_id,
95
+ location=self.location,
96
+ impersonation_chain=self.impersonation_chain,
97
+ )
98
+ # Set source_project_dataset_table here, after hooks are initialized and project_id is available
99
+ project_id = self.bigquery_hook.project_id
100
+ self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
101
+
87
102
  self.persist_links(context)
88
- sql_hook: PostgresHook = self.get_sql_hook()
89
103
  for rows in bigquery_get_data(
90
104
  self.log,
91
105
  self.dataset_id,
92
106
  self.table_id,
93
- big_query_hook,
107
+ self.bigquery_hook,
94
108
  self.batch_size,
95
109
  self.selected_fields,
96
110
  ):
97
- sql_hook.insert_rows(
111
+ self.postgres_hook.insert_rows(
98
112
  table=self.target_table_name,
99
113
  rows=rows,
100
114
  target_fields=self.selected_fields,