github.com/pachyderm/pachyderm@v1.13.4/src/server/pkg/storage/chunk/metadata.go (about)

     1  package chunk
     2  
     3  import (
     4  	"context"
     5  	"crypto/sha512"
     6  	"database/sql"
     7  	"encoding/hex"
     8  
     9  	"github.com/jmoiron/sqlx"
    10  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
    11  )
    12  
    13  // ID uniquely identifies a chunk. It is the hash of its content
    14  type ID []byte
    15  
    16  // Hash produces an ID by hashing data
    17  func Hash(data []byte) ID {
    18  	h := sha512.New()
    19  	h.Write(data)
    20  	return h.Sum(nil)[:32]
    21  }
    22  
    23  // IDFromHex parses a hex string into an ID
    24  func IDFromHex(h string) (ID, error) {
    25  	return hex.DecodeString(h)
    26  }
    27  
    28  // HexString hex encodes the ID
    29  func (id ID) HexString() string {
    30  	return hex.EncodeToString(id)
    31  }
    32  
    33  // Metadata holds metadata about a chunk
    34  type Metadata struct {
    35  	Size     int
    36  	PointsTo []ID
    37  }
    38  
    39  var (
    40  	// ErrMetadataExists metadata exists
    41  	ErrMetadataExists = errors.Errorf("metadata exists")
    42  	// ErrChunkNotExists chunk does not exist
    43  	ErrChunkNotExists = errors.Errorf("chunk does not exist")
    44  )
    45  
    46  // MetadataStore stores metadata about chunks
    47  type MetadataStore interface {
    48  	// Set adds chunk metadata to the tracker
    49  	Set(ctx context.Context, chunkID ID, md Metadata) error
    50  	// Get returns info about the chunk if it exists
    51  	Get(ctx context.Context, chunkID ID) (*Metadata, error)
    52  	// Delete removes chunk metadata from the tracker
    53  	Delete(ctx context.Context, chunkID ID) error
    54  }
    55  
    56  var _ MetadataStore = &postgresStore{}
    57  
    58  type postgresStore struct {
    59  	db *sqlx.DB
    60  }
    61  
    62  // NewPostgresStore returns a Metadata backed by db
    63  func NewPostgresStore(db *sqlx.DB) MetadataStore {
    64  	return &postgresStore{db: db}
    65  }
    66  
    67  func (s *postgresStore) Set(ctx context.Context, chunkID ID, md Metadata) error {
    68  	res, err := s.db.ExecContext(ctx,
    69  		`INSERT INTO storage.chunks (hash_id, size) VALUES ($1, $2)
    70  		ON CONFLICT DO NOTHING
    71  		`, chunkID, md.Size)
    72  	if err != nil {
    73  		return err
    74  	}
    75  	n, err := res.RowsAffected()
    76  	if err != nil {
    77  		return err
    78  	}
    79  	if n == 0 {
    80  		return ErrMetadataExists
    81  	}
    82  	return nil
    83  }
    84  
    85  func (s *postgresStore) Get(ctx context.Context, chunkID ID) (*Metadata, error) {
    86  	type chunkRow struct {
    87  		size int `db:"size"`
    88  	}
    89  	var x chunkRow
    90  	if err := s.db.GetContext(ctx, &x, `SELECT size FROM storage.chunks WHERE hash_id = $1`, chunkID); err != nil {
    91  		if err == sql.ErrNoRows {
    92  			err = ErrChunkNotExists
    93  		}
    94  		return nil, err
    95  	}
    96  	return &Metadata{
    97  		Size: x.size,
    98  	}, nil
    99  }
   100  
   101  func (s *postgresStore) Delete(ctx context.Context, chunkID ID) error {
   102  	_, err := s.db.ExecContext(ctx, `DELETE FROM storage.chunks WHERE hash_id = $1`, chunkID)
   103  	return err
   104  }
   105  
   106  // SetupPostgresStore sets up tables in db
   107  func SetupPostgresStore(db *sqlx.DB) {
   108  	db.MustExec(schema)
   109  }
   110  
   111  const schema = `
   112  	CREATE SCHEMA IF NOT EXISTS storage;
   113  
   114  	CREATE TABLE IF NOT EXISTS storage.chunks (
   115  		hash_id BYTEA NOT NULL UNIQUE,
   116  		size INT8 NOT NULL,
   117  		created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
   118  	);	
   119  `