github.com/google/syzkaller@v0.0.0-20251211124644-a066d2bc4b02/pkg/coveragedb/spannerclient/spanner_client.go (about)

     1  // Copyright 2024 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package spannerclient
     5  
     6  import (
     7  	"context"
     8  	"time"
     9  
    10  	"cloud.google.com/go/spanner"
    11  )
    12  
    13  type SpannerClient interface {
    14  	Close()
    15  	Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption) (commitTimestamp time.Time, err error)
    16  	Single() ReadOnlyTransaction
    17  }
    18  
    19  type ReadOnlyTransaction interface {
    20  	Query(ctx context.Context, statement spanner.Statement) RowIterator
    21  }
    22  
    23  type RowIterator interface {
    24  	Next() (Row, error)
    25  	Stop()
    26  }
    27  
    28  type Row interface {
    29  	ToStruct(p interface{}) error
    30  }
    31  
    32  type SpannerClientProxy struct {
    33  	client *spanner.Client
    34  }
    35  
    36  func (proxy *SpannerClientProxy) Close() {
    37  	proxy.client.Close()
    38  }
    39  
    40  func (proxy *SpannerClientProxy) Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption,
    41  ) (commitTimestamp time.Time, err error) {
    42  	return proxy.client.Apply(ctx, ms, opts...)
    43  }
    44  
    45  func (proxy *SpannerClientProxy) Single() ReadOnlyTransaction {
    46  	return &SpannerReadOnlyTransactionProxy{
    47  		readOnlyTransaction: proxy.client.Single(),
    48  	}
    49  }
    50  
    51  type SpannerReadOnlyTransactionProxy struct {
    52  	readOnlyTransaction *spanner.ReadOnlyTransaction
    53  }
    54  
    55  func (proxy *SpannerReadOnlyTransactionProxy) Query(ctx context.Context, statement spanner.Statement) RowIterator {
    56  	return &SpannerRowIteratorProxy{
    57  		rowIterator: proxy.readOnlyTransaction.Query(ctx, statement),
    58  	}
    59  }
    60  
    61  type SpannerRowIteratorProxy struct {
    62  	rowIterator *spanner.RowIterator
    63  }
    64  
    65  func (proxy *SpannerRowIteratorProxy) Next() (Row, error) {
    66  	return proxy.rowIterator.Next()
    67  }
    68  
    69  func (proxy *SpannerRowIteratorProxy) Stop() {
    70  	proxy.rowIterator.Stop()
    71  }
    72  
    73  func NewClient(ctx context.Context, projectID string) (SpannerClient, error) {
    74  	database := "projects/" + projectID + "/instances/syzbot/databases/coverage"
    75  	client, err := spanner.NewClient(ctx, database)
    76  	return &SpannerClientProxy{
    77  		client: client,
    78  	}, err
    79  }