feat(datastore): added newTransaction option for runQuery (#925)

Author: mwaehnerdialpad

datastore/gcloud/aio/datastore/query.py

@@ -282,9 +282,11 @@ class QueryResult:
     query_result_batch_kind = QueryResultBatch
 
     def __init__(self, result_batch: Optional[QueryResultBatch] = None,
-                 explain_metrics: Optional[ExplainMetrics] = None):
+                 explain_metrics: Optional[ExplainMetrics] = None,
+                 transaction: Optional[str] = None):
         self.result_batch = result_batch
         self.explain_metrics = explain_metrics
+        self.transaction = transaction
 
     def __repr__(self) -> str:
         return str(self.to_repr())
@@ -293,27 +295,33 @@ def __eq__(self, other: object) -> bool:
         if not isinstance(other, QueryResult):
             return False
         return (self.result_batch == other.result_batch
-                and self.explain_metrics == other.explain_metrics)
+                and self.explain_metrics == other.explain_metrics
+                and self.transaction == other.transaction)
 
     @classmethod
     def from_repr(cls, data: Dict[str, Any]) -> 'QueryResult':
         result_batch = None
         explain_metrics = None
+        transaction = None
 
         if 'batch' in data:
             result_batch = cls.query_result_batch_kind.from_repr(data['batch'])
         if 'explainMetrics' in data:
             explain_metrics = ExplainMetrics.from_repr(data['explainMetrics'])
+        if 'transaction' in data:
+            transaction = data['transaction']
 
-        return cls(result_batch=result_batch, explain_metrics=explain_metrics)
+        return cls(result_batch=result_batch,
+                   explain_metrics=explain_metrics, transaction=transaction)
 
     def to_repr(self) -> Dict[str, Any]:
-        result = {}
+        result: Dict[str, Any] = {}
         if self.result_batch:
             result['batch'] = self.result_batch.to_repr()
         if self.explain_metrics:
             result['explainMetrics'] = self.explain_metrics.to_repr()
-
+        if self.transaction:
+            result['transaction'] = self.transaction
         return result
 
     def get_explain_metrics(self) -> Optional[ExplainMetrics]:

datastore/tests/integration/smoke_test.py

@@ -135,6 +135,53 @@ async def test_start_transaction_on_lookup(creds: str,
         actual = await ds.lookup([key], session=s)
         assert actual['found'][0].entity.properties == {'animal': 'aardvark'}
 
+        # Clean up test data
+        await ds.delete(key, s)
+
+
+@pytest.mark.asyncio
+async def test_start_transaction_on_query(
+        creds: str, kind: str, project: str,
+) -> None:
+    key = Key(project, [PathElement(kind, name=f'test_record_{uuid.uuid4()}')])
+
+    async with Session() as s:
+        ds = Datastore(project=project, service_file=creds, session=s)
+
+        # Test query with newTransaction parameter
+        property_filter = PropertyFilter(
+            prop='animal',
+            operator=PropertyFilterOperator.EQUAL,
+            value=Value('three-toed sloth'),
+        )
+        query = Query(kind=kind, query_filter=Filter(property_filter))
+
+        # Use newTransaction parameter
+        options = TransactionOptions(ReadWrite())
+        result = await ds.runQuery(query, newTransaction=options, session=s)
+        assert result.transaction is not None and result.transaction
+
+        mutations = [
+            ds.make_mutation(
+                Operation.INSERT, key,
+                properties={'animal': 'three-toed sloth'},
+            ),
+            ds.make_mutation(
+                Operation.UPDATE, key,
+                properties={'animal': 'aardvark'},
+            ),
+        ]
+        await ds.commit(
+            mutations,
+            transaction=result.transaction,
+            session=s)
+
+        actual = await ds.lookup([key], session=s)
+        assert actual['found'][0].entity.properties == {'animal': 'aardvark'}
+
+        # Clean up test data
+        await ds.delete(key, s)
+
 
 @pytest.mark.asyncio
 async def test_transaction(creds: str, kind: str, project: str) -> None: