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

     1  package fileset
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"testing"
     7  
     8  	"github.com/gogo/protobuf/proto"
     9  	"github.com/jmoiron/sqlx"
    10  )
    11  
    12  var _ Store = &postgresStore{}
    13  
    14  type postgresStore struct {
    15  	db *sqlx.DB
    16  }
    17  
    18  // NewPostgresStore returns a Store backed by db
    19  func NewPostgresStore(db *sqlx.DB) Store {
    20  	return &postgresStore{db: db}
    21  }
    22  
    23  func (s *postgresStore) Set(ctx context.Context, p string, md *Metadata) error {
    24  	if md == nil {
    25  		md = &Metadata{}
    26  	}
    27  	data, err := proto.Marshal(md)
    28  	if err != nil {
    29  		return err
    30  	}
    31  	res, err := s.db.ExecContext(ctx,
    32  		`INSERT INTO storage.filesets (path, metadata_pb)
    33  		VALUES ($1, $2)
    34  		ON CONFLICT (path) DO NOTHING
    35  		`, p, data)
    36  	if err != nil {
    37  		return err
    38  	}
    39  	n, err := res.RowsAffected()
    40  	if err != nil {
    41  		return err
    42  	}
    43  	if n == 0 {
    44  		return ErrPathExists
    45  	}
    46  	return nil
    47  }
    48  
    49  func (s *postgresStore) Get(ctx context.Context, p string) (*Metadata, error) {
    50  	var mdData []byte
    51  	if err := s.db.GetContext(ctx, &mdData, `SELECT metadata_pb FROM storage.filesets WHERE path = $1`, p); err != nil {
    52  		if err == sql.ErrNoRows {
    53  			return nil, ErrPathNotExists
    54  		}
    55  		return nil, err
    56  	}
    57  	md := &Metadata{}
    58  	if err := proto.Unmarshal(mdData, md); err != nil {
    59  		return nil, err
    60  	}
    61  	return md, nil
    62  }
    63  
    64  func (s *postgresStore) Walk(ctx context.Context, prefix string, cb func(string) error) (retErr error) {
    65  	rows, err := s.db.QueryContext(ctx, `SELECT path from storage.filesets WHERE path LIKE $1 || '%'`, prefix)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	defer func() {
    70  		if err := rows.Close(); retErr == nil {
    71  			retErr = err
    72  		}
    73  	}()
    74  	var p string
    75  	for rows.Next() {
    76  		if err := rows.Scan(&p); err != nil {
    77  			return err
    78  		}
    79  		if err := cb(p); err != nil {
    80  			return err
    81  		}
    82  	}
    83  	return rows.Err()
    84  }
    85  
    86  func (s *postgresStore) Delete(ctx context.Context, p string) error {
    87  	_, err := s.db.ExecContext(ctx, `DELETE FROM storage.filesets WHERE path = $1`, p)
    88  	return err
    89  }
    90  
    91  const schema = `
    92  	CREATE SCHEMA IF NOT EXISTS storage;
    93  
    94  	CREATE TABLE IF NOT EXISTS storage.filesets (
    95  		path VARCHAR(250) PRIMARY KEY,
    96  		metadata_pb BYTEA NOT NULL,
    97  		created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
    98  	);
    99  `
   100  
   101  // SetupPostgresStore sets up the tables for a Store
   102  func SetupPostgresStore(db *sqlx.DB) {
   103  	db.MustExec(schema)
   104  }
   105  
   106  // NewTestStore returns a Store scoped to the lifetime of the test.
   107  func NewTestStore(t testing.TB, db *sqlx.DB) Store {
   108  	SetupPostgresStore(db)
   109  	return NewPostgresStore(db)
   110  }