k3s/vendor/github.com/ibuildthecloud/kvsql/clientv3/driver/generic.go

package driver

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"sync/atomic"
	"time"

	"github.com/ibuildthecloud/kvsql/pkg/broadcast"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	utiltrace "k8s.io/utils/trace"
)

type Generic struct {
	// revision must be first to ensure that this is properly aligned for atomic.LoadInt64
	revision        int64

	db *sql.DB

	CleanupSQL      string
	GetSQL          string
	ListSQL         string
	ListRevisionSQL string
	ListResumeSQL   string
	ReplaySQL       string
	InsertSQL       string
	GetRevisionSQL  string
	ToDeleteSQL     string
	DeleteOldSQL    string

	changes     chan *KeyValue
	broadcaster broadcast.Broadcaster
	cancel      func()
}

func (g *Generic) Start(ctx context.Context, db *sql.DB) error {
	g.db = db
	g.changes = make(chan *KeyValue, 1024)

	row := db.QueryRowContext(ctx, g.GetRevisionSQL)
	rev := sql.NullInt64{}
	if err := row.Scan(&rev); err != nil {
		return errors.Wrap(err, "Failed to initialize revision")
	}
	if rev.Int64 == 0 {
		g.revision = 1
	} else {
		g.revision = rev.Int64
	}

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-time.After(time.Minute):
				_, err := g.ExecContext(ctx, g.CleanupSQL, time.Now().Unix())
				if err != nil {
					logrus.Errorf("Failed to purge expired TTL entries")
				}

				err = g.cleanup(ctx)
				if err != nil {
					logrus.Errorf("Failed to cleanup duplicate entries")
				}
			}
		}
	}()

	return nil
}

func (g *Generic) cleanup(ctx context.Context) error {
	rows, err := g.QueryContext(ctx, g.ToDeleteSQL)
	if err != nil {
		return err
	}
	defer rows.Close()

	toDelete := map[string]int64{}
	for rows.Next() {
		var (
			count, revision int64
			name            string
		)
		err := rows.Scan(&count, &name, &revision)
		if err != nil {
			return err
		}
		toDelete[name] = revision
	}

	rows.Close()

	for name, rev := range toDelete {
		_, err = g.ExecContext(ctx, g.DeleteOldSQL, name, rev, rev)
		if err != nil {
			return err
		}
	}

	return nil
}

func (g *Generic) Get(ctx context.Context, key string) (*KeyValue, error) {
	kvs, _, err := g.List(ctx, 0, 1, key, "")
	if err != nil {
		return nil, err
	}
	if len(kvs) > 0 {
		return kvs[0], nil
	}
	return nil, nil
}

func (g *Generic) replayEvents(ctx context.Context, key string, revision int64) ([]*KeyValue, error) {
	rows, err := g.QueryContext(ctx, g.ReplaySQL, key, revision)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var resp []*KeyValue
	for rows.Next() {
		value := KeyValue{}
		if err := scan(rows.Scan, &value); err != nil {
			return nil, err
		}
		resp = append(resp, &value)
	}

	return resp, nil
}

func (g *Generic) List(ctx context.Context, revision, limit int64, rangeKey, startKey string) ([]*KeyValue, int64, error) {
	var (
		rows *sql.Rows
		err  error
	)

	if limit == 0 {
		limit = 1000000
	} else {
		limit = limit + 1
	}

	listRevision := atomic.LoadInt64(&g.revision)
	if !strings.HasSuffix(rangeKey, "%") && revision <= 0 {
		rows, err = g.QueryContext(ctx, g.GetSQL, rangeKey, 1)
	} else if revision <= 0 {
		rows, err = g.QueryContext(ctx, g.ListSQL, rangeKey, limit)
	} else if len(startKey) > 0 {
		listRevision = revision
		rows, err = g.QueryContext(ctx, g.ListResumeSQL, revision, rangeKey, startKey, limit)
	} else {
		rows, err = g.QueryContext(ctx, g.ListRevisionSQL, revision, rangeKey, limit)
	}

	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	var resp []*KeyValue
	for rows.Next() {
		value := KeyValue{}
		if err := scan(rows.Scan, &value); err != nil {
			return nil, 0, err
		}
		if value.Revision > listRevision {
			listRevision = value.Revision
		}
		if value.Del == 0 {
			resp = append(resp, &value)
		}
	}

	return resp, listRevision, nil
}

func (g *Generic) Delete(ctx context.Context, key string, revision int64) ([]*KeyValue, error) {
	if strings.HasSuffix(key, "%") {
		panic("can not delete list revision")
	}

	_, err := g.mod(ctx, true, key, []byte{}, revision, 0)
	return nil, err
}

func (g *Generic) Update(ctx context.Context, key string, value []byte, revision, ttl int64) (*KeyValue, *KeyValue, error) {
	kv, err := g.mod(ctx, false, key, value, revision, ttl)
	if err != nil {
		return nil, nil, err
	}

	if kv.Version == 1 {
		return nil, kv, nil
	}

	oldKv := *kv
	oldKv.Revision = oldKv.OldRevision
	oldKv.Value = oldKv.OldValue
	return &oldKv, kv, nil
}

func (g *Generic) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	trace := utiltrace.New(fmt.Sprintf("SQL DB ExecContext query: %s keys: %v", query, args))
	defer trace.LogIfLong(500 * time.Millisecond)

	return g.db.ExecContext(ctx, query, args...)
}

func (g *Generic) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	trace := utiltrace.New(fmt.Sprintf("SQL DB QueryContext query: %s keys: %v", query, args))
	defer trace.LogIfLong(500 * time.Millisecond)

	return g.db.QueryContext(ctx, query, args...)
}

func (g *Generic) mod(ctx context.Context, delete bool, key string, value []byte, revision int64, ttl int64) (*KeyValue, error) {
	oldKv, err := g.Get(ctx, key)
	if err != nil {
		return nil, err
	}

	if revision > 0 && oldKv == nil {
		return nil, ErrNotExists
	}

	if revision > 0 && oldKv.Revision != revision {
		return nil, ErrRevisionMatch
	}

	if ttl > 0 {
		ttl = int64(time.Now().Unix()) + ttl
	}

	newRevision := atomic.AddInt64(&g.revision, 1)
	result := &KeyValue{
		Key:            key,
		Value:          value,
		Revision:       newRevision,
		TTL:            int64(ttl),
		CreateRevision: newRevision,
		Version:        1,
	}
	if oldKv != nil {
		result.OldRevision = oldKv.Revision
		result.OldValue = oldKv.Value
		result.TTL = oldKv.TTL
		result.CreateRevision = oldKv.CreateRevision
		result.Version = oldKv.Version + 1
	}

	if delete {
		result.Del = 1
	}

	_, err = g.ExecContext(ctx, g.InsertSQL,
		result.Key,
		result.Value,
		result.OldValue,
		result.OldRevision,
		result.CreateRevision,
		result.Revision,
		result.TTL,
		result.Version,
		result.Del,
	)
	if err != nil {
		return nil, err
	}

	g.changes <- result
	return result, nil
}

type scanner func(dest ...interface{}) error

func scan(s scanner, out *KeyValue) error {
	return s(
		&out.ID,
		&out.Key,
		&out.Value,
		&out.OldValue,
		&out.OldRevision,
		&out.CreateRevision,
		&out.Revision,
		&out.TTL,
		&out.Version,
		&out.Del)
}