tasks/tools/cassandra/cqlclient.go

package cassandra

import (
	"context"
	"fmt"
	"time"

	"github.com/gocql/gocql"
	"go.temporal.io/server/common/auth"
	"go.temporal.io/server/common/config"
	"go.temporal.io/server/common/log"
	"go.temporal.io/server/common/log/tag"
	"go.temporal.io/server/common/metrics"
	commongocql "go.temporal.io/server/common/persistence/nosql/nosqlplugin/cassandra/gocql"
	"go.temporal.io/server/common/resolver"
	"go.temporal.io/server/tools/common/schema"
)

type (
	cqlClient struct {
		nReplicas  int
		datacenter string
		keyspace   string
		timeout    time.Duration
		session    commongocql.Session
		logger     log.Logger
	}
	// CQLClientConfig contains the configuration for cql client
	CQLClientConfig struct {
		Hosts                    string
		Port                     int
		User                     string
		Password                 string
		AllowedAuthenticators    []string
		Keyspace                 string
		Timeout                  int
		numReplicas              int
		Datacenter               string
		Consistency              string
		TLS                      *auth.TLS
		DisableInitialHostLookup bool
		AddressTranslator        *config.CassandraAddressTranslator
	}
)

const (
	defaultTimeout = 30 // Timeout in seconds
	systemKeyspace = "system"
	dbType         = "cassandra"
)

const (
	readSchemaVersionCQL        = `SELECT curr_version from schema_version where keyspace_name=?`
	listTablesCQL               = `SELECT table_name from system_schema.tables where keyspace_name=?`
	listTypesCQL                = `SELECT type_name from system_schema.types where keyspace_name=?`
	writeSchemaVersionCQL       = `INSERT into schema_version(keyspace_name, creation_time, curr_version, min_compatible_version) VALUES (?,?,?,?)`
	writeSchemaUpdateHistoryCQL = `INSERT into schema_update_history(year, month, update_time, old_version, new_version, manifest_md5, description) VALUES(?,?,?,?,?,?,?)`

	createSchemaVersionTableCQL = `CREATE TABLE IF NOT EXISTS schema_version(keyspace_name text PRIMARY KEY, ` +
		`creation_time timestamp, ` +
		`curr_version text, ` +
		`min_compatible_version text);`

	createSchemaUpdateHistoryTableCQL = `CREATE TABLE IF NOT EXISTS schema_update_history(` +
		`year int, ` +
		`month int, ` +
		`update_time timestamp, ` +
		`description text, ` +
		`manifest_md5 text, ` +
		`new_version text, ` +
		`old_version text, ` +
		`PRIMARY KEY ((year, month), update_time));`

	createKeyspaceCQL = `CREATE KEYSPACE IF NOT EXISTS %v ` +
		`WITH replication = { 'class' : 'SimpleStrategy', 'replication_factor' : %v};`

	createKeyspaceNetworkTopologyCQL = `CREATE KEYSPACE IF NOT EXISTS %v ` +
		`WITH replication = { 'class' : 'NetworkTopologyStrategy', '%v' : %v};`
)

var _ schema.DB = (*cqlClient)(nil)

// newCQLClient returns a new instance of CQLClient
func newCQLClient(cfg *CQLClientConfig, logger log.Logger) (*cqlClient, error) {
	var err error

	cassandraConfig := cfg.toCassandraConfig()

	logger.Info("Validating connection to cassandra cluster.")
	session, err := commongocql.NewSession(
		func() (*gocql.ClusterConfig, error) {
			return commongocql.NewCassandraCluster(*cassandraConfig, resolver.NewNoopResolver())
		},
		logger,
		metrics.NoopMetricsHandler,
	)
	if err != nil {
		logger.Error("Connection validation failed.", tag.Error(err))
		return nil, err
	}
	logger.Info("Connection validation succeeded.")

	return &cqlClient{
		keyspace:   cfg.Keyspace,
		nReplicas:  cfg.numReplicas,
		datacenter: cfg.Datacenter,
		timeout:    time.Duration(cfg.Timeout) * time.Second,
		session:    session,
		logger:     logger,
	}, nil
}

func (cfg *CQLClientConfig) toCassandraConfig() *config.Cassandra {
	cassandraConfig := config.Cassandra{
		Hosts:                    cfg.Hosts,
		Port:                     cfg.Port,
		User:                     cfg.User,
		Password:                 cfg.Password,
		AllowedAuthenticators:    cfg.AllowedAuthenticators,
		Keyspace:                 cfg.Keyspace,
		TLS:                      cfg.TLS,
		Datacenter:               cfg.Datacenter,
		DisableInitialHostLookup: cfg.DisableInitialHostLookup,
		Consistency: &config.CassandraStoreConsistency{
			Default: &config.CassandraConsistencySettings{
				Consistency: cfg.Consistency,
			},
		},
		AddressTranslator: cfg.AddressTranslator,
		ConnectTimeout:    time.Duration(cfg.Timeout) * time.Second,
	}

	return &cassandraConfig
}

func (client *cqlClient) CreateDatabase(name string) error {
	return client.createKeyspace(name)
}

func (client *cqlClient) DropDatabase(name string) error {
	return client.dropKeyspace(name)
}

// createKeyspace creates a cassandra Keyspace if it doesn't exist
func (client *cqlClient) createKeyspace(name string) error {
	if client.datacenter != "" {
		client.logger.Info(fmt.Sprintf("Creating Keyspace %v using NetworkTopologyStrategy in Datacenter %v with RF=%v.", name, client.datacenter, client.nReplicas))
		return client.Exec(fmt.Sprintf(createKeyspaceNetworkTopologyCQL, name, client.datacenter, client.nReplicas))
	}
	client.logger.Info(fmt.Sprintf("Creating Keyspace %v using SimpleStrategy with RF=%v.", name, client.nReplicas))
	return client.Exec(fmt.Sprintf(createKeyspaceCQL, name, client.nReplicas))
}

// dropKeyspace drops a Keyspace
func (client *cqlClient) dropKeyspace(name string) error {
	return client.Exec(fmt.Sprintf("DROP KEYSPACE IF EXISTS %v", name))
}

func (client *cqlClient) DropAllTables() error {
	return client.dropAllTablesTypes()
}

// CreateSchemaVersionTables sets up the schema version tables
func (client *cqlClient) CreateSchemaVersionTables() error {
	if err := client.Exec(createSchemaVersionTableCQL); err != nil {
		return err
	}
	return client.Exec(createSchemaUpdateHistoryTableCQL)
}

// ReadSchemaVersion returns the current schema version for the Keyspace
func (client *cqlClient) ReadSchemaVersion() (string, error) {
	query := client.session.Query(readSchemaVersionCQL, client.keyspace)

	iter := query.Iter()
	var version string
	success := iter.Scan(&version)
	err := iter.Close()
	if err == nil && !success {
		err = fmt.Errorf("no schema version found for keyspace %q", client.keyspace)
	}
	if err != nil {
		return "", fmt.Errorf("unable to get current schema version from Cassandra: %w", err)
	}
	return version, nil
}

// UpdateShemaVersion updates the schema version for the Keyspace
func (client *cqlClient) UpdateSchemaVersion(newVersion string, minCompatibleVersion string) error {
	query := client.session.Query(writeSchemaVersionCQL, client.keyspace, time.Now().UTC(), newVersion, minCompatibleVersion)
	return query.Exec()
}

// WriteSchemaUpdateLog adds an entry to the schema update history table
func (client *cqlClient) WriteSchemaUpdateLog(oldVersion string, newVersion string, manifestMD5 string, desc string) error {
	now := time.Now().UTC()
	query := client.session.Query(writeSchemaUpdateHistoryCQL)
	query.Bind(now.Year(), int(now.Month()), now, oldVersion, newVersion, manifestMD5, desc)
	return query.Exec()
}

// Exec executes a cql statement
func (client *cqlClient) Exec(stmt string, args ...interface{}) error {
	if err := client.session.Query(stmt, args...).Exec(); err != nil {
		return err
	}
	return client.waitSchemaAgreement()
}

// Close closes the cql client
func (client *cqlClient) Close() {
	if client.session != nil {
		client.session.Close()
	}
}

// ListTables lists the table names in a Keyspace
func (client *cqlClient) ListTables() ([]string, error) {
	query := client.session.Query(listTablesCQL, client.keyspace)
	iter := query.Iter()
	var names []string
	var name string
	for iter.Scan(&name) {
		names = append(names, name)
	}
	if err := iter.Close(); err != nil {
		return nil, err
	}
	return names, nil
}

// listTypes lists the User defined types in a Keyspace
func (client *cqlClient) listTypes() ([]string, error) {
	qry := client.session.Query(listTypesCQL, client.keyspace)
	iter := qry.Iter()
	var names []string
	var name string
	for iter.Scan(&name) {
		names = append(names, name)
	}
	if err := iter.Close(); err != nil {
		return nil, err
	}
	return names, nil
}

// dropTable drops a given table from the Keyspace
func (client *cqlClient) dropTable(name string) error {
	return client.Exec(fmt.Sprintf("DROP TABLE %v", name))
}

// dropType drops a given type from the Keyspace
func (client *cqlClient) dropType(name string) error {
	return client.Exec(fmt.Sprintf("DROP TYPE %v", name))
}

// dropAllTablesTypes deletes all tables/types in the
// Keyspace without deleting the Keyspace
func (client *cqlClient) dropAllTablesTypes() error {
	tables, err := client.ListTables()
	if err != nil {
		return err
	}
	client.logger.Info(fmt.Sprintf("Dropping following tables: %v.", tables))
	for _, table := range tables {
		err1 := client.dropTable(table)
		if err1 != nil {
			client.logger.Error(fmt.Sprintf("Error dropping table %v.", table), tag.Error(err1))
		}
	}

	types, err := client.listTypes()
	if err != nil {
		return err
	}
	client.logger.Info(fmt.Sprintf("Dropping following types: %v.", types))
	numOfTypes := len(types)
	for i := 0; i < numOfTypes && len(types) > 0; i++ {
		var erroredTypes []string
		for _, t := range types {
			err = client.dropType(t)
			if err != nil {
				client.logger.Error(fmt.Sprintf("Error dropping type %v.", t), tag.Error(err))
				erroredTypes = append(erroredTypes, t)
			}
		}
		types = erroredTypes
	}
	if len(types) > 0 {
		return err
	}
	return nil
}

// waitSchemaAgreement wait for schema change agreements
func (client *cqlClient) waitSchemaAgreement() error {
	ctx, cancel := context.WithTimeout(context.Background(), client.timeout)
	defer cancel()
	return client.session.AwaitSchemaAgreement(ctx)
}

// Type gives the type of db
func (client *cqlClient) Type() string {
	return dbType
}