github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/chunk/client/gcp/fixtures.go (about)

     1  package gcp
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  
     8  	"cloud.google.com/go/bigtable"
     9  	"cloud.google.com/go/bigtable/bttest"
    10  	"cloud.google.com/go/storage"
    11  	"github.com/fsouza/fake-gcs-server/fakestorage"
    12  	"google.golang.org/api/option"
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/credentials/insecure"
    15  
    16  	"github.com/grafana/loki/pkg/storage/chunk/client"
    17  	"github.com/grafana/loki/pkg/storage/chunk/client/hedging"
    18  	"github.com/grafana/loki/pkg/storage/chunk/client/testutils"
    19  	"github.com/grafana/loki/pkg/storage/config"
    20  	"github.com/grafana/loki/pkg/storage/stores/series/index"
    21  )
    22  
    23  const (
    24  	proj, instance = "proj", "instance"
    25  )
    26  
    27  type fixture struct {
    28  	btsrv  *bttest.Server
    29  	gcssrv *fakestorage.Server
    30  
    31  	name string
    32  
    33  	gcsObjectClient bool
    34  	columnKeyClient bool
    35  	hashPrefix      bool
    36  }
    37  
    38  func (f *fixture) Name() string {
    39  	return f.name
    40  }
    41  
    42  func (f *fixture) Clients() (
    43  	iClient index.Client, cClient client.Client, tClient index.TableClient,
    44  	schemaConfig config.SchemaConfig, closer io.Closer, err error,
    45  ) {
    46  	f.btsrv, err = bttest.NewServer("localhost:0")
    47  	if err != nil {
    48  		return
    49  	}
    50  
    51  	f.gcssrv = fakestorage.NewServer(nil)
    52  	f.gcssrv.CreateBucket("chunks")
    53  
    54  	conn, err := grpc.Dial(f.btsrv.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    55  	if err != nil {
    56  		return
    57  	}
    58  
    59  	ctx := context.Background()
    60  	adminClient, err := bigtable.NewAdminClient(ctx, proj, instance, option.WithGRPCConn(conn))
    61  	if err != nil {
    62  		return
    63  	}
    64  
    65  	schemaConfig = testutils.DefaultSchemaConfig("gcp-columnkey")
    66  	tClient = &tableClient{
    67  		client: adminClient,
    68  	}
    69  
    70  	c, err := bigtable.NewClient(ctx, proj, instance, option.WithGRPCConn(conn))
    71  	if err != nil {
    72  		return
    73  	}
    74  
    75  	cfg := Config{
    76  		DistributeKeys: f.hashPrefix,
    77  	}
    78  	if f.columnKeyClient {
    79  		iClient = newStorageClientColumnKey(cfg, schemaConfig, c)
    80  	} else {
    81  		iClient = newStorageClientV1(cfg, schemaConfig, c)
    82  	}
    83  
    84  	if f.gcsObjectClient {
    85  		var c *GCSObjectClient
    86  		c, err = newGCSObjectClient(ctx, GCSConfig{
    87  			BucketName: "chunks",
    88  			Insecure:   true,
    89  		}, hedging.Config{}, func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) {
    90  			return f.gcssrv.Client(), nil
    91  		})
    92  		if err != nil {
    93  			return
    94  		}
    95  		cClient = client.NewClient(c, nil, config.SchemaConfig{})
    96  	} else {
    97  		cClient = newBigtableObjectClient(Config{}, schemaConfig, c)
    98  	}
    99  
   100  	closer = testutils.CloserFunc(func() error {
   101  		conn.Close()
   102  		return nil
   103  	})
   104  
   105  	return
   106  }
   107  
   108  // Fixtures for unit testing GCP storage.
   109  var Fixtures = func() []testutils.Fixture {
   110  	fixtures := []testutils.Fixture{}
   111  	for _, gcsObjectClient := range []bool{true, false} {
   112  		for _, columnKeyClient := range []bool{true, false} {
   113  			for _, hashPrefix := range []bool{true, false} {
   114  				fixtures = append(fixtures, &fixture{
   115  					name:            fmt.Sprintf("bigtable-columnkey:%v-gcsObjectClient:%v-hashPrefix:%v", columnKeyClient, gcsObjectClient, hashPrefix),
   116  					columnKeyClient: columnKeyClient,
   117  					gcsObjectClient: gcsObjectClient,
   118  					hashPrefix:      hashPrefix,
   119  				})
   120  			}
   121  		}
   122  	}
   123  	return fixtures
   124  }()