From 5e6c350b2ac94787934380e930af2cb2094fa8f1 Mon Sep 17 00:00:00 2001
From: Aleksandra Bogoslavetc <alexasha470@gmail.com>
Date: Fri, 18 Sep 2020 21:34:51 +0300
Subject: [PATCH] feat(datastore): add opencensus tracing/stats support (#2804)

Co-authored-by: Chris Cotter <cjcotter@google.com>
---
 datastore/client.go      | 19 +++++++++++++++++++
 datastore/transaction.go |  8 +++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/datastore/client.go b/datastore/client.go
index 22d6c95e02e..2792fb5209e 100644
--- a/datastore/client.go
+++ b/datastore/client.go
@@ -20,6 +20,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/internal"
+	"cloud.google.com/go/internal/trace"
 	"cloud.google.com/go/internal/version"
 	gax "github.com/googleapis/gax-go/v2"
 	pb "google.golang.org/genproto/googleapis/datastore/v1"
@@ -50,6 +51,9 @@ func newDatastoreClient(conn grpc.ClientConnInterface, projectID string) pb.Data
 }
 
 func (dc *datastoreClient) Lookup(ctx context.Context, in *pb.LookupRequest, opts ...grpc.CallOption) (res *pb.LookupResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.Lookup")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	err = dc.invoke(ctx, func(ctx context.Context) error {
 		res, err = dc.c.Lookup(ctx, in, opts...)
 		return err
@@ -58,6 +62,9 @@ func (dc *datastoreClient) Lookup(ctx context.Context, in *pb.LookupRequest, opt
 }
 
 func (dc *datastoreClient) RunQuery(ctx context.Context, in *pb.RunQueryRequest, opts ...grpc.CallOption) (res *pb.RunQueryResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.RunQuery")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	err = dc.invoke(ctx, func(ctx context.Context) error {
 		res, err = dc.c.RunQuery(ctx, in, opts...)
 		return err
@@ -66,6 +73,9 @@ func (dc *datastoreClient) RunQuery(ctx context.Context, in *pb.RunQueryRequest,
 }
 
 func (dc *datastoreClient) BeginTransaction(ctx context.Context, in *pb.BeginTransactionRequest, opts ...grpc.CallOption) (res *pb.BeginTransactionResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.BeginTransaction")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	err = dc.invoke(ctx, func(ctx context.Context) error {
 		res, err = dc.c.BeginTransaction(ctx, in, opts...)
 		return err
@@ -74,6 +84,9 @@ func (dc *datastoreClient) BeginTransaction(ctx context.Context, in *pb.BeginTra
 }
 
 func (dc *datastoreClient) Commit(ctx context.Context, in *pb.CommitRequest, opts ...grpc.CallOption) (res *pb.CommitResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.Commit")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	err = dc.invoke(ctx, func(ctx context.Context) error {
 		res, err = dc.c.Commit(ctx, in, opts...)
 		return err
@@ -82,6 +95,9 @@ func (dc *datastoreClient) Commit(ctx context.Context, in *pb.CommitRequest, opt
 }
 
 func (dc *datastoreClient) Rollback(ctx context.Context, in *pb.RollbackRequest, opts ...grpc.CallOption) (res *pb.RollbackResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.Rollback")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	err = dc.invoke(ctx, func(ctx context.Context) error {
 		res, err = dc.c.Rollback(ctx, in, opts...)
 		return err
@@ -90,6 +106,9 @@ func (dc *datastoreClient) Rollback(ctx context.Context, in *pb.RollbackRequest,
 }
 
 func (dc *datastoreClient) AllocateIds(ctx context.Context, in *pb.AllocateIdsRequest, opts ...grpc.CallOption) (res *pb.AllocateIdsResponse, err error) {
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.datastoreClient.AllocateIds")
+	defer func() { trace.EndSpan(ctx, err) }()
+
 	err = dc.invoke(ctx, func(ctx context.Context) error {
 		res, err = dc.c.AllocateIds(ctx, in, opts...)
 		return err
diff --git a/datastore/transaction.go b/datastore/transaction.go
index a90e81b4c30..08420f26ae6 100644
--- a/datastore/transaction.go
+++ b/datastore/transaction.go
@@ -110,13 +110,19 @@ func (c *Client) NewTransaction(ctx context.Context, opts ...TransactionOption)
 	return c.newTransaction(ctx, newTransactionSettings(opts))
 }
 
-func (c *Client) newTransaction(ctx context.Context, s *transactionSettings) (*Transaction, error) {
+func (c *Client) newTransaction(ctx context.Context, s *transactionSettings) (_ *Transaction, err error) {
 	req := &pb.BeginTransactionRequest{ProjectId: c.dataset}
 	if s.readOnly {
+		ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Transaction.ReadOnlyTransaction")
+		defer func() { trace.EndSpan(ctx, err) }()
+
 		req.TransactionOptions = &pb.TransactionOptions{
 			Mode: &pb.TransactionOptions_ReadOnly_{ReadOnly: &pb.TransactionOptions_ReadOnly{}},
 		}
 	} else if s.prevID != nil {
+		ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Transaction.ReadWriteTransaction")
+		defer func() { trace.EndSpan(ctx, err) }()
+
 		req.TransactionOptions = &pb.TransactionOptions{
 			Mode: &pb.TransactionOptions_ReadWrite_{ReadWrite: &pb.TransactionOptions_ReadWrite{
 				PreviousTransaction: s.prevID,