package scylladb import ( "errors" "fmt" "github.com/gocql/gocql" "strings" "time" ) // Storage interface that is implemented by storage providers type Storage struct { session *gocql.Session tableName string selectQuery string insertQuery string deleteQuery string resetQuery string } var ( checkSchemaErrorMsg = errors.New("the `value` row has an incorrect data type. " + "The message should be BLOB, but it is instead %s. This could lead to encoding-related issues if the database is not migrated (refer to https://github.com/gofiber/storage/blob/main/MIGRATE.md)") keyspaceErrorMsg = errors.New(`keyspace cannot be empty`) createKeyspaceQuery = `CREATE KEYSPACE IF NOT EXISTS %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};` dropQuery = `DROP TABLE IF EXISTS %s.%s;` createTableQuery = `CREATE TABLE IF NOT EXISTS %s.%s (key TEXT PRIMARY KEY, value BLOB)` checkSchemaQuery = `SELECT type FROM system_schema.columns WHERE keyspace_name = '%s' AND table_name = '%s' AND column_name = 'value';` selectQuery = `SELECT value FROM %s.%s WHERE key = ?` insertQuery = `INSERT INTO %s.%s (key, value) VALUES (?, ?) USING TTL ?` deleteQuery = `DELETE FROM %s.%s WHERE key = ?` resetQuery = `TRUNCATE %s.%s` ) // New creates a new storage func New(config ...Config) *Storage { var err error var session *gocql.Session // Set default config cfg := configDefault(config...) if len(strings.TrimSpace(cfg.Keyspace)) == 0 { panic(keyspaceErrorMsg) } if cfg.Session == nil { // Create a new cluster cluster := gocql.NewCluster(cfg.Hosts...) cluster.Consistency = gocql.ParseConsistency(cfg.Consistency) cluster.Port = cfg.Port // Set credentials if provided if len(strings.TrimSpace(cfg.Username)) > 0 && len(strings.TrimSpace(cfg.Password)) > 0 { cluster.Authenticator = gocql.PasswordAuthenticator{ Username: cfg.Username, Password: cfg.Password, } } // Set cfg.SslOpts if provided. if cfg.SslOpts != nil { cluster.SslOpts = cfg.SslOpts } // Create session session, err = cluster.CreateSession() if err != nil { panic(err) } } else { // Set session if provided session = cfg.Session } // Create keyspace if it does not exist if err = session.Query(fmt.Sprintf(createKeyspaceQuery, cfg.Keyspace)).Exec(); err != nil { session.Close() panic(err) } // Drop table if reset is true if cfg.Reset { if err = session.Query(fmt.Sprintf(dropQuery, cfg.Keyspace, cfg.Table)).Exec(); err != nil { session.Close() panic(err) } } // Create the storage store := &Storage{ session: session, tableName: cfg.Table, selectQuery: fmt.Sprintf(selectQuery, cfg.Keyspace, cfg.Table), insertQuery: fmt.Sprintf(insertQuery, cfg.Keyspace, cfg.Table), deleteQuery: fmt.Sprintf(deleteQuery, cfg.Keyspace, cfg.Table), resetQuery: fmt.Sprintf(resetQuery, cfg.Keyspace, cfg.Table), } // Create table if not exists if err = store.createTableIfNotExists(cfg.Keyspace); err != nil { session.Close() panic(err) } // Check schema store.checkSchema(cfg.Keyspace) return store // Return storage } func (s *Storage) createTableIfNotExists(keyspace string) error { // Create table if not exists query := fmt.Sprintf(createTableQuery, keyspace, s.tableName) if err := s.session.Query(query).Exec(); err != nil { return err } return nil } func (s *Storage) checkSchema(keyspace string) { // Check schema for value column type (should be blob) var dataType string query := fmt.Sprintf(checkSchemaQuery, keyspace, s.tableName) if err := s.session.Query(query).Scan(&dataType); err != nil { panic(err) } if dataType != "blob" { panic(fmt.Errorf(checkSchemaErrorMsg.Error(), dataType)) } } // Get retrieves a value by key func (s *Storage) Get(key string) ([]byte, error) { var value []byte if err := s.session.Query(s.selectQuery, key).Scan(&value); err != nil { if errors.Is(err, gocql.ErrNotFound) { return nil, nil } return nil, err } return value, nil } // Set sets a value by key func (s *Storage) Set(key string, value []byte, expire time.Duration) error { var expiration int if expire != 0 { expiration = int(expire.Round(time.Second).Seconds()) } return s.session.Query(s.insertQuery, key, value, expiration).Exec() } // Delete removes a value by key func (s *Storage) Delete(key string) error { return s.session.Query(s.deleteQuery, key).Exec() } // Reset resets all values func (s *Storage) Reset() error { return s.session.Query(s.resetQuery).Exec() } // Close closes the storage func (s *Storage) Close() error { s.session.Close() return nil } // Conn returns the underlying gocql session func (s *Storage) Conn() *gocql.Session { return s.session }