github.com/sequix/cortex@v1.1.6/pkg/chunk/gcp/table_client.go (about)

     1  package gcp
     2  
     3  import (
     4  	"context"
     5  
     6  	"google.golang.org/grpc/codes"
     7  
     8  	"cloud.google.com/go/bigtable"
     9  	"google.golang.org/grpc/status"
    10  
    11  	"github.com/sequix/cortex/pkg/chunk"
    12  	"github.com/pkg/errors"
    13  )
    14  
    15  type tableClient struct {
    16  	cfg    Config
    17  	client *bigtable.AdminClient
    18  }
    19  
    20  // NewTableClient returns a new TableClient.
    21  func NewTableClient(ctx context.Context, cfg Config) (chunk.TableClient, error) {
    22  	opts := toOptions(cfg.GRPCClientConfig.DialOption(bigtableInstrumentation()))
    23  	client, err := bigtable.NewAdminClient(ctx, cfg.Project, cfg.Instance, opts...)
    24  	if err != nil {
    25  		return nil, err
    26  	}
    27  	return &tableClient{
    28  		cfg:    cfg,
    29  		client: client,
    30  	}, nil
    31  }
    32  
    33  func (c *tableClient) ListTables(ctx context.Context) ([]string, error) {
    34  	tables, err := c.client.Tables(ctx)
    35  	if err != nil {
    36  		return nil, errors.Wrap(err, "client.Tables")
    37  	}
    38  
    39  	// Check each table has the right column family.  If not, omit it.
    40  	output := make([]string, 0, len(tables))
    41  	for _, table := range tables {
    42  		info, err := c.client.TableInfo(ctx, table)
    43  		if err != nil {
    44  			return nil, errors.Wrap(err, "client.TableInfo")
    45  		}
    46  
    47  		if hasColumnFamily(info.FamilyInfos) {
    48  			output = append(output, table)
    49  		}
    50  	}
    51  
    52  	return output, nil
    53  }
    54  
    55  func hasColumnFamily(infos []bigtable.FamilyInfo) bool {
    56  	for _, family := range infos {
    57  		if family.Name == columnFamily {
    58  			return true
    59  		}
    60  	}
    61  	return false
    62  }
    63  
    64  func (c *tableClient) CreateTable(ctx context.Context, desc chunk.TableDesc) error {
    65  	if err := c.client.CreateTable(ctx, desc.Name); err != nil {
    66  		if !alreadyExistsError(err) {
    67  			return errors.Wrap(err, "client.CreateTable")
    68  		}
    69  	}
    70  
    71  	if err := c.client.CreateColumnFamily(ctx, desc.Name, columnFamily); err != nil {
    72  		if !alreadyExistsError(err) {
    73  			return errors.Wrap(err, "client.CreateColumnFamily")
    74  		}
    75  	}
    76  
    77  	return nil
    78  }
    79  
    80  func alreadyExistsError(err error) bool {
    81  	serr, ok := status.FromError(err)
    82  	return ok && serr.Code() == codes.AlreadyExists
    83  }
    84  
    85  func (c *tableClient) DeleteTable(ctx context.Context, name string) error {
    86  	if err := c.client.DeleteTable(ctx, name); err != nil {
    87  		return errors.Wrap(err, "client.DeleteTable")
    88  	}
    89  
    90  	return nil
    91  }
    92  
    93  func (c *tableClient) DescribeTable(ctx context.Context, name string) (desc chunk.TableDesc, isActive bool, err error) {
    94  	return chunk.TableDesc{
    95  		Name: name,
    96  	}, true, nil
    97  }
    98  
    99  func (c *tableClient) UpdateTable(ctx context.Context, current, expected chunk.TableDesc) error {
   100  	return nil
   101  }