go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/bqlog/writer.go (about)

     1  // Copyright 2021 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package bqlog
    16  
    17  import (
    18  	"context"
    19  
    20  	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
    21  	gax "github.com/googleapis/gax-go/v2"
    22  	"google.golang.org/grpc/metadata"
    23  )
    24  
    25  // BigQueryWriter is a subset of BigQueryWriteClient API used by Bundler.
    26  //
    27  // Use storage.NewBigQueryWriteClient to construct a production instance.
    28  type BigQueryWriter interface {
    29  	AppendRows(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error)
    30  	Close() error
    31  }
    32  
    33  // FakeBigQueryWriter is a fake for BigQueryWriter.
    34  //
    35  // Calls given callbacks if they are non-nil. Useful in tests.
    36  type FakeBigQueryWriter struct {
    37  	Send func(*storagepb.AppendRowsRequest) error
    38  	Recv func() (*storagepb.AppendRowsResponse, error)
    39  }
    40  
    41  func (f *FakeBigQueryWriter) AppendRows(ctx context.Context, opts ...gax.CallOption) (storagepb.BigQueryWrite_AppendRowsClient, error) {
    42  	return &noopAppendRowsClient{send: f.Send, recv: f.Recv}, nil
    43  }
    44  
    45  func (*FakeBigQueryWriter) Close() error {
    46  	return nil
    47  }
    48  
    49  type noopAppendRowsClient struct {
    50  	send func(*storagepb.AppendRowsRequest) error
    51  	recv func() (*storagepb.AppendRowsResponse, error)
    52  }
    53  
    54  func (c *noopAppendRowsClient) Send(r *storagepb.AppendRowsRequest) error {
    55  	if c.send != nil {
    56  		return c.send(r)
    57  	}
    58  	return nil
    59  }
    60  
    61  func (c *noopAppendRowsClient) Recv() (*storagepb.AppendRowsResponse, error) {
    62  	if c.recv != nil {
    63  		return c.recv()
    64  	}
    65  	return &storagepb.AppendRowsResponse{}, nil
    66  }
    67  
    68  func (*noopAppendRowsClient) Header() (metadata.MD, error) { panic("unused") }
    69  func (*noopAppendRowsClient) Trailer() metadata.MD         { panic("unused") }
    70  func (*noopAppendRowsClient) CloseSend() error             { return nil }
    71  func (*noopAppendRowsClient) Context() context.Context     { panic("unused") }
    72  func (*noopAppendRowsClient) SendMsg(m any) error          { panic("unused") }
    73  func (*noopAppendRowsClient) RecvMsg(m any) error          { panic("unused") }