This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fcf27fcf86c Fix Datastore rollback system test transaction (#73194)
fcf27fcf86c is described below
commit fcf27fcf86c193f9af9fd834ee0fa6d807e888f4
Author: Ulada Zakharava <[email protected]>
AuthorDate: Wed Sep 16 07:14:56 2026 +0200
Fix Datastore rollback system test transaction (#73194)
---
.../google/docs/operators/cloud/datastore.rst | 42 ++++++++++++-------
.../cloud/datastore/example_datastore_commit.py | 16 -------
.../cloud/datastore/example_datastore_query.py | 14 +------
.../cloud/datastore/example_datastore_rollback.py | 49 ++++++++++++++++++----
4 files changed, 68 insertions(+), 53 deletions(-)
diff --git a/providers/google/docs/operators/cloud/datastore.rst
b/providers/google/docs/operators/cloud/datastore.rst
index a3856ea587a..d900d63d37b 100644
--- a/providers/google/docs/operators/cloud/datastore.rst
+++ b/providers/google/docs/operators/cloud/datastore.rst
@@ -88,19 +88,24 @@ Begin transaction
To begin a new transaction use
:class:`~airflow.providers.google.cloud.operators.datastore.CloudDatastoreBeginTransactionOperator`
-.. exampleinclude::
/../../google/tests/system/google/cloud/datastore/example_datastore_commit.py
- :language: python
- :dedent: 4
- :start-after: [START how_to_begin_transaction]
- :end-before: [END how_to_begin_transaction]
+.. code-block:: python
-An example of a transaction options required by the operator:
+ TRANSACTION_OPTIONS = {"readWrite": {}}
-.. exampleinclude::
/../../google/tests/system/google/cloud/datastore/example_datastore_commit.py
- :language: python
- :dedent: 0
- :start-after: [START how_to_transaction_def]
- :end-before: [END how_to_transaction_def]
+ begin_transaction = CloudDatastoreBeginTransactionOperator(
+ task_id="begin_transaction",
+ transaction_options=TRANSACTION_OPTIONS,
+ project_id=PROJECT_ID,
+ )
+
+.. warning::
+
+ Datastore transactions expire after 270 seconds or after 60 seconds of
inactivity.
+ Airflow does not guarantee that a downstream task will start before those
limits.
+ Do not pass a transaction handle between tasks. Run all operations
belonging to a
+ transaction within one task using
+ :class:`~airflow.providers.google.cloud.hooks.datastore.DatastoreHook` or
a Datastore client.
+ See `Datastore transaction limits
<https://cloud.google.com/datastore/docs/concepts/transactions>`__.
.. _howto/operator:CloudDatastoreCommitOperator:
@@ -154,11 +159,16 @@ Roll back transaction
To roll back a transaction
use
:class:`~airflow.providers.google.cloud.operators.datastore.CloudDatastoreRollbackOperator`
-.. exampleinclude::
/../../google/tests/system/google/cloud/datastore/example_datastore_rollback.py
- :language: python
- :dedent: 4
- :start-after: [START how_to_rollback_transaction]
- :end-before: [END how_to_rollback_transaction]
+This low-level operator requires a transaction handle that is still active
when task execution begins.
+Given such a handle, configure the operator as follows:
+
+.. code-block:: python
+
+ rollback_transaction = CloudDatastoreRollbackOperator(
+ task_id="rollback_transaction",
+ transaction=TRANSACTION_ID,
+ project_id=PROJECT_ID,
+ )
.. _howto/operator:CloudDatastoreGetOperationOperator:
diff --git
a/providers/google/tests/system/google/cloud/datastore/example_datastore_commit.py
b/providers/google/tests/system/google/cloud/datastore/example_datastore_commit.py
index ae39f1fe10f..b138caf2d07 100644
---
a/providers/google/tests/system/google/cloud/datastore/example_datastore_commit.py
+++
b/providers/google/tests/system/google/cloud/datastore/example_datastore_commit.py
@@ -23,13 +23,11 @@ from __future__ import annotations
import os
from datetime import datetime
-from typing import Any
from airflow.models.baseoperator import chain
from airflow.models.dag import DAG
from airflow.providers.google.cloud.operators.datastore import (
CloudDatastoreAllocateIdsOperator,
- CloudDatastoreBeginTransactionOperator,
CloudDatastoreCommitOperator,
CloudDatastoreDeleteOperationOperator,
CloudDatastoreExportEntitiesOperator,
@@ -60,11 +58,6 @@ KEYS = [
]
# [END how_to_keys_def]
-# [START how_to_transaction_def]
-TRANSACTION_OPTIONS: dict[str, Any] = {"readWrite": {}}
-# [END how_to_transaction_def]
-
-
with DAG(
DAG_ID,
schedule="@once",
@@ -82,14 +75,6 @@ with DAG(
)
# [END how_to_allocate_ids]
- # [START how_to_begin_transaction]
- begin_transaction_commit = CloudDatastoreBeginTransactionOperator(
- task_id="begin_transaction_commit",
- transaction_options=TRANSACTION_OPTIONS,
- project_id=PROJECT_ID,
- )
- # [END how_to_begin_transaction]
-
# [START how_to_commit_def]
COMMIT_BODY = {
"mode": "TRANSACTIONAL",
@@ -156,7 +141,6 @@ with DAG(
create_bucket,
# TEST BODY
allocate_ids,
- begin_transaction_commit,
commit_task,
export_task,
import_task,
diff --git
a/providers/google/tests/system/google/cloud/datastore/example_datastore_query.py
b/providers/google/tests/system/google/cloud/datastore/example_datastore_query.py
index 7a7f9b7c25a..d69bc90209b 100644
---
a/providers/google/tests/system/google/cloud/datastore/example_datastore_query.py
+++
b/providers/google/tests/system/google/cloud/datastore/example_datastore_query.py
@@ -23,12 +23,10 @@ from __future__ import annotations
import os
from datetime import datetime
-from typing import Any
from airflow.models.dag import DAG
from airflow.providers.google.cloud.operators.datastore import (
CloudDatastoreAllocateIdsOperator,
- CloudDatastoreBeginTransactionOperator,
CloudDatastoreRunQueryOperator,
)
@@ -46,9 +44,6 @@ KEYS = [
}
]
-TRANSACTION_OPTIONS: dict[str, Any] = {"readWrite": {}}
-
-
with DAG(
DAG_ID,
schedule="@once",
@@ -60,16 +55,9 @@ with DAG(
task_id="allocate_ids", partial_keys=KEYS, project_id=PROJECT_ID
)
- begin_transaction_query = CloudDatastoreBeginTransactionOperator(
- task_id="begin_transaction_query",
- transaction_options=TRANSACTION_OPTIONS,
- project_id=PROJECT_ID,
- )
-
# [START how_to_query_def]
QUERY = {
"partitionId": {"projectId": PROJECT_ID, "namespaceId": "query"},
- "readOptions": {"transaction": begin_transaction_query.output},
"query": {},
}
# [END how_to_query_def]
@@ -78,7 +66,7 @@ with DAG(
run_query = CloudDatastoreRunQueryOperator(task_id="run_query",
body=QUERY, project_id=PROJECT_ID)
# [END how_to_run_query]
- allocate_ids >> begin_transaction_query >> run_query
+ allocate_ids >> run_query
from tests_common.test_utils.watcher import watcher
diff --git
a/providers/google/tests/system/google/cloud/datastore/example_datastore_rollback.py
b/providers/google/tests/system/google/cloud/datastore/example_datastore_rollback.py
index 00dfecba9f4..85b81ab2f4c 100644
---
a/providers/google/tests/system/google/cloud/datastore/example_datastore_rollback.py
+++
b/providers/google/tests/system/google/cloud/datastore/example_datastore_rollback.py
@@ -16,16 +16,17 @@
# specific language governing permissions and limitations
# under the License.
"""
-Airflow System Test DAG that verifies Datastore rollback operators.
+Airflow System Test DAG that verifies Datastore transaction operators.
"""
from __future__ import annotations
import os
from datetime import datetime
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
from airflow.models.dag import DAG
+from airflow.providers.google.cloud.hooks.datastore import DatastoreHook
from airflow.providers.google.cloud.operators.datastore import (
CloudDatastoreBeginTransactionOperator,
CloudDatastoreRollbackOperator,
@@ -33,6 +34,9 @@ from airflow.providers.google.cloud.operators.datastore
import (
from system.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or
DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
@@ -41,6 +45,33 @@ DAG_ID = "datastore_rollback"
TRANSACTION_OPTIONS: dict[str, Any] = {"readWrite": {}}
+def begin_transaction_for_rollback(context: Context) -> None:
+ task = cast("CloudDatastoreRollbackOperator", context["task"])
+
+ hook = DatastoreHook(
+ gcp_conn_id=task.gcp_conn_id,
+ impersonation_chain=task.impersonation_chain,
+ )
+
+ task.transaction = hook.begin_transaction(
+ transaction_options=TRANSACTION_OPTIONS,
+ project_id=task.project_id,
+ )
+
+ task.log.info("Created Datastore transaction immediately before rollback")
+
+
+def rollback_after_begin(context: Context, transaction: str) -> None:
+ task = cast("CloudDatastoreBeginTransactionOperator", context["task"])
+
+ hook = DatastoreHook(
+ gcp_conn_id=task.gcp_conn_id,
+ impersonation_chain=task.impersonation_chain,
+ )
+
+ hook.rollback(transaction=transaction, project_id=task.project_id)
+
+
with DAG(
DAG_ID,
schedule="@once",
@@ -48,20 +79,22 @@ with DAG(
catchup=False,
tags=["datastore", "example"],
) as dag:
- begin_transaction_to_rollback = CloudDatastoreBeginTransactionOperator(
- task_id="begin_transaction_to_rollback",
+ begin_transaction = CloudDatastoreBeginTransactionOperator(
+ task_id="begin_transaction",
transaction_options=TRANSACTION_OPTIONS,
project_id=PROJECT_ID,
+ post_execute=rollback_after_begin,
)
- # [START how_to_rollback_transaction]
rollback_transaction = CloudDatastoreRollbackOperator(
task_id="rollback_transaction",
- transaction=begin_transaction_to_rollback.output,
+ transaction="defined_in_begin_transaction_for_rollback",
+ pre_execute=begin_transaction_for_rollback,
+ project_id=PROJECT_ID,
)
- # [END how_to_rollback_transaction]
- begin_transaction_to_rollback >> rollback_transaction
+ # Each task owns its transaction so scheduler delays cannot expire a
handle between tasks.
+ begin_transaction >> rollback_transaction
from tests_common.test_utils.watcher import watcher
