From 4565eb7fe730eade294fb3baa85bd255df008bfa Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Thu, 22 Jul 2021 15:33:45 -0700
Subject: [PATCH] feat(bigquery): support multistatement transaction statistics
 in jobs (#4485)

* feat(bigquery): support multistatement transaction statistics in jobs
---
 bigquery/integration_test.go |  8 ++++++++
 bigquery/job.go              | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 25ba131c1f0..09f72d6be8a 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -2435,6 +2435,7 @@ func TestIntegration_Scripting(t *testing.T) {
 	sql := `
 	-- Declare a variable to hold names as an array.
 	DECLARE top_names ARRAY<STRING>;
+	BEGIN TRANSACTION;
 	-- Build an array of the top 100 names from the year 2017.
 	SET top_names = (
 	  SELECT ARRAY_AGG(name ORDER BY number DESC LIMIT 100)
@@ -2449,6 +2450,7 @@ func TestIntegration_Scripting(t *testing.T) {
 	  SELECT word
 	  FROM ` + "`bigquery-public-data`" + `.samples.shakespeare
 	);
+	COMMIT TRANSACTION;
 	`
 	q := client.Query(sql)
 	job, err := q.Run(ctx)
@@ -2506,6 +2508,12 @@ func TestIntegration_Scripting(t *testing.T) {
 		if cStatus.Statistics.ScriptStatistics.EvaluationKind == "" {
 			t.Errorf("child job %q didn't indicate evaluation kind", cj.ID())
 		}
+		if cStatus.Statistics.TransactionInfo == nil {
+			t.Errorf("child job %q didn't have transaction info present", cj.ID())
+		}
+		if cStatus.Statistics.TransactionInfo.TransactionID == "" {
+			t.Errorf("child job %q didn't have transactionID present", cj.ID())
+		}
 	}
 
 }
diff --git a/bigquery/job.go b/bigquery/job.go
index 0499da7d129..657d338e4b4 100644
--- a/bigquery/job.go
+++ b/bigquery/job.go
@@ -372,6 +372,9 @@ type JobStatistics struct {
 
 	// ReservationUsage attributes slot consumption to reservations.
 	ReservationUsage []*ReservationUsage
+
+	// TransactionInfo indicates the transaction ID associated with the job, if any.
+	TransactionInfo *TransactionInfo
 }
 
 // Statistics is one of ExtractStatistics, LoadStatistics or QueryStatistics.
@@ -880,6 +883,7 @@ func (j *Job) setStatistics(s *bq.JobStatistics, c *Client) {
 		ParentJobID:         s.ParentJobId,
 		ScriptStatistics:    bqToScriptStatistics(s.ScriptStatistics),
 		ReservationUsage:    bqToReservationUsage(s.ReservationUsage),
+		TransactionInfo:     bqToTransactionInfo(s.TransactionInfo),
 	}
 	switch {
 	case s.Extract != nil:
@@ -983,3 +987,18 @@ func timelineFromProto(timeline []*bq.QueryTimelineSample) []*QueryTimelineSampl
 	}
 	return res
 }
+
+// TransactionInfo contains information about a multi-statement transaction that may have associated with a job.
+type TransactionInfo struct {
+	// TransactionID is the system-generated identifier for the transaction.
+	TransactionID string
+}
+
+func bqToTransactionInfo(in *bq.TransactionInfo) *TransactionInfo {
+	if in == nil {
+		return nil
+	}
+	return &TransactionInfo{
+		TransactionID: in.TransactionId,
+	}
+}