Skip to content

Commit 7f7b1a8

Browse files
plamutplamut
authored
feat: add support for transaction statistics (googleapis#849)
* feat: add support for transaction statistics * Hoist transaction_info into base job class * Add versionadded directive to new property and class * Include new class in docs reference
1 parent 9c6614f commit 7f7b1a8

8 files changed

Lines changed: 112 additions & 0 deletions

File tree

‎docs/reference.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ Job-Related Types
6868
job.SourceFormat
6969
job.WriteDisposition
7070
job.SchemaUpdateOption
71+
job.TransactionInfo
7172

7273

7374
Dataset

‎google/cloud/bigquery/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
from google.cloud.bigquery.job import ScriptOptions
7171
from google.cloud.bigquery.job import SourceFormat
7272
from google.cloud.bigquery.job import UnknownJob
73+
from google.cloud.bigquery.job import TransactionInfo
7374
from google.cloud.bigquery.job import WriteDisposition
7475
from google.cloud.bigquery.model import Model
7576
from google.cloud.bigquery.model import ModelReference
@@ -149,6 +150,7 @@
149150
"GoogleSheetsOptions",
150151
"ParquetOptions",
151152
"ScriptOptions",
153+
"TransactionInfo",
152154
"DEFAULT_RETRY",
153155
# Enum Constants
154156
"enums",

‎google/cloud/bigquery/job/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from google.cloud.bigquery.job.base import ReservationUsage
2323
from google.cloud.bigquery.job.base import ScriptStatistics
2424
from google.cloud.bigquery.job.base import ScriptStackFrame
25+
from google.cloud.bigquery.job.base import TransactionInfo
2526
from google.cloud.bigquery.job.base import UnknownJob
2627
from google.cloud.bigquery.job.copy_ import CopyJob
2728
from google.cloud.bigquery.job.copy_ import CopyJobConfig
@@ -81,5 +82,6 @@
8182
"QueryPriority",
8283
"SchemaUpdateOption",
8384
"SourceFormat",
85+
"TransactionInfo",
8486
"WriteDisposition",
8587
]

‎google/cloud/bigquery/job/base.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import http
2020
import threading
2121
import typing
22+
from typing import Dict, Optional
2223

2324
from google.api_core import exceptions
2425
import google.api_core.future.polling
@@ -88,6 +89,22 @@ def _error_result_to_exception(error_result):
8889
)
8990

9091

92+
class TransactionInfo(typing.NamedTuple):
93+
"""[Alpha] Information of a multi-statement transaction.
94+
95+
https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#TransactionInfo
96+
97+
.. versionadded:: 2.24.0
98+
"""
99+
100+
transaction_id: str
101+
"""Output only. ID of the transaction."""
102+
103+
@classmethod
104+
def from_api_repr(cls, transaction_info: Dict[str, str]) -> "TransactionInfo":
105+
return cls(transaction_info["transactionId"])
106+
107+
91108
class _JobReference(object):
92109
"""A reference to a job.
93110
@@ -336,6 +353,18 @@ def reservation_usage(self):
336353
for usage in usage_stats_raw
337354
]
338355

356+
@property
357+
def transaction_info(self) -> Optional[TransactionInfo]:
358+
"""Information of the multi-statement transaction if this job is part of one.
359+
360+
.. versionadded:: 2.24.0
361+
"""
362+
info = self._properties.get("statistics", {}).get("transactionInfo")
363+
if info is None:
364+
return None
365+
else:
366+
return TransactionInfo.from_api_repr(info)
367+
339368
@property
340369
def error_result(self):
341370
"""Error information about the job as a whole.

‎tests/system/test_client.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1557,6 +1557,40 @@ def test_dml_statistics(self):
15571557
assert query_job.dml_stats.updated_row_count == 0
15581558
assert query_job.dml_stats.deleted_row_count == 3
15591559

1560+
def test_transaction_info(self):
1561+
table_schema = (
1562+
bigquery.SchemaField("foo", "STRING"),
1563+
bigquery.SchemaField("bar", "INTEGER"),
1564+
)
1565+
1566+
dataset_id = _make_dataset_id("bq_system_test")
1567+
self.temp_dataset(dataset_id)
1568+
table_id = f"{Config.CLIENT.project}.{dataset_id}.test_dml_statistics"
1569+
1570+
# Create the table before loading so that the column order is deterministic.
1571+
table = helpers.retry_403(Config.CLIENT.create_table)(
1572+
Table(table_id, schema=table_schema)
1573+
)
1574+
self.to_delete.insert(0, table)
1575+
1576+
# Insert a few rows and check the stats.
1577+
sql = f"""
1578+
BEGIN TRANSACTION;
1579+
INSERT INTO `{table_id}`
1580+
VALUES ("one", 1), ("two", 2), ("three", 3), ("four", 4);
1581+
1582+
UPDATE `{table_id}`
1583+
SET bar = bar + 1
1584+
WHERE bar > 2;
1585+
COMMIT TRANSACTION;
1586+
"""
1587+
query_job = Config.CLIENT.query(sql)
1588+
query_job.result()
1589+
1590+
# Transaction ID set by the server should be accessible
1591+
assert query_job.transaction_info is not None
1592+
assert query_job.transaction_info.transaction_id != ""
1593+
15601594
def test_dbapi_w_standard_sql_types(self):
15611595
for sql, expected in helpers.STANDARD_SQL_EXAMPLES:
15621596
Config.CURSOR.execute(sql)

‎tests/unit/job/helpers.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ def _verifyInitialReadonlyProperties(self, job):
162162
self.assertIsNone(job.created)
163163
self.assertIsNone(job.started)
164164
self.assertIsNone(job.ended)
165+
self.assertIsNone(job.transaction_info)
165166

166167
# derived from resource['status']
167168
self.assertIsNone(job.error_result)

‎tests/unit/job/test_base.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,20 @@ def test_script_statistics(self):
227227
self.assertEqual(stack_frame.end_column, 14)
228228
self.assertEqual(stack_frame.text, "QUERY TEXT")
229229

230+
def test_transaction_info(self):
231+
from google.cloud.bigquery.job.base import TransactionInfo
232+
233+
client = _make_client(project=self.PROJECT)
234+
job = self._make_one(self.JOB_ID, client)
235+
assert job.transaction_info is None
236+
237+
statistics = job._properties["statistics"] = {}
238+
assert job.transaction_info is None
239+
240+
statistics["transactionInfo"] = {"transactionId": "123-abc-xyz"}
241+
assert isinstance(job.transaction_info, TransactionInfo)
242+
assert job.transaction_info.transaction_id == "123-abc-xyz"
243+
230244
def test_num_child_jobs(self):
231245
client = _make_client(project=self.PROJECT)
232246
job = self._make_one(self.JOB_ID, client)

‎tests/unit/job/test_query.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,18 @@ def _verify_dml_stats_resource_properties(self, job, resource):
128128
else:
129129
assert job.dml_stats is None
130130

131+
def _verify_transaction_info_resource_properties(self, job, resource):
132+
resource_stats = resource.get("statistics", {})
133+
134+
if "transactionInfo" in resource_stats:
135+
resource_transaction_info = resource_stats["transactionInfo"]
136+
job_transaction_info = job.transaction_info
137+
assert job_transaction_info.transaction_id == resource_transaction_info.get(
138+
"transactionId"
139+
)
140+
else:
141+
assert job.transaction_info is None
142+
131143
def _verify_configuration_properties(self, job, configuration):
132144
if "dryRun" in configuration:
133145
self.assertEqual(job.dry_run, configuration["dryRun"])
@@ -137,6 +149,7 @@ def _verify_configuration_properties(self, job, configuration):
137149
def _verifyResourceProperties(self, job, resource):
138150
self._verifyReadonlyResourceProperties(job, resource)
139151
self._verify_dml_stats_resource_properties(job, resource)
152+
self._verify_transaction_info_resource_properties(job, resource)
140153

141154
configuration = resource.get("configuration", {})
142155
self._verify_configuration_properties(job, configuration)
@@ -325,6 +338,22 @@ def test_from_api_repr_with_dml_stats(self):
325338
self.assertIs(job._client, client)
326339
self._verifyResourceProperties(job, RESOURCE)
327340

341+
def test_from_api_repr_with_transaction_info(self):
342+
self._setUpConstants()
343+
client = _make_client(project=self.PROJECT)
344+
RESOURCE = {
345+
"id": self.JOB_ID,
346+
"jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
347+
"configuration": {"query": {"query": self.QUERY}},
348+
"statistics": {"transactionInfo": {"transactionId": "1a2b-3c4d"}},
349+
}
350+
klass = self._get_target_class()
351+
352+
job = klass.from_api_repr(RESOURCE, client=client)
353+
354+
self.assertIs(job._client, client)
355+
self._verifyResourceProperties(job, RESOURCE)
356+
328357
def test_from_api_repr_w_properties(self):
329358
from google.cloud.bigquery.job import CreateDisposition
330359
from google.cloud.bigquery.job import SchemaUpdateOption

0 commit comments

Comments
 (0)