github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/shipper/index/querier.go (about)

     1  package index
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/grafana/dskit/tenant"
     8  	"go.etcd.io/bbolt"
     9  
    10  	shipper_index "github.com/grafana/loki/pkg/storage/stores/indexshipper/index"
    11  	"github.com/grafana/loki/pkg/storage/stores/series/index"
    12  	"github.com/grafana/loki/pkg/storage/stores/shipper/index/indexfile"
    13  	"github.com/grafana/loki/pkg/storage/stores/shipper/util"
    14  )
    15  
    16  type Writer interface {
    17  	ForEach(ctx context.Context, tableName string, callback func(boltdb *bbolt.DB) error) error
    18  }
    19  
    20  type Querier interface {
    21  	QueryPages(ctx context.Context, queries []index.Query, callback index.QueryPagesCallback) error
    22  }
    23  
    24  type querier struct {
    25  	writer       Writer
    26  	indexShipper Shipper
    27  }
    28  
    29  func NewQuerier(writer Writer, indexShipper Shipper) Querier {
    30  	return &querier{
    31  		writer:       writer,
    32  		indexShipper: indexShipper,
    33  	}
    34  }
    35  
    36  // QueryPages queries both the writer and indexShipper for the given queries.
    37  func (q *querier) QueryPages(ctx context.Context, queries []index.Query, callback index.QueryPagesCallback) error {
    38  	userID, err := tenant.TenantID(ctx)
    39  	if err != nil {
    40  		return err
    41  	}
    42  
    43  	userIDBytes := util.GetUnsafeBytes(userID)
    44  	queriesByTable := util.QueriesByTable(queries)
    45  	for table, queries := range queriesByTable {
    46  		err := util.DoParallelQueries(ctx, func(ctx context.Context, queries []index.Query, callback index.QueryPagesCallback) error {
    47  			// writer could be nil when running in ReadOnly mode
    48  			if q.writer != nil {
    49  				err := q.writer.ForEach(ctx, table, func(boltdb *bbolt.DB) error {
    50  					return indexfile.QueryBoltDB(ctx, boltdb, userIDBytes, queries, callback)
    51  				})
    52  				if err != nil {
    53  					return err
    54  				}
    55  			}
    56  
    57  			return q.indexShipper.ForEach(ctx, table, userID, func(_ bool, idx shipper_index.Index) error {
    58  				boltdbIndexFile, ok := idx.(*indexfile.IndexFile)
    59  				if !ok {
    60  					return fmt.Errorf("unexpected index type %T", idx)
    61  				}
    62  
    63  				return indexfile.QueryBoltDB(ctx, boltdbIndexFile.GetBoltDB(), userIDBytes, queries, callback)
    64  			})
    65  		}, queries, callback)
    66  		if err != nil {
    67  			return err
    68  		}
    69  	}
    70  
    71  	return nil
    72  }