apache-airflow-providers-google 10.17.0rc1__py3-none-any.whl → 10.18.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.
- airflow/providers/google/__init__.py +5 -8
- airflow/providers/google/cloud/hooks/automl.py +35 -1
- airflow/providers/google/cloud/hooks/bigquery.py +126 -41
- airflow/providers/google/cloud/hooks/cloud_composer.py +250 -2
- airflow/providers/google/cloud/hooks/cloud_sql.py +154 -7
- airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +7 -2
- airflow/providers/google/cloud/hooks/compute_ssh.py +2 -1
- airflow/providers/google/cloud/hooks/dataflow.py +246 -32
- airflow/providers/google/cloud/hooks/dataplex.py +6 -2
- airflow/providers/google/cloud/hooks/dlp.py +14 -14
- airflow/providers/google/cloud/hooks/gcs.py +6 -2
- airflow/providers/google/cloud/hooks/gdm.py +2 -2
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +2 -2
- airflow/providers/google/cloud/hooks/mlengine.py +8 -4
- airflow/providers/google/cloud/hooks/pubsub.py +1 -1
- airflow/providers/google/cloud/hooks/secret_manager.py +252 -4
- airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +1431 -74
- airflow/providers/google/cloud/hooks/vertex_ai/prediction_service.py +91 -0
- airflow/providers/google/cloud/links/vertex_ai.py +2 -1
- airflow/providers/google/cloud/log/gcs_task_handler.py +2 -1
- airflow/providers/google/cloud/operators/automl.py +243 -37
- airflow/providers/google/cloud/operators/bigquery.py +164 -62
- airflow/providers/google/cloud/operators/bigquery_dts.py +4 -3
- airflow/providers/google/cloud/operators/bigtable.py +7 -6
- airflow/providers/google/cloud/operators/cloud_build.py +12 -11
- airflow/providers/google/cloud/operators/cloud_composer.py +147 -2
- airflow/providers/google/cloud/operators/cloud_memorystore.py +17 -16
- airflow/providers/google/cloud/operators/cloud_sql.py +60 -17
- airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +35 -16
- airflow/providers/google/cloud/operators/compute.py +12 -11
- airflow/providers/google/cloud/operators/datacatalog.py +21 -20
- airflow/providers/google/cloud/operators/dataflow.py +59 -42
- airflow/providers/google/cloud/operators/datafusion.py +11 -10
- airflow/providers/google/cloud/operators/datapipeline.py +3 -2
- airflow/providers/google/cloud/operators/dataprep.py +5 -4
- airflow/providers/google/cloud/operators/dataproc.py +20 -17
- airflow/providers/google/cloud/operators/datastore.py +8 -7
- airflow/providers/google/cloud/operators/dlp.py +31 -30
- airflow/providers/google/cloud/operators/functions.py +4 -3
- airflow/providers/google/cloud/operators/gcs.py +66 -41
- airflow/providers/google/cloud/operators/kubernetes_engine.py +256 -49
- airflow/providers/google/cloud/operators/life_sciences.py +2 -1
- airflow/providers/google/cloud/operators/mlengine.py +11 -10
- airflow/providers/google/cloud/operators/pubsub.py +6 -5
- airflow/providers/google/cloud/operators/spanner.py +7 -6
- airflow/providers/google/cloud/operators/speech_to_text.py +2 -1
- airflow/providers/google/cloud/operators/stackdriver.py +11 -10
- airflow/providers/google/cloud/operators/tasks.py +14 -13
- airflow/providers/google/cloud/operators/text_to_speech.py +2 -1
- airflow/providers/google/cloud/operators/translate_speech.py +2 -1
- airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +333 -26
- airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +20 -12
- airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +0 -1
- airflow/providers/google/cloud/operators/vision.py +13 -12
- airflow/providers/google/cloud/operators/workflows.py +12 -14
- airflow/providers/google/cloud/secrets/secret_manager.py +2 -1
- airflow/providers/google/cloud/sensors/bigquery_dts.py +2 -1
- airflow/providers/google/cloud/sensors/bigtable.py +2 -1
- airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +2 -1
- airflow/providers/google/cloud/sensors/dataflow.py +239 -52
- airflow/providers/google/cloud/sensors/datafusion.py +2 -1
- airflow/providers/google/cloud/sensors/dataproc.py +3 -2
- airflow/providers/google/cloud/sensors/gcs.py +14 -12
- airflow/providers/google/cloud/sensors/tasks.py +2 -1
- airflow/providers/google/cloud/sensors/workflows.py +2 -1
- airflow/providers/google/cloud/transfers/adls_to_gcs.py +8 -2
- airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +7 -1
- airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +7 -1
- airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +2 -1
- airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +1 -1
- airflow/providers/google/cloud/transfers/bigquery_to_sql.py +1 -0
- airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +5 -6
- airflow/providers/google/cloud/transfers/gcs_to_gcs.py +22 -12
- airflow/providers/google/cloud/triggers/bigquery.py +75 -6
- airflow/providers/google/cloud/triggers/cloud_composer.py +68 -0
- airflow/providers/google/cloud/triggers/cloud_sql.py +2 -1
- airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +2 -1
- airflow/providers/google/cloud/triggers/dataflow.py +504 -4
- airflow/providers/google/cloud/triggers/dataproc.py +190 -27
- airflow/providers/google/cloud/triggers/kubernetes_engine.py +2 -3
- airflow/providers/google/cloud/triggers/mlengine.py +2 -1
- airflow/providers/google/cloud/triggers/vertex_ai.py +94 -0
- airflow/providers/google/common/hooks/base_google.py +45 -7
- airflow/providers/google/firebase/hooks/firestore.py +2 -2
- airflow/providers/google/firebase/operators/firestore.py +2 -1
- airflow/providers/google/get_provider_info.py +5 -3
- {apache_airflow_providers_google-10.17.0rc1.dist-info → apache_airflow_providers_google-10.18.0.dist-info}/METADATA +18 -18
- {apache_airflow_providers_google-10.17.0rc1.dist-info → apache_airflow_providers_google-10.18.0.dist-info}/RECORD +90 -90
- airflow/providers/google/cloud/example_dags/example_cloud_sql_query.py +0 -289
- {apache_airflow_providers_google-10.17.0rc1.dist-info → apache_airflow_providers_google-10.18.0.dist-info}/WHEEL +0 -0
- {apache_airflow_providers_google-10.17.0rc1.dist-info → apache_airflow_providers_google-10.18.0.dist-info}/entry_points.txt +0 -0
@@ -40,6 +40,7 @@ from airflow.providers.google.cloud.links.cloud_build import (
|
|
40
40
|
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
|
41
41
|
from airflow.providers.google.cloud.triggers.cloud_build import CloudBuildCreateBuildTrigger
|
42
42
|
from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
|
43
|
+
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
43
44
|
from airflow.utils import yaml
|
44
45
|
from airflow.utils.helpers import exactly_one
|
45
46
|
|
@@ -87,7 +88,7 @@ class CloudBuildCancelBuildOperator(GoogleCloudBaseOperator):
|
|
87
88
|
self,
|
88
89
|
*,
|
89
90
|
id_: str,
|
90
|
-
project_id: str
|
91
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
91
92
|
retry: Retry | _MethodDefault = DEFAULT,
|
92
93
|
timeout: float | None = None,
|
93
94
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -171,7 +172,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
|
|
171
172
|
self,
|
172
173
|
*,
|
173
174
|
build: dict | Build,
|
174
|
-
project_id: str
|
175
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
175
176
|
wait: bool = True,
|
176
177
|
retry: Retry | _MethodDefault = DEFAULT,
|
177
178
|
timeout: float | None = None,
|
@@ -316,7 +317,7 @@ class CloudBuildCreateBuildTriggerOperator(GoogleCloudBaseOperator):
|
|
316
317
|
self,
|
317
318
|
*,
|
318
319
|
trigger: dict | BuildTrigger,
|
319
|
-
project_id: str
|
320
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
320
321
|
retry: Retry | _MethodDefault = DEFAULT,
|
321
322
|
timeout: float | None = None,
|
322
323
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -399,7 +400,7 @@ class CloudBuildDeleteBuildTriggerOperator(GoogleCloudBaseOperator):
|
|
399
400
|
self,
|
400
401
|
*,
|
401
402
|
trigger_id: str,
|
402
|
-
project_id: str
|
403
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
403
404
|
retry: Retry | _MethodDefault = DEFAULT,
|
404
405
|
timeout: float | None = None,
|
405
406
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -473,7 +474,7 @@ class CloudBuildGetBuildOperator(GoogleCloudBaseOperator):
|
|
473
474
|
self,
|
474
475
|
*,
|
475
476
|
id_: str,
|
476
|
-
project_id: str
|
477
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
477
478
|
retry: Retry | _MethodDefault = DEFAULT,
|
478
479
|
timeout: float | None = None,
|
479
480
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -549,7 +550,7 @@ class CloudBuildGetBuildTriggerOperator(GoogleCloudBaseOperator):
|
|
549
550
|
self,
|
550
551
|
*,
|
551
552
|
trigger_id: str,
|
552
|
-
project_id: str
|
553
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
553
554
|
retry: Retry | _MethodDefault = DEFAULT,
|
554
555
|
timeout: float | None = None,
|
555
556
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -627,7 +628,7 @@ class CloudBuildListBuildTriggersOperator(GoogleCloudBaseOperator):
|
|
627
628
|
self,
|
628
629
|
*,
|
629
630
|
location: str = "global",
|
630
|
-
project_id: str
|
631
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
631
632
|
page_size: int | None = None,
|
632
633
|
page_token: str | None = None,
|
633
634
|
retry: Retry | _MethodDefault = DEFAULT,
|
@@ -707,7 +708,7 @@ class CloudBuildListBuildsOperator(GoogleCloudBaseOperator):
|
|
707
708
|
self,
|
708
709
|
*,
|
709
710
|
location: str = "global",
|
710
|
-
project_id: str
|
711
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
711
712
|
page_size: int | None = None,
|
712
713
|
filter_: str | None = None,
|
713
714
|
retry: Retry | _MethodDefault = DEFAULT,
|
@@ -783,7 +784,7 @@ class CloudBuildRetryBuildOperator(GoogleCloudBaseOperator):
|
|
783
784
|
self,
|
784
785
|
*,
|
785
786
|
id_: str,
|
786
|
-
project_id: str
|
787
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
787
788
|
wait: bool = True,
|
788
789
|
retry: Retry | _MethodDefault = DEFAULT,
|
789
790
|
timeout: float | None = None,
|
@@ -868,7 +869,7 @@ class CloudBuildRunBuildTriggerOperator(GoogleCloudBaseOperator):
|
|
868
869
|
*,
|
869
870
|
trigger_id: str,
|
870
871
|
source: dict | RepoSource,
|
871
|
-
project_id: str
|
872
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
872
873
|
wait: bool = True,
|
873
874
|
retry: Retry | _MethodDefault = DEFAULT,
|
874
875
|
timeout: float | None = None,
|
@@ -953,7 +954,7 @@ class CloudBuildUpdateBuildTriggerOperator(GoogleCloudBaseOperator):
|
|
953
954
|
*,
|
954
955
|
trigger_id: str,
|
955
956
|
trigger: dict | BuildTrigger,
|
956
|
-
project_id: str
|
957
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
957
958
|
retry: Retry | _MethodDefault = DEFAULT,
|
958
959
|
timeout: float | None = None,
|
959
960
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -17,19 +17,23 @@
|
|
17
17
|
# under the License.
|
18
18
|
from __future__ import annotations
|
19
19
|
|
20
|
+
import shlex
|
20
21
|
from typing import TYPE_CHECKING, Sequence
|
21
22
|
|
22
23
|
from google.api_core.exceptions import AlreadyExists
|
23
24
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
24
25
|
from google.cloud.orchestration.airflow.service_v1 import ImageVersion
|
25
|
-
from google.cloud.orchestration.airflow.service_v1.types import Environment
|
26
|
+
from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
|
26
27
|
|
27
28
|
from airflow.configuration import conf
|
28
29
|
from airflow.exceptions import AirflowException
|
29
30
|
from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
|
30
31
|
from airflow.providers.google.cloud.links.base import BaseGoogleLink
|
31
32
|
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
|
32
|
-
from airflow.providers.google.cloud.triggers.cloud_composer import
|
33
|
+
from airflow.providers.google.cloud.triggers.cloud_composer import (
|
34
|
+
CloudComposerAirflowCLICommandTrigger,
|
35
|
+
CloudComposerExecutionTrigger,
|
36
|
+
)
|
33
37
|
from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
|
34
38
|
|
35
39
|
if TYPE_CHECKING:
|
@@ -651,3 +655,144 @@ class CloudComposerListImageVersionsOperator(GoogleCloudBaseOperator):
|
|
651
655
|
metadata=self.metadata,
|
652
656
|
)
|
653
657
|
return [ImageVersion.to_dict(image) for image in result]
|
658
|
+
|
659
|
+
|
660
|
+
class CloudComposerRunAirflowCLICommandOperator(GoogleCloudBaseOperator):
|
661
|
+
"""
|
662
|
+
Run Airflow command for provided Composer environment.
|
663
|
+
|
664
|
+
:param project_id: The ID of the Google Cloud project that the service belongs to.
|
665
|
+
:param region: The ID of the Google Cloud region that the service belongs to.
|
666
|
+
:param environment_id: The ID of the Google Cloud environment that the service belongs to.
|
667
|
+
:param command: Airflow command.
|
668
|
+
:param retry: Designation of what errors, if any, should be retried.
|
669
|
+
:param timeout: The timeout for this request.
|
670
|
+
:param metadata: Strings which should be sent along with the request as metadata.
|
671
|
+
:param gcp_conn_id: The connection ID used to connect to Google Cloud Platform.
|
672
|
+
:param impersonation_chain: Optional service account to impersonate using short-term
|
673
|
+
credentials, or chained list of accounts required to get the access_token
|
674
|
+
of the last account in the list, which will be impersonated in the request.
|
675
|
+
If set as a string, the account must grant the originating account
|
676
|
+
the Service Account Token Creator IAM role.
|
677
|
+
If set as a sequence, the identities from the list must grant
|
678
|
+
Service Account Token Creator IAM role to the directly preceding identity, with first
|
679
|
+
account from the list granting this role to the originating account (templated).
|
680
|
+
:param deferrable: Run operator in the deferrable mode
|
681
|
+
:param poll_interval: Optional: Control the rate of the poll for the result of deferrable run.
|
682
|
+
By default, the trigger will poll every 10 seconds.
|
683
|
+
"""
|
684
|
+
|
685
|
+
template_fields = (
|
686
|
+
"project_id",
|
687
|
+
"region",
|
688
|
+
"environment_id",
|
689
|
+
"command",
|
690
|
+
"impersonation_chain",
|
691
|
+
)
|
692
|
+
|
693
|
+
def __init__(
|
694
|
+
self,
|
695
|
+
*,
|
696
|
+
project_id: str,
|
697
|
+
region: str,
|
698
|
+
environment_id: str,
|
699
|
+
command: str,
|
700
|
+
retry: Retry | _MethodDefault = DEFAULT,
|
701
|
+
timeout: float | None = None,
|
702
|
+
metadata: Sequence[tuple[str, str]] = (),
|
703
|
+
gcp_conn_id: str = "google_cloud_default",
|
704
|
+
impersonation_chain: str | Sequence[str] | None = None,
|
705
|
+
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
|
706
|
+
poll_interval: int = 10,
|
707
|
+
**kwargs,
|
708
|
+
) -> None:
|
709
|
+
super().__init__(**kwargs)
|
710
|
+
self.project_id = project_id
|
711
|
+
self.region = region
|
712
|
+
self.environment_id = environment_id
|
713
|
+
self.command = command
|
714
|
+
self.retry = retry
|
715
|
+
self.timeout = timeout
|
716
|
+
self.metadata = metadata
|
717
|
+
self.gcp_conn_id = gcp_conn_id
|
718
|
+
self.impersonation_chain = impersonation_chain
|
719
|
+
self.deferrable = deferrable
|
720
|
+
self.poll_interval = poll_interval
|
721
|
+
|
722
|
+
def execute(self, context: Context):
|
723
|
+
hook = CloudComposerHook(
|
724
|
+
gcp_conn_id=self.gcp_conn_id,
|
725
|
+
impersonation_chain=self.impersonation_chain,
|
726
|
+
)
|
727
|
+
|
728
|
+
self.log.info("Executing the command: [ airflow %s ]...", self.command)
|
729
|
+
|
730
|
+
cmd, subcommand, parameters = self._parse_cmd_to_args(self.command)
|
731
|
+
execution_cmd_info = hook.execute_airflow_command(
|
732
|
+
project_id=self.project_id,
|
733
|
+
region=self.region,
|
734
|
+
environment_id=self.environment_id,
|
735
|
+
command=cmd,
|
736
|
+
subcommand=subcommand,
|
737
|
+
parameters=parameters,
|
738
|
+
retry=self.retry,
|
739
|
+
timeout=self.timeout,
|
740
|
+
metadata=self.metadata,
|
741
|
+
)
|
742
|
+
execution_cmd_info_dict = ExecuteAirflowCommandResponse.to_dict(execution_cmd_info)
|
743
|
+
|
744
|
+
self.log.info("Command has been started. execution_id=%s", execution_cmd_info_dict["execution_id"])
|
745
|
+
|
746
|
+
if self.deferrable:
|
747
|
+
self.defer(
|
748
|
+
trigger=CloudComposerAirflowCLICommandTrigger(
|
749
|
+
project_id=self.project_id,
|
750
|
+
region=self.region,
|
751
|
+
environment_id=self.environment_id,
|
752
|
+
execution_cmd_info=execution_cmd_info_dict,
|
753
|
+
gcp_conn_id=self.gcp_conn_id,
|
754
|
+
impersonation_chain=self.impersonation_chain,
|
755
|
+
poll_interval=self.poll_interval,
|
756
|
+
),
|
757
|
+
method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
|
758
|
+
)
|
759
|
+
return
|
760
|
+
|
761
|
+
result = hook.wait_command_execution_result(
|
762
|
+
project_id=self.project_id,
|
763
|
+
region=self.region,
|
764
|
+
environment_id=self.environment_id,
|
765
|
+
execution_cmd_info=execution_cmd_info_dict,
|
766
|
+
retry=self.retry,
|
767
|
+
timeout=self.timeout,
|
768
|
+
metadata=self.metadata,
|
769
|
+
poll_interval=self.poll_interval,
|
770
|
+
)
|
771
|
+
result_str = self._merge_cmd_output_result(result)
|
772
|
+
self.log.info("Command execution result:\n%s", result_str)
|
773
|
+
return result
|
774
|
+
|
775
|
+
def execute_complete(self, context: Context, event: dict) -> dict:
|
776
|
+
if event and event["status"] == "error":
|
777
|
+
raise AirflowException(event["message"])
|
778
|
+
result: dict = event["result"]
|
779
|
+
result_str = self._merge_cmd_output_result(result)
|
780
|
+
self.log.info("Command execution result:\n%s", result_str)
|
781
|
+
return result
|
782
|
+
|
783
|
+
def _parse_cmd_to_args(self, cmd: str) -> tuple:
|
784
|
+
"""Parse user command to command, subcommand and parameters."""
|
785
|
+
cmd_dict = shlex.split(cmd)
|
786
|
+
if not cmd_dict:
|
787
|
+
raise AirflowException("The provided command is empty.")
|
788
|
+
|
789
|
+
command = cmd_dict[0] if len(cmd_dict) >= 1 else None
|
790
|
+
subcommand = cmd_dict[1] if len(cmd_dict) >= 2 else None
|
791
|
+
parameters = cmd_dict[2:] if len(cmd_dict) >= 3 else None
|
792
|
+
|
793
|
+
return command, subcommand, parameters
|
794
|
+
|
795
|
+
def _merge_cmd_output_result(self, result) -> str:
|
796
|
+
"""Merge output to one string."""
|
797
|
+
result_str = "\n".join(line_dict["content"] for line_dict in result["output"])
|
798
|
+
return result_str
|
@@ -43,6 +43,7 @@ from airflow.providers.google.cloud.links.cloud_memorystore import (
|
|
43
43
|
RedisInstanceListLink,
|
44
44
|
)
|
45
45
|
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
|
46
|
+
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
46
47
|
|
47
48
|
if TYPE_CHECKING:
|
48
49
|
from google.api_core.retry import Retry
|
@@ -112,7 +113,7 @@ class CloudMemorystoreCreateInstanceOperator(GoogleCloudBaseOperator):
|
|
112
113
|
location: str,
|
113
114
|
instance_id: str,
|
114
115
|
instance: dict | Instance,
|
115
|
-
project_id: str
|
116
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
116
117
|
retry: Retry | _MethodDefault = DEFAULT,
|
117
118
|
timeout: float | None = None,
|
118
119
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -198,7 +199,7 @@ class CloudMemorystoreDeleteInstanceOperator(GoogleCloudBaseOperator):
|
|
198
199
|
*,
|
199
200
|
location: str,
|
200
201
|
instance: str,
|
201
|
-
project_id: str
|
202
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
202
203
|
retry: Retry | _MethodDefault = DEFAULT,
|
203
204
|
timeout: float | None = None,
|
204
205
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -283,7 +284,7 @@ class CloudMemorystoreExportInstanceOperator(GoogleCloudBaseOperator):
|
|
283
284
|
location: str,
|
284
285
|
instance: str,
|
285
286
|
output_config: dict | OutputConfig,
|
286
|
-
project_id: str
|
287
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
287
288
|
retry: Retry | _MethodDefault = DEFAULT,
|
288
289
|
timeout: float | None = None,
|
289
290
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -376,7 +377,7 @@ class CloudMemorystoreFailoverInstanceOperator(GoogleCloudBaseOperator):
|
|
376
377
|
location: str,
|
377
378
|
instance: str,
|
378
379
|
data_protection_mode: FailoverInstanceRequest.DataProtectionMode,
|
379
|
-
project_id: str
|
380
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
380
381
|
retry: Retry | _MethodDefault = DEFAULT,
|
381
382
|
timeout: float | None = None,
|
382
383
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -462,7 +463,7 @@ class CloudMemorystoreGetInstanceOperator(GoogleCloudBaseOperator):
|
|
462
463
|
*,
|
463
464
|
location: str,
|
464
465
|
instance: str,
|
465
|
-
project_id: str
|
466
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
466
467
|
retry: Retry | _MethodDefault = DEFAULT,
|
467
468
|
timeout: float | None = None,
|
468
469
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -556,7 +557,7 @@ class CloudMemorystoreImportOperator(GoogleCloudBaseOperator):
|
|
556
557
|
location: str,
|
557
558
|
instance: str,
|
558
559
|
input_config: dict | InputConfig,
|
559
|
-
project_id: str
|
560
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
560
561
|
retry: Retry | _MethodDefault = DEFAULT,
|
561
562
|
timeout: float | None = None,
|
562
563
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -646,7 +647,7 @@ class CloudMemorystoreListInstancesOperator(GoogleCloudBaseOperator):
|
|
646
647
|
*,
|
647
648
|
location: str,
|
648
649
|
page_size: int,
|
649
|
-
project_id: str
|
650
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
650
651
|
retry: Retry | _MethodDefault = DEFAULT,
|
651
652
|
timeout: float | None = None,
|
652
653
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -749,7 +750,7 @@ class CloudMemorystoreUpdateInstanceOperator(GoogleCloudBaseOperator):
|
|
749
750
|
instance: dict | Instance,
|
750
751
|
location: str | None = None,
|
751
752
|
instance_id: str | None = None,
|
752
|
-
project_id: str
|
753
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
753
754
|
retry: Retry | _MethodDefault = DEFAULT,
|
754
755
|
timeout: float | None = None,
|
755
756
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -842,7 +843,7 @@ class CloudMemorystoreScaleInstanceOperator(GoogleCloudBaseOperator):
|
|
842
843
|
memory_size_gb: int,
|
843
844
|
location: str | None = None,
|
844
845
|
instance_id: str | None = None,
|
845
|
-
project_id: str
|
846
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
846
847
|
retry: Retry | _MethodDefault = DEFAULT,
|
847
848
|
timeout: float | None = None,
|
848
849
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -954,7 +955,7 @@ class CloudMemorystoreCreateInstanceAndImportOperator(GoogleCloudBaseOperator):
|
|
954
955
|
instance_id: str,
|
955
956
|
instance: dict | Instance,
|
956
957
|
input_config: dict | InputConfig,
|
957
|
-
project_id: str
|
958
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
958
959
|
retry: Retry | _MethodDefault = DEFAULT,
|
959
960
|
timeout: float | None = None,
|
960
961
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -1061,7 +1062,7 @@ class CloudMemorystoreExportAndDeleteInstanceOperator(GoogleCloudBaseOperator):
|
|
1061
1062
|
location: str,
|
1062
1063
|
instance: str,
|
1063
1064
|
output_config: dict | OutputConfig,
|
1064
|
-
project_id: str
|
1065
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1065
1066
|
retry: Retry | _MethodDefault = DEFAULT,
|
1066
1067
|
timeout: float | None = None,
|
1067
1068
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -1243,7 +1244,7 @@ class CloudMemorystoreMemcachedCreateInstanceOperator(GoogleCloudBaseOperator):
|
|
1243
1244
|
location: str,
|
1244
1245
|
instance_id: str,
|
1245
1246
|
instance: dict | cloud_memcache.Instance,
|
1246
|
-
project_id: str
|
1247
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1247
1248
|
retry: Retry | _MethodDefault = DEFAULT,
|
1248
1249
|
timeout: float | None = None,
|
1249
1250
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -1316,7 +1317,7 @@ class CloudMemorystoreMemcachedDeleteInstanceOperator(GoogleCloudBaseOperator):
|
|
1316
1317
|
self,
|
1317
1318
|
location: str,
|
1318
1319
|
instance: str,
|
1319
|
-
project_id: str
|
1320
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1320
1321
|
retry: Retry | _MethodDefault = DEFAULT,
|
1321
1322
|
timeout: float | None = None,
|
1322
1323
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -1390,7 +1391,7 @@ class CloudMemorystoreMemcachedGetInstanceOperator(GoogleCloudBaseOperator):
|
|
1390
1391
|
*,
|
1391
1392
|
location: str,
|
1392
1393
|
instance: str,
|
1393
|
-
project_id: str
|
1394
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1394
1395
|
retry: Retry | _MethodDefault = DEFAULT,
|
1395
1396
|
timeout: float | None = None,
|
1396
1397
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -1474,7 +1475,7 @@ class CloudMemorystoreMemcachedListInstancesOperator(GoogleCloudBaseOperator):
|
|
1474
1475
|
self,
|
1475
1476
|
*,
|
1476
1477
|
location: str,
|
1477
|
-
project_id: str
|
1478
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1478
1479
|
retry: Retry | _MethodDefault = DEFAULT,
|
1479
1480
|
timeout: float | None = None,
|
1480
1481
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -1572,7 +1573,7 @@ class CloudMemorystoreMemcachedUpdateInstanceOperator(GoogleCloudBaseOperator):
|
|
1572
1573
|
instance: dict | cloud_memcache.Instance,
|
1573
1574
|
location: str | None = None,
|
1574
1575
|
instance_id: str | None = None,
|
1575
|
-
project_id: str
|
1576
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1576
1577
|
retry: Retry | _MethodDefault = DEFAULT,
|
1577
1578
|
timeout: float | None = None,
|
1578
1579
|
metadata: Sequence[tuple[str, str]] = (),
|
@@ -19,6 +19,7 @@
|
|
19
19
|
|
20
20
|
from __future__ import annotations
|
21
21
|
|
22
|
+
from functools import cached_property
|
22
23
|
from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence
|
23
24
|
|
24
25
|
from googleapiclient.errors import HttpError
|
@@ -31,7 +32,7 @@ from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatab
|
|
31
32
|
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
|
32
33
|
from airflow.providers.google.cloud.triggers.cloud_sql import CloudSQLExportTrigger
|
33
34
|
from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
|
34
|
-
from airflow.providers.google.common.hooks.base_google import get_field
|
35
|
+
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, get_field
|
35
36
|
from airflow.providers.google.common.links.storage import FileDetailsLink
|
36
37
|
|
37
38
|
if TYPE_CHECKING:
|
@@ -244,7 +245,7 @@ class CloudSQLBaseOperator(GoogleCloudBaseOperator):
|
|
244
245
|
self,
|
245
246
|
*,
|
246
247
|
instance: str,
|
247
|
-
project_id: str
|
248
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
248
249
|
gcp_conn_id: str = "google_cloud_default",
|
249
250
|
api_version: str = "v1beta4",
|
250
251
|
impersonation_chain: str | Sequence[str] | None = None,
|
@@ -337,7 +338,7 @@ class CloudSQLCreateInstanceOperator(CloudSQLBaseOperator):
|
|
337
338
|
*,
|
338
339
|
body: dict,
|
339
340
|
instance: str,
|
340
|
-
project_id: str
|
341
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
341
342
|
gcp_conn_id: str = "google_cloud_default",
|
342
343
|
api_version: str = "v1beta4",
|
343
344
|
validate_body: bool = True,
|
@@ -440,7 +441,7 @@ class CloudSQLInstancePatchOperator(CloudSQLBaseOperator):
|
|
440
441
|
*,
|
441
442
|
body: dict,
|
442
443
|
instance: str,
|
443
|
-
project_id: str
|
444
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
444
445
|
gcp_conn_id: str = "google_cloud_default",
|
445
446
|
api_version: str = "v1beta4",
|
446
447
|
impersonation_chain: str | Sequence[str] | None = None,
|
@@ -572,7 +573,7 @@ class CloudSQLCloneInstanceOperator(CloudSQLBaseOperator):
|
|
572
573
|
instance: str,
|
573
574
|
destination_instance_name: str,
|
574
575
|
clone_context: dict | None = None,
|
575
|
-
project_id: str
|
576
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
576
577
|
gcp_conn_id: str = "google_cloud_default",
|
577
578
|
api_version: str = "v1beta4",
|
578
579
|
impersonation_chain: str | Sequence[str] | None = None,
|
@@ -663,7 +664,7 @@ class CloudSQLCreateInstanceDatabaseOperator(CloudSQLBaseOperator):
|
|
663
664
|
*,
|
664
665
|
instance: str,
|
665
666
|
body: dict,
|
666
|
-
project_id: str
|
667
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
667
668
|
gcp_conn_id: str = "google_cloud_default",
|
668
669
|
api_version: str = "v1beta4",
|
669
670
|
validate_body: bool = True,
|
@@ -771,7 +772,7 @@ class CloudSQLPatchInstanceDatabaseOperator(CloudSQLBaseOperator):
|
|
771
772
|
instance: str,
|
772
773
|
database: str,
|
773
774
|
body: dict,
|
774
|
-
project_id: str
|
775
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
775
776
|
gcp_conn_id: str = "google_cloud_default",
|
776
777
|
api_version: str = "v1beta4",
|
777
778
|
validate_body: bool = True,
|
@@ -867,7 +868,7 @@ class CloudSQLDeleteInstanceDatabaseOperator(CloudSQLBaseOperator):
|
|
867
868
|
*,
|
868
869
|
instance: str,
|
869
870
|
database: str,
|
870
|
-
project_id: str
|
871
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
871
872
|
gcp_conn_id: str = "google_cloud_default",
|
872
873
|
api_version: str = "v1beta4",
|
873
874
|
impersonation_chain: str | Sequence[str] | None = None,
|
@@ -957,7 +958,7 @@ class CloudSQLExportInstanceOperator(CloudSQLBaseOperator):
|
|
957
958
|
*,
|
958
959
|
instance: str,
|
959
960
|
body: dict,
|
960
|
-
project_id: str
|
961
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
961
962
|
gcp_conn_id: str = "google_cloud_default",
|
962
963
|
api_version: str = "v1beta4",
|
963
964
|
validate_body: bool = True,
|
@@ -1104,7 +1105,7 @@ class CloudSQLImportInstanceOperator(CloudSQLBaseOperator):
|
|
1104
1105
|
*,
|
1105
1106
|
instance: str,
|
1106
1107
|
body: dict,
|
1107
|
-
project_id: str
|
1108
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
1108
1109
|
gcp_conn_id: str = "google_cloud_default",
|
1109
1110
|
api_version: str = "v1beta4",
|
1110
1111
|
validate_body: bool = True,
|
@@ -1181,10 +1182,35 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
|
|
1181
1182
|
details on how to define ``gcpcloudsql://`` connection.
|
1182
1183
|
:param sql_proxy_binary_path: (optional) Path to the cloud-sql-proxy binary.
|
1183
1184
|
is not specified or the binary is not present, it is automatically downloaded.
|
1185
|
+
:param ssl_cert: (optional) Path to client certificate to authenticate when SSL is used. Overrides the
|
1186
|
+
connection field ``sslcert``.
|
1187
|
+
:param ssl_key: (optional) Path to client private key to authenticate when SSL is used. Overrides the
|
1188
|
+
connection field ``sslkey``.
|
1189
|
+
:param ssl_root_cert: (optional) Path to server's certificate to authenticate when SSL is used. Overrides
|
1190
|
+
the connection field ``sslrootcert``.
|
1191
|
+
:param ssl_secret_id: (optional) ID of the secret in Google Cloud Secret Manager that stores SSL
|
1192
|
+
certificate in the format below:
|
1193
|
+
|
1194
|
+
{'sslcert': '',
|
1195
|
+
'sslkey': '',
|
1196
|
+
'sslrootcert': ''}
|
1197
|
+
|
1198
|
+
Overrides the connection fields ``sslcert``, ``sslkey``, ``sslrootcert``.
|
1199
|
+
Note that according to the Secret Manager requirements, the mentioned dict should be saved as a
|
1200
|
+
string, and encoded with base64.
|
1201
|
+
Note that this parameter is incompatible with parameters ``ssl_cert``, ``ssl_key``, ``ssl_root_cert``.
|
1184
1202
|
"""
|
1185
1203
|
|
1186
1204
|
# [START gcp_sql_query_template_fields]
|
1187
|
-
template_fields: Sequence[str] = (
|
1205
|
+
template_fields: Sequence[str] = (
|
1206
|
+
"sql",
|
1207
|
+
"gcp_cloudsql_conn_id",
|
1208
|
+
"gcp_conn_id",
|
1209
|
+
"ssl_server_cert",
|
1210
|
+
"ssl_client_cert",
|
1211
|
+
"ssl_client_key",
|
1212
|
+
"ssl_secret_id",
|
1213
|
+
)
|
1188
1214
|
template_ext: Sequence[str] = (".sql",)
|
1189
1215
|
template_fields_renderers = {"sql": "sql"}
|
1190
1216
|
# [END gcp_sql_query_template_fields]
|
@@ -1199,6 +1225,10 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
|
|
1199
1225
|
gcp_conn_id: str = "google_cloud_default",
|
1200
1226
|
gcp_cloudsql_conn_id: str = "google_cloud_sql_default",
|
1201
1227
|
sql_proxy_binary_path: str | None = None,
|
1228
|
+
ssl_server_cert: str | None = None,
|
1229
|
+
ssl_client_cert: str | None = None,
|
1230
|
+
ssl_client_key: str | None = None,
|
1231
|
+
ssl_secret_id: str | None = None,
|
1202
1232
|
**kwargs,
|
1203
1233
|
) -> None:
|
1204
1234
|
super().__init__(**kwargs)
|
@@ -1209,6 +1239,10 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
|
|
1209
1239
|
self.parameters = parameters
|
1210
1240
|
self.gcp_connection: Connection | None = None
|
1211
1241
|
self.sql_proxy_binary_path = sql_proxy_binary_path
|
1242
|
+
self.ssl_server_cert = ssl_server_cert
|
1243
|
+
self.ssl_client_cert = ssl_client_cert
|
1244
|
+
self.ssl_client_key = ssl_client_key
|
1245
|
+
self.ssl_secret_id = ssl_secret_id
|
1212
1246
|
|
1213
1247
|
def _execute_query(self, hook: CloudSQLDatabaseHook, database_hook: PostgresHook | MySqlHook) -> None:
|
1214
1248
|
cloud_sql_proxy_runner = None
|
@@ -1228,12 +1262,8 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
|
|
1228
1262
|
|
1229
1263
|
def execute(self, context: Context):
|
1230
1264
|
self.gcp_connection = BaseHook.get_connection(self.gcp_conn_id)
|
1231
|
-
|
1232
|
-
|
1233
|
-
gcp_conn_id=self.gcp_conn_id,
|
1234
|
-
default_gcp_project_id=get_field(self.gcp_connection.extra_dejson, "project"),
|
1235
|
-
sql_proxy_binary_path=self.sql_proxy_binary_path,
|
1236
|
-
)
|
1265
|
+
|
1266
|
+
hook = self.hook
|
1237
1267
|
hook.validate_ssl_certs()
|
1238
1268
|
connection = hook.create_connection()
|
1239
1269
|
hook.validate_socket_path_length()
|
@@ -1242,3 +1272,16 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
|
|
1242
1272
|
self._execute_query(hook, database_hook)
|
1243
1273
|
finally:
|
1244
1274
|
hook.cleanup_database_hook()
|
1275
|
+
|
1276
|
+
@cached_property
|
1277
|
+
def hook(self):
|
1278
|
+
return CloudSQLDatabaseHook(
|
1279
|
+
gcp_cloudsql_conn_id=self.gcp_cloudsql_conn_id,
|
1280
|
+
gcp_conn_id=self.gcp_conn_id,
|
1281
|
+
default_gcp_project_id=get_field(self.gcp_connection.extra_dejson, "project"),
|
1282
|
+
sql_proxy_binary_path=self.sql_proxy_binary_path,
|
1283
|
+
ssl_root_cert=self.ssl_server_cert,
|
1284
|
+
ssl_cert=self.ssl_client_cert,
|
1285
|
+
ssl_key=self.ssl_client_key,
|
1286
|
+
ssl_secret_id=self.ssl_secret_id,
|
1287
|
+
)
|