Skip to content

Commit 6efa1b9

Browse files
turbaszekturbaszek
authored
Add additional Cloud Datastore operators (#10032)
This PR adds more operators for Google Cloud Datastore service. It also adds missing tests and how-to guides.
1 parent d0776cf commit 6efa1b9

10 files changed

Lines changed: 835 additions & 10 deletions

File tree

‎airflow/providers/google/cloud/example_dags/example_datastore.py‎

Lines changed: 103 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,13 @@
2323
"""
2424

2525
import os
26+
from typing import Any, Dict
2627

2728
from airflow import models
2829
from airflow.providers.google.cloud.operators.datastore import (
30+
CloudDatastoreAllocateIdsOperator, CloudDatastoreBeginTransactionOperator, CloudDatastoreCommitOperator,
2931
CloudDatastoreExportEntitiesOperator, CloudDatastoreImportEntitiesOperator,
32+
CloudDatastoreRollbackOperator, CloudDatastoreRunQueryOperator,
3033
)
3134
from airflow.utils import dates
3235

@@ -37,20 +40,118 @@
3740
"example_gcp_datastore",
3841
schedule_interval=None, # Override to match your needs
3942
start_date=dates.days_ago(1),
40-
tags=['example'],
43+
tags=["example"],
4144
) as dag:
45+
# [START how_to_export_task]
4246
export_task = CloudDatastoreExportEntitiesOperator(
4347
task_id="export_task",
4448
bucket=BUCKET,
4549
project_id=GCP_PROJECT_ID,
4650
overwrite_existing=True,
4751
)
52+
# [END how_to_export_task]
4853

54+
# [START how_to_import_task]
4955
import_task = CloudDatastoreImportEntitiesOperator(
5056
task_id="import_task",
5157
bucket="{{ task_instance.xcom_pull('export_task')['response']['outputUrl'].split('/')[2] }}",
5258
file="{{ '/'.join(task_instance.xcom_pull('export_task')['response']['outputUrl'].split('/')[3:]) }}",
53-
project_id=GCP_PROJECT_ID
59+
project_id=GCP_PROJECT_ID,
5460
)
61+
# [END how_to_import_task]
5562

5663
export_task >> import_task
64+
65+
# [START how_to_keys_def]
66+
KEYS = [
67+
{
68+
"partitionId": {"projectId": GCP_PROJECT_ID, "namespaceId": ""},
69+
"path": {"kind": "airflow"},
70+
}
71+
]
72+
# [END how_to_keys_def]
73+
74+
# [START how_to_transaction_def]
75+
TRANSACTION_OPTIONS: Dict[str, Any] = {"readWrite": {}}
76+
# [END how_to_transaction_def]
77+
78+
# [START how_to_commit_def]
79+
COMMIT_BODY = {
80+
"mode": "TRANSACTIONAL",
81+
"mutations": [
82+
{
83+
"insert": {
84+
"key": KEYS[0],
85+
"properties": {"string": {"stringValue": "airflow is awesome!"}},
86+
}
87+
}
88+
],
89+
"transaction": "{{ task_instance.xcom_pull('begin_transaction_commit') }}",
90+
}
91+
# [END how_to_commit_def]
92+
93+
# [START how_to_query_def]
94+
QUERY = {
95+
"partitionId": {"projectId": GCP_PROJECT_ID, "namespaceId": ""},
96+
"readOptions": {
97+
"transaction": "{{ task_instance.xcom_pull('begin_transaction_query') }}"
98+
},
99+
"query": {},
100+
}
101+
# [END how_to_query_def]
102+
103+
with models.DAG(
104+
"example_gcp_datastore_operations",
105+
start_date=dates.days_ago(1),
106+
schedule_interval=None, # Override to match your needs
107+
tags=["example"],
108+
) as dag2:
109+
# [START how_to_allocate_ids]
110+
allocate_ids = CloudDatastoreAllocateIdsOperator(
111+
task_id="allocate_ids", partial_keys=KEYS, project_id=GCP_PROJECT_ID
112+
)
113+
# [END how_to_allocate_ids]
114+
115+
# [START how_to_begin_transaction]
116+
begin_transaction_commit = CloudDatastoreBeginTransactionOperator(
117+
task_id="begin_transaction_commit",
118+
transaction_options=TRANSACTION_OPTIONS,
119+
project_id=GCP_PROJECT_ID,
120+
)
121+
# [END how_to_begin_transaction]
122+
123+
# [START how_to_commit_task]
124+
commit_task = CloudDatastoreCommitOperator(
125+
task_id="commit_task", body=COMMIT_BODY, project_id=GCP_PROJECT_ID
126+
)
127+
# [END how_to_commit_task]
128+
129+
allocate_ids >> begin_transaction_commit >> commit_task
130+
131+
begin_transaction_query = CloudDatastoreBeginTransactionOperator(
132+
task_id="begin_transaction_query",
133+
transaction_options=TRANSACTION_OPTIONS,
134+
project_id=GCP_PROJECT_ID,
135+
)
136+
137+
# [START how_to_run_query]
138+
run_query = CloudDatastoreRunQueryOperator(
139+
task_id="run_query", body=QUERY, project_id=GCP_PROJECT_ID
140+
)
141+
# [END how_to_run_query]
142+
143+
allocate_ids >> begin_transaction_query >> run_query
144+
145+
begin_transaction_to_rollback = CloudDatastoreBeginTransactionOperator(
146+
task_id="begin_transaction_to_rollback",
147+
transaction_options=TRANSACTION_OPTIONS,
148+
project_id=GCP_PROJECT_ID,
149+
)
150+
151+
# [START how_to_rollback_transaction]
152+
rollback_transaction = CloudDatastoreRollbackOperator(
153+
task_id="rollback_transaction",
154+
transaction="{{ task_instance.xcom_pull('begin_transaction_to_rollback') }}",
155+
)
156+
begin_transaction_to_rollback >> rollback_transaction
157+
# [END how_to_rollback_transaction]

‎airflow/providers/google/cloud/hooks/datastore.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ def allocate_ids(self, partial_keys: List, project_id: str) -> List:
100100
return resp['keys']
101101

102102
@GoogleBaseHook.fallback_to_default_project_id
103-
def begin_transaction(self, project_id: str) -> str:
103+
def begin_transaction(self, project_id: str, transaction_options: Dict[str, Any]) -> str:
104104
"""
105105
Begins a new transaction.
106106
@@ -109,14 +109,16 @@ def begin_transaction(self, project_id: str) -> str:
109109
110110
:param project_id: Google Cloud Platform project ID against which to make the request.
111111
:type project_id: str
112+
:param transaction_options: Options for a new transaction.
113+
:type transaction_options: Dict[str, Any]
112114
:return: a transaction handle.
113115
:rtype: str
114116
"""
115117
conn = self.get_conn() # type: Any
116118

117119
resp = (conn # pylint: disable=no-member
118120
.projects()
119-
.beginTransaction(projectId=project_id, body={})
121+
.beginTransaction(projectId=project_id, body={"transactionOptions": transaction_options})
120122
.execute(num_retries=self.num_retries))
121123

122124
return resp['transaction']

0 commit comments

Comments
 (0)