micro/store/cockroach/cockroach.go

443 lines
9.3 KiB
Go
Raw Normal View History

2019-12-16 18:09:59 +03:00
// Package cockroach implements the cockroach store
package cockroach
2019-11-01 17:13:21 +03:00
import (
"database/sql"
"fmt"
"net/url"
"regexp"
2019-12-16 18:09:59 +03:00
"strings"
"sync"
2019-11-01 17:13:21 +03:00
"time"
"github.com/lib/pq"
"github.com/micro/go-micro/v2/store"
2019-12-16 17:38:51 +03:00
"github.com/pkg/errors"
2019-11-01 17:13:21 +03:00
)
// DefaultDatabase is the namespace that the sql store
2019-11-01 17:13:21 +03:00
// will use if no namespace is provided.
var (
DefaultDatabase = "micro"
DefaultTable = "micro"
)
2019-11-01 17:13:21 +03:00
var (
statements = map[string]string{
"list": "SELECT key, value, expiry FROM %s.%s;",
"read": "SELECT key, value, expiry FROM %s.%s WHERE key = $1;",
"readMany": "SELECT key, value, expiry FROM %s.%s WHERE key LIKE $1;",
"readOffset": "SELECT key, value, expiry FROM %s.%s WHERE key LIKE $1 ORDER BY key DESC LIMIT $2 OFFSET $3;",
"write": "INSERT INTO %s.%s(key, value, expiry) VALUES ($1, $2::bytea, $3) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, expiry = EXCLUDED.expiry;",
"delete": "DELETE FROM %s.%s WHERE key = $1;",
}
)
2019-11-01 17:13:21 +03:00
type sqlStore struct {
options store.Options
db *sql.DB
2019-11-01 17:13:21 +03:00
sync.RWMutex
// known databases
databases map[string]bool
}
2019-12-16 17:38:51 +03:00
func (s *sqlStore) createDB(database, table string) {
if len(database) == 0 {
database = s.options.Database
}
if len(table) == 0 {
table = s.options.Table
}
s.Lock()
_, ok := s.databases[database+":"+table]
if !ok {
s.initDB(database, table)
s.databases[database+":"+table] = true
}
s.Unlock()
}
func (s *sqlStore) initDB(database, table string) error {
// Create the namespace's database
_, err := s.db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s ;", database))
if err != nil {
return err
}
_, err = s.db.Exec(fmt.Sprintf("SET DATABASE = %s ;", database))
if err != nil {
return errors.Wrap(err, "Couldn't set database")
}
// Create a table for the namespace's prefix
_, err = s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s
(
key text NOT NULL,
value bytea,
expiry timestamp with time zone,
CONSTRAINT %s_pkey PRIMARY KEY (key)
);`, table, table))
if err != nil {
return errors.Wrap(err, "Couldn't create table")
}
// Create Index
_, err = s.db.Exec(fmt.Sprintf(`CREATE INDEX IF NOT EXISTS "%s" ON %s.%s USING btree ("key");`, "key_index_"+table, database, table))
if err != nil {
return err
}
return nil
}
func (s *sqlStore) configure() error {
if len(s.options.Nodes) == 0 {
s.options.Nodes = []string{"postgresql://root@localhost:26257?sslmode=disable"}
}
database := s.options.Database
if len(database) == 0 {
database = DefaultDatabase
}
table := s.options.Table
if len(table) == 0 {
table = DefaultTable
}
// store.namespace must only contain letters, numbers and underscores
reg, err := regexp.Compile("[^a-zA-Z0-9]+")
if err != nil {
return errors.New("error compiling regex for namespace")
}
database = reg.ReplaceAllString(database, "_")
table = reg.ReplaceAllString(table, "_")
source := s.options.Nodes[0]
// check if it is a standard connection string eg: host=%s port=%d user=%s password=%s dbname=%s sslmode=disable
// if err is nil which means it would be a URL like postgre://xxxx?yy=zz
_, err = url.Parse(source)
if err != nil {
if !strings.Contains(source, " ") {
source = fmt.Sprintf("host=%s", source)
}
}
// create source from first node
db, err := sql.Open("postgres", source)
if err != nil {
return err
}
if err := db.Ping(); err != nil {
return err
}
if s.db != nil {
s.db.Close()
}
// save the values
s.db = db
// initialise the database
return s.initDB(s.options.Database, s.options.Table)
}
func (s *sqlStore) prepare(database, table, query string) (*sql.Stmt, error) {
st, ok := statements[query]
if !ok {
return nil, errors.New("unsupported statement")
}
if len(database) == 0 {
database = s.options.Database
}
if len(table) == 0 {
table = s.options.Table
}
q := fmt.Sprintf(st, database, table)
stmt, err := s.db.Prepare(q)
if err != nil {
return nil, err
}
return stmt, nil
2019-11-01 17:13:21 +03:00
}
func (s *sqlStore) Close() error {
if s.db != nil {
return s.db.Close()
}
return nil
}
func (s *sqlStore) Init(opts ...store.Option) error {
for _, o := range opts {
o(&s.options)
}
// reconfigure
return s.configure()
}
2019-11-01 17:13:21 +03:00
// List all the known records
func (s *sqlStore) List(opts ...store.ListOption) ([]string, error) {
var options store.ListOptions
for _, o := range opts {
o(&options)
}
// create the db if not exists
s.createDB(options.Database, options.Table)
st, err := s.prepare(options.Database, options.Table, "list")
if err != nil {
return nil, err
}
defer st.Close()
rows, err := st.Query()
2019-11-01 17:13:21 +03:00
if err != nil {
if err == sql.ErrNoRows {
return nil, nil
2019-11-01 17:13:21 +03:00
}
return nil, err
}
defer rows.Close()
var keys []string
var timehelper pq.NullTime
2019-11-01 17:13:21 +03:00
for rows.Next() {
record := &store.Record{}
if err := rows.Scan(&record.Key, &record.Value, &timehelper); err != nil {
return keys, err
2019-11-01 17:13:21 +03:00
}
if timehelper.Valid {
if timehelper.Time.Before(time.Now()) {
// record has expired
go s.Delete(record.Key)
} else {
record.Expiry = time.Until(timehelper.Time)
keys = append(keys, record.Key)
2019-11-01 17:13:21 +03:00
}
} else {
keys = append(keys, record.Key)
2019-11-01 17:13:21 +03:00
}
}
rowErr := rows.Close()
if rowErr != nil {
// transaction rollback or something
return keys, rowErr
2019-11-01 17:13:21 +03:00
}
if err := rows.Err(); err != nil {
return keys, err
2019-11-01 17:13:21 +03:00
}
return keys, nil
2019-11-01 17:13:21 +03:00
}
// Read a single key
func (s *sqlStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
var options store.ReadOptions
for _, o := range opts {
o(&options)
}
// create the db if not exists
s.createDB(options.Database, options.Table)
if options.Prefix || options.Suffix {
return s.read(key, options)
}
2019-11-01 17:13:21 +03:00
var records []*store.Record
var timehelper pq.NullTime
st, err := s.prepare(options.Database, options.Table, "read")
if err != nil {
return nil, err
}
defer st.Close()
row := st.QueryRow(key)
record := &store.Record{}
if err := row.Scan(&record.Key, &record.Value, &timehelper); err != nil {
if err == sql.ErrNoRows {
return records, store.ErrNotFound
2019-11-01 17:13:21 +03:00
}
return records, err
}
if timehelper.Valid {
if timehelper.Time.Before(time.Now()) {
// record has expired
go s.Delete(key)
return records, store.ErrNotFound
2019-11-01 17:13:21 +03:00
}
record.Expiry = time.Until(timehelper.Time)
records = append(records, record)
} else {
records = append(records, record)
2019-11-01 17:13:21 +03:00
}
2019-11-01 17:13:21 +03:00
return records, nil
}
// Read Many records
func (s *sqlStore) read(key string, options store.ReadOptions) ([]*store.Record, error) {
pattern := "%"
if options.Prefix {
pattern = key + pattern
}
if options.Suffix {
pattern = pattern + key
}
var rows *sql.Rows
var err error
if options.Limit != 0 {
st, err := s.prepare(options.Database, options.Table, "readOffset")
if err != nil {
return nil, err
}
defer st.Close()
rows, err = st.Query(pattern, options.Limit, options.Offset)
} else {
st, err := s.prepare(options.Database, options.Table, "readMany")
if err != nil {
return nil, err
}
defer st.Close()
rows, err = st.Query(pattern)
}
if err != nil {
if err == sql.ErrNoRows {
return []*store.Record{}, nil
}
return []*store.Record{}, errors.Wrap(err, "sqlStore.read failed")
}
defer rows.Close()
var records []*store.Record
var timehelper pq.NullTime
for rows.Next() {
record := &store.Record{}
if err := rows.Scan(&record.Key, &record.Value, &timehelper); err != nil {
return records, err
}
if timehelper.Valid {
if timehelper.Time.Before(time.Now()) {
// record has expired
go s.Delete(record.Key)
} else {
record.Expiry = time.Until(timehelper.Time)
records = append(records, record)
}
} else {
records = append(records, record)
}
}
rowErr := rows.Close()
if rowErr != nil {
// transaction rollback or something
return records, rowErr
}
if err := rows.Err(); err != nil {
return records, err
}
return records, nil
}
2019-11-01 17:13:21 +03:00
// Write records
func (s *sqlStore) Write(r *store.Record, opts ...store.WriteOption) error {
var options store.WriteOptions
for _, o := range opts {
o(&options)
}
// create the db if not exists
s.createDB(options.Database, options.Table)
st, err := s.prepare(options.Database, options.Table, "write")
if err != nil {
return err
}
defer st.Close()
if r.Expiry != 0 {
_, err = st.Exec(r.Key, r.Value, time.Now().Add(r.Expiry))
} else {
_, err = st.Exec(r.Key, r.Value, nil)
}
if err != nil {
return errors.Wrap(err, "Couldn't insert record "+r.Key)
2019-11-01 17:13:21 +03:00
}
return nil
}
// Delete records with keys
func (s *sqlStore) Delete(key string, opts ...store.DeleteOption) error {
var options store.DeleteOptions
for _, o := range opts {
o(&options)
}
// create the db if not exists
s.createDB(options.Database, options.Table)
2019-11-01 17:13:21 +03:00
st, err := s.prepare(options.Database, options.Table, "delete")
2019-11-01 17:13:21 +03:00
if err != nil {
return err
2019-11-01 17:13:21 +03:00
}
defer st.Close()
2019-12-16 17:38:51 +03:00
result, err := st.Exec(key)
if err != nil {
return err
}
_, err = result.RowsAffected()
if err != nil {
return err
}
return nil
}
2019-11-01 17:13:21 +03:00
func (s *sqlStore) Options() store.Options {
return s.options
}
func (s *sqlStore) String() string {
return "cockroach"
}
// NewStore returns a new micro Store backed by sql
func NewStore(opts ...store.Option) store.Store {
var options store.Options
for _, o := range opts {
o(&options)
}
// new store
s := new(sqlStore)
// set the options
s.options = options
// mark known databases
s.databases = make(map[string]bool)
2020-03-24 20:16:38 +03:00
// best-effort configure the store
s.configure()
// return store
2019-12-16 17:38:51 +03:00
return s
2019-11-01 17:13:21 +03:00
}