github.com/blong14/gache@v0.0.0-20240124023949-89416fd8bbfa/internal/db/query.go (about)

     1  package db
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  )
     7  
     8  type QueryInstruction int
     9  
    10  const (
    11  	AddTable QueryInstruction = iota
    12  	BatchSetValue
    13  	Count
    14  	GetValue
    15  	GetRange
    16  	Load
    17  	Print
    18  	Range
    19  	SetValue
    20  )
    21  
    22  func (i QueryInstruction) String() string {
    23  	switch i {
    24  	case AddTable:
    25  		return "AddTable"
    26  	case BatchSetValue:
    27  		return "BatchSetValue"
    28  	case Count:
    29  		return "Count"
    30  	case GetValue:
    31  		return "GetValue"
    32  	case GetRange:
    33  		return "GetRange"
    34  	case Load:
    35  		return "Load"
    36  	case Print:
    37  		return "Print"
    38  	case Range:
    39  		return "Range"
    40  	case SetValue:
    41  		return "SetValue"
    42  	default:
    43  		return "unknown"
    44  	}
    45  }
    46  
    47  type QueryHeader struct {
    48  	DataDir   []byte
    49  	TableName []byte
    50  	Opts      *TableOpts
    51  	FileName  []byte
    52  	Inst      QueryInstruction
    53  }
    54  
    55  type QueryStats struct {
    56  	Count uint
    57  }
    58  
    59  type QueryResponse struct {
    60  	Key         []byte
    61  	Value       []byte
    62  	RangeValues [][][]byte
    63  	Stats       QueryStats
    64  	Success     bool
    65  }
    66  
    67  type KeyRange struct {
    68  	Start []byte
    69  	End   []byte
    70  	Limit int
    71  }
    72  
    73  func (kr *KeyRange) String() string {
    74  	return fmt.Sprintf("%s %s", kr.Start, kr.End)
    75  }
    76  
    77  type Query struct {
    78  	ctx      context.Context
    79  	done     chan QueryResponse
    80  	Header   QueryHeader
    81  	KeyRange KeyRange
    82  	Key      []byte
    83  	Value    []byte
    84  	Values   []KeyValue
    85  }
    86  
    87  func NewQuery(ctx context.Context, outbox chan QueryResponse) *Query {
    88  	if outbox == nil {
    89  		outbox = make(chan QueryResponse, 1)
    90  	}
    91  	return &Query{ctx: ctx, done: outbox}
    92  }
    93  
    94  func (m *Query) String() string {
    95  	return fmt.Sprintf(
    96  		"%s %s %s %s %s %s",
    97  		m.Header.FileName, m.Header.TableName,
    98  		m.Header.Inst, m.Key, m.Value, m.KeyRange.String(),
    99  	)
   100  }
   101  
   102  func (m *Query) Done(r QueryResponse) {
   103  	select {
   104  	case <-m.ctx.Done():
   105  	case m.done <- r:
   106  		close(m.done)
   107  	}
   108  }
   109  
   110  func (m *Query) Context() context.Context {
   111  	if m.ctx == nil {
   112  		return context.Background()
   113  	}
   114  	return m.ctx
   115  }
   116  
   117  func (m *Query) GetResponse() *QueryResponse {
   118  	resp := <-m.done
   119  	return &resp
   120  }
   121  
   122  func NewGetValueQuery(ctx context.Context, db []byte, key []byte) (*Query, chan QueryResponse) {
   123  	done := make(chan QueryResponse, 1)
   124  	query := NewQuery(ctx, done)
   125  	query.Header = QueryHeader{
   126  		TableName: db,
   127  		Inst:      GetValue,
   128  	}
   129  	query.Key = key
   130  	return query, done
   131  }
   132  
   133  func NewLoadFromFileQuery(ctx context.Context, db []byte, filename []byte) (*Query, chan QueryResponse) {
   134  	done := make(chan QueryResponse, 1)
   135  	query := NewQuery(ctx, done)
   136  	query.Header = QueryHeader{
   137  		TableName: db,
   138  		FileName:  filename,
   139  		Inst:      Load,
   140  	}
   141  	return query, done
   142  }
   143  
   144  func NewSetValueQuery(ctx context.Context, db, key, value []byte) (*Query, chan QueryResponse) {
   145  	done := make(chan QueryResponse, 1)
   146  	query := NewQuery(ctx, done)
   147  	query.Header = QueryHeader{
   148  		TableName: db,
   149  		Inst:      SetValue,
   150  	}
   151  	query.Key = key
   152  	query.Value = value
   153  	return query, done
   154  }
   155  
   156  func NewBatchSetValueQuery(ctx context.Context, db []byte, values []KeyValue) (*Query, chan QueryResponse) {
   157  	done := make(chan QueryResponse, 1)
   158  	query := NewQuery(ctx, done)
   159  	query.Header = QueryHeader{
   160  		TableName: db,
   161  		Inst:      BatchSetValue,
   162  	}
   163  	query.Values = values
   164  
   165  	return query, done
   166  }
   167  
   168  func NewAddTableQuery(ctx context.Context, db []byte) (*Query, chan QueryResponse) {
   169  	done := make(chan QueryResponse, 1)
   170  	query := NewQuery(ctx, done)
   171  	query.Header = QueryHeader{
   172  		TableName: db,
   173  		Inst:      AddTable,
   174  	}
   175  	return query, done
   176  }