kythe.io@v0.0.68-0.20240422202219-7225dbc01741/kythe/go/storage/table/table.go (about)

     1  /*
     2   * Copyright 2015 The Kythe Authors. All rights reserved.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *   http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  // Package table implements lookup table interfaces for protobufs.
    18  package table // import "kythe.io/kythe/go/storage/table"
    19  
    20  import (
    21  	"context"
    22  	"errors"
    23  	"fmt"
    24  	"io"
    25  
    26  	"kythe.io/kythe/go/storage/keyvalue"
    27  
    28  	"google.golang.org/protobuf/proto"
    29  )
    30  
    31  // Proto is a key-value direct lookup table with protobuf values.
    32  type Proto interface {
    33  	ProtoLookup
    34  
    35  	// Put marshals msg and writes it as the value for the given key.
    36  	Put(ctx context.Context, key []byte, msg proto.Message) error
    37  
    38  	// Buffered returns a buffered write interface.
    39  	Buffered() BufferedProto
    40  
    41  	// Close release the underlying resources for the table.
    42  	Close(context.Context) error
    43  }
    44  
    45  // ProtoLookup is a read-only key-value direct lookup table with protobuf values.
    46  type ProtoLookup interface {
    47  	// Lookup unmarshals the value for the given key into msg, returning any
    48  	// error.  If the key was not found, ErrNoSuchKey is returned.
    49  	Lookup(ctx context.Context, key []byte, msg proto.Message) error
    50  
    51  	// Lookup unmarshals the values for the given key into new proto.Message using
    52  	// m, returning any error.  If the key was not found, f is never called.
    53  	LookupValues(ctx context.Context, key []byte, m proto.Message, f func(msg proto.Message) error) error
    54  }
    55  
    56  // BufferedProto buffers calls to Put to provide a high throughput write
    57  // interface to a Proto table.
    58  type BufferedProto interface {
    59  	// Put marshals msg and writes it as the value for the given key.
    60  	Put(ctx context.Context, key []byte, msg proto.Message) error
    61  
    62  	// Flush sends all buffered writes to the underlying table.
    63  	Flush(ctx context.Context) error
    64  }
    65  
    66  // KVProto implements a Proto table using a keyvalue.DB.
    67  type KVProto struct{ keyvalue.DB }
    68  
    69  // ErrNoSuchKey is returned when a value was not found for a particular key.
    70  var ErrNoSuchKey = errors.New("no such key")
    71  
    72  // ErrStopLookup should be returned from the function passed to LookupValues
    73  // when the client wants no further values.
    74  var ErrStopLookup = errors.New("stop lookup")
    75  
    76  // Lookup implements part of the Proto interface.
    77  func (t *KVProto) Lookup(ctx context.Context, key []byte, msg proto.Message) error {
    78  	v, err := t.Get(ctx, key, nil)
    79  	if err == io.EOF {
    80  		return ErrNoSuchKey
    81  	} else if err != nil {
    82  		return err
    83  	} else if err := proto.Unmarshal(v, msg); err != nil {
    84  		return fmt.Errorf("proto unmarshal error: %v", err)
    85  	}
    86  	return nil
    87  }
    88  
    89  // LookupValues implements part of the ProtoLookup interface.
    90  func (t *KVProto) LookupValues(ctx context.Context, key []byte, m proto.Message, f func(proto.Message) error) error {
    91  	it, err := t.ScanRange(ctx, keyvalue.KeyRange(key), nil)
    92  	if err != nil {
    93  		return err
    94  	}
    95  	for {
    96  		_, val, err := it.Next()
    97  		if err == io.EOF {
    98  			return nil
    99  		} else if err != nil {
   100  			return err
   101  		}
   102  
   103  		msg := m.ProtoReflect().New().Interface()
   104  		if err := proto.Unmarshal(val, msg); err != nil {
   105  			return err
   106  		}
   107  
   108  		if err := f(msg); err == ErrStopLookup {
   109  			return nil
   110  		} else if err != nil {
   111  			return err
   112  		}
   113  	}
   114  	return nil
   115  }
   116  
   117  // Put implements part of the Proto interface.
   118  func (t *KVProto) Put(ctx context.Context, key []byte, msg proto.Message) error {
   119  	b := t.Buffered()
   120  	if err := b.Put(ctx, key, msg); err != nil {
   121  		return err
   122  	}
   123  	return b.Flush(ctx)
   124  }
   125  
   126  type kvProtoBuffer struct{ pool *keyvalue.WritePool }
   127  
   128  // Put implements part of the BufferedProto interface.
   129  func (b *kvProtoBuffer) Put(ctx context.Context, key []byte, msg proto.Message) error {
   130  	rec, err := proto.Marshal(msg)
   131  	if err != nil {
   132  		return err
   133  	}
   134  	return b.pool.Write(ctx, key, rec)
   135  }
   136  
   137  // Flush implements part of the BufferedProto interface.
   138  func (b *kvProtoBuffer) Flush(_ context.Context) error { return b.pool.Flush() }
   139  
   140  // Buffered implements part of the Proto interface.
   141  func (t *KVProto) Buffered() BufferedProto { return &kvProtoBuffer{keyvalue.NewPool(t.DB, nil)} }
   142  
   143  // Close implements part of the Proto interface.
   144  func (t *KVProto) Close(ctx context.Context) error { return t.DB.Close(ctx) }