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

     1  package gcp
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"google.golang.org/grpc/codes"
     8  
     9  	"cloud.google.com/go/bigtable"
    10  	"google.golang.org/grpc/status"
    11  
    12  	"github.com/pkg/errors"
    13  
    14  	"github.com/grafana/loki/pkg/storage/config"
    15  	"github.com/grafana/loki/pkg/storage/stores/series/index"
    16  )
    17  
    18  type tableClient struct {
    19  	cfg    Config
    20  	client *bigtable.AdminClient
    21  
    22  	tableInfo       map[string]*bigtable.TableInfo
    23  	tableExpiration time.Time
    24  }
    25  
    26  // NewTableClient returns a new TableClient.
    27  func NewTableClient(ctx context.Context, cfg Config) (index.TableClient, error) {
    28  	dialOpts, err := cfg.GRPCClientConfig.DialOption(bigtableInstrumentation())
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  	client, err := bigtable.NewAdminClient(ctx, cfg.Project, cfg.Instance, toOptions(dialOpts)...)
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  	return &tableClient{
    37  		cfg:    cfg,
    38  		client: client,
    39  
    40  		tableInfo: map[string]*bigtable.TableInfo{},
    41  	}, nil
    42  }
    43  
    44  // ListTables lists all of the correctly specified cortex tables in bigtable
    45  func (c *tableClient) ListTables(ctx context.Context) ([]string, error) {
    46  	tables, err := c.client.Tables(ctx)
    47  	if err != nil {
    48  		return nil, errors.Wrap(err, "client.Tables")
    49  	}
    50  
    51  	if c.tableExpiration.Before(time.Now()) {
    52  		c.tableInfo = map[string]*bigtable.TableInfo{}
    53  		c.tableExpiration = time.Now().Add(c.cfg.TableCacheExpiration)
    54  	}
    55  
    56  	output := make([]string, 0, len(tables))
    57  	for _, table := range tables {
    58  		info, exists := c.tableInfo[table]
    59  		if !c.cfg.TableCacheEnabled || !exists {
    60  			info, err = c.client.TableInfo(ctx, table)
    61  			if err != nil {
    62  				return nil, errors.Wrap(err, "client.TableInfo")
    63  			}
    64  		}
    65  
    66  		// Check each table has the right column family.  If not, omit it.
    67  		if hasColumnFamily(info.FamilyInfos) {
    68  			output = append(output, table)
    69  			c.tableInfo[table] = info
    70  		}
    71  	}
    72  
    73  	return output, nil
    74  }
    75  
    76  func hasColumnFamily(infos []bigtable.FamilyInfo) bool {
    77  	for _, family := range infos {
    78  		if family.Name == columnFamily {
    79  			return true
    80  		}
    81  	}
    82  	return false
    83  }
    84  
    85  func (c *tableClient) CreateTable(ctx context.Context, desc config.TableDesc) error {
    86  	if err := c.client.CreateTable(ctx, desc.Name); err != nil {
    87  		if !alreadyExistsError(err) {
    88  			return errors.Wrap(err, "client.CreateTable")
    89  		}
    90  	}
    91  
    92  	if err := c.client.CreateColumnFamily(ctx, desc.Name, columnFamily); err != nil {
    93  		if !alreadyExistsError(err) {
    94  			return errors.Wrap(err, "client.CreateColumnFamily")
    95  		}
    96  	}
    97  
    98  	return nil
    99  }
   100  
   101  func alreadyExistsError(err error) bool {
   102  	serr, ok := status.FromError(err)
   103  	return ok && serr.Code() == codes.AlreadyExists
   104  }
   105  
   106  func (c *tableClient) DeleteTable(ctx context.Context, name string) error {
   107  	if err := c.client.DeleteTable(ctx, name); err != nil {
   108  		return errors.Wrap(err, "client.DeleteTable")
   109  	}
   110  	delete(c.tableInfo, name)
   111  
   112  	return nil
   113  }
   114  
   115  func (c *tableClient) DescribeTable(ctx context.Context, name string) (desc config.TableDesc, isActive bool, err error) {
   116  	return config.TableDesc{
   117  		Name: name,
   118  	}, true, nil
   119  }
   120  
   121  func (c *tableClient) UpdateTable(ctx context.Context, current, expected config.TableDesc) error {
   122  	return nil
   123  }
   124  
   125  func (c *tableClient) Stop() {
   126  	c.client.Close()
   127  }