Create Cassandra db schema on session initialization

pkg/cassandra/config/config.go

@@ -6,6 +6,7 @@
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -58,6 +59,18 @@
 	// while connecting to the Cassandra Cluster. This is useful for connecting to clusters, like Azure Cosmos DB,
 	// that do not support SnappyCompression.
 	DisableCompression bool `mapstructure:"disable_compression"`
+	// CreateSchema tells if the schema ahould be created during session initialization based on the configs provided
+	CreateSchema bool `mapstructure:"create" valid:"optional"`
+	// Datacenter is the name for network topology
+	Datacenter string `mapstructure:"datacenter" valid:"optional"`
+	// TraceTTL is Time To Live (TTL) for the trace data. Should at least be 1 second
+	TraceTTL time.Duration `mapstructure:"trace_ttl" valid:"optional"`
+	// DependenciesTTL is Time To Live (TTL) for dependencies data. Should at least be 1 second
+	DependenciesTTL time.Duration `mapstructure:"dependencies_ttl" valid:"optional"`
+	// Replication factor for the db
+	ReplicationFactor int `mapstructure:"replication_factor" valid:"optional"`
+	// CompactionWindow of format tells the compaction window of the db. Should atleast be 1 minute
+	CompactionWindow time.Duration `mapstructure:"compaction_window" valid:"optional"`
 }
 
 type Query struct {
@@ -86,7 +99,13 @@
 func DefaultConfiguration() Configuration {
 	return Configuration{
 		Schema: Schema{
-			Keyspace: "jaeger_v1_test",
+			CreateSchema:      false,
+			Keyspace:          "jaeger_dc1",
+			Datacenter:        "dc1",
+			TraceTTL:          2 * 24 * time.Hour,
+			DependenciesTTL:   2 * 24 * time.Hour,
+			ReplicationFactor: 1,
+			CompactionWindow:  time.Minute,
 		},
 		Connection: Connection{
 			Servers:            []string{"127.0.0.1"},
@@ -106,6 +125,27 @@
 	if c.Schema.Keyspace == "" {
 		c.Schema.Keyspace = source.Schema.Keyspace
 	}
+
+	if c.Schema.Datacenter == "" {
+		c.Schema.Datacenter = source.Schema.Datacenter
+	}
+
+	if c.Schema.TraceTTL == 0 {
+		c.Schema.TraceTTL = source.Schema.TraceTTL
+	}
+
+	if c.Schema.DependenciesTTL == 0 {
+		c.Schema.DependenciesTTL = source.Schema.DependenciesTTL
+	}
+
+	if c.Schema.ReplicationFactor == 0 {
+		c.Schema.ReplicationFactor = source.Schema.ReplicationFactor
+	}
+
+	if c.Schema.CompactionWindow == 0 {
+		c.Schema.CompactionWindow = source.Schema.CompactionWindow
+	}
+
 	if c.Connection.ConnectionsPerHost == 0 {
 		c.Connection.ConnectionsPerHost = source.Connection.ConnectionsPerHost
 	}
@@ -134,19 +174,48 @@
 	NewSession() (cassandra.Session, error)
 }
 
-// NewSession creates a new Cassandra session
-func (c *Configuration) NewSession() (cassandra.Session, error) {
+// createSession creates session from a configuration
+func createSession(c *Configuration) (cassandra.Session, error) {
 	cluster, err := c.NewCluster()
 	if err != nil {
 		return nil, err
 	}
+
 	session, err := cluster.CreateSession()
 	if err != nil {
 		return nil, err
 	}
+
 	return gocqlw.WrapCQLSession(session), nil
 }
 
+// newSessionPrerequisites creates tables and types before creating a session
+func (c *Configuration) newSessionPrerequisites() error {
+	cfg := *c // clone because we need to connect without specifying a keyspace
+	cfg.Schema.Keyspace = ""
+	if !c.Schema.CreateSchema {
+		return nil
+	}
+
+	c.Schema.Keyspace = ""
+
+	session, err := createSession(c)
+	if err != nil {
+		return err
+	}
+
+	return generateSchemaIfNotPresent(session, &c.Schema)
+}
+
+// NewSession creates a new Cassandra session
+func (c *Configuration) NewSession() (cassandra.Session, error) {
+	if err := c.newSessionPrerequisites(); err != nil {
+		return nil, err
+	}
+
+	return createSession(c)
+}
+
 // NewCluster creates a new gocql cluster from the configuration
 func (c *Configuration) NewCluster() (*gocql.ClusterConfig, error) {
 	cluster := gocql.NewCluster(c.Connection.Servers...)
@@ -210,7 +279,27 @@
 	return fmt.Sprintf("%+v", *c)
 }
 
+func isValidTTL(duration time.Duration) bool {
+	return duration == 0 || duration >= time.Second
+}
+
 func (c *Configuration) Validate() error {
 	_, err := govalidator.ValidateStruct(c)
-	return err
+	if err != nil {
+		return err
+	}
+
+	if !isValidTTL(c.Schema.TraceTTL) {
+		return errors.New("trace_ttl can either be 0 or greater than or equal to 1 second")
+	}
+
+	if !isValidTTL(c.Schema.DependenciesTTL) {
+		return errors.New("dependencies_ttl can either be 0 or greater than or equal to 1 second")
+	}
+
+	if c.Schema.CompactionWindow < time.Minute {
+		return errors.New("compaction_window should at least be 1 minute")
+	}
+
+	return nil
 }

pkg/cassandra/config/config_test.go

@@ -5,6 +5,7 @@ package config
 
 import (
 	"testing"
+	"time"
 
 	"github.com/gocql/gocql"
 	"github.com/stretchr/testify/assert"
@@ -43,6 +44,9 @@ func TestValidate_DoesNotReturnErrorWhenRequiredFieldsSet(t *testing.T) {
 		Connection: Connection{
 			Servers: []string{"localhost:9200"},
 		},
+		Schema: Schema{
+			CompactionWindow: time.Minute,
+		},
 	}
 
 	err := cfg.Validate()
@@ -94,3 +98,18 @@ func TestToString(t *testing.T) {
 	s := cfg.String()
 	assert.Contains(t, s, "Keyspace:test")
 }
+
+func TestConfigSchemaValidation(t *testing.T) {
+	cfg := DefaultConfiguration()
+	err := cfg.Validate()
+	require.NoError(t, err)
+
+	cfg.Schema.TraceTTL = time.Millisecond
+	err = cfg.Validate()
+	require.Error(t, err)
+
+	cfg.Schema.TraceTTL = time.Second
+	cfg.Schema.CompactionWindow = time.Minute - 1
+	err = cfg.Validate()
+	require.Error(t, err)
+}

pkg/cassandra/config/schema.go

@@ -0,0 +1,138 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package config
+
+import (
+	"bytes"
+	"embed"
+	"errors"
+	"fmt"
+	"text/template"
+	"time"
+
+	"github.com/jaegertracing/jaeger/pkg/cassandra"
+)
+
+//go:embed v004-go-tmpl.cql.tmpl
+var schemaFile embed.FS
+
+type TemplateParams struct {
+	// Keyspace in which tables and types will be created for storage
+	Keyspace string
+	// Replication is the replication strategy used. Ex: "{'class': 'NetworkTopologyStrategy', 'replication_factor': '1' }"
+	Replication string
+	// CompactionWindowInMinutes is constructed from CompactionWindow for using in template
+	CompactionWindowInMinutes int64
+	// TraceTTLInSeconds is constructed from TraceTTL for using in template
+	TraceTTLInSeconds int64
+	// DependenciesTTLInSeconds is constructed from DependenciesTTL for using in template
+	DependenciesTTLInSeconds int64
+}
+
+func constructTemplateParams(cfg Schema) TemplateParams {
+	return TemplateParams{
+		Keyspace:                  cfg.Keyspace,
+		Replication:               fmt.Sprintf("{'class': 'NetworkTopologyStrategy', 'replication_factor': '%v' }", cfg.ReplicationFactor),
+		CompactionWindowInMinutes: int64(cfg.CompactionWindow / time.Minute),
+		TraceTTLInSeconds:         int64(cfg.TraceTTL / time.Second),
+		DependenciesTTLInSeconds:  int64(cfg.DependenciesTTL / time.Second),
+	}
+}
+
+func getQueryFileAsBytes(fileName string, params TemplateParams) ([]byte, error) {
+	tmpl, err := template.ParseFS(schemaFile, fileName)
+	if err != nil {
+		return nil, err
+	}
+
+	var result bytes.Buffer
+	err = tmpl.Execute(&result, params)
+	if err != nil {
+		return nil, err
+	}
+
+	return result.Bytes(), nil
+}
+
+func getQueriesFromBytes(queryFile []byte) ([]string, error) {
+	lines := bytes.Split(queryFile, []byte("\n"))
+
+	var extractedLines [][]byte
+
+	for _, line := range lines {
+		// Remove any comments, if at the end of the line
+		commentIndex := bytes.Index(line, []byte(`--`))
+		if commentIndex != -1 {
+			// remove everything after comment
+			line = line[0:commentIndex]
+		}
+
+		if len(line) == 0 {
+			continue
+		}
+
+		extractedLines = append(extractedLines, bytes.TrimSpace(line))
+	}
+
+	var queries []string
+
+	// Construct individual queries strings
+	var queryString string
+	for _, line := range extractedLines {
+		queryString += string(line) + "\n"
+		if bytes.HasSuffix(line, []byte(";")) {
+			queries = append(queries, queryString)
+			queryString = ""
+		}
+	}
+
+	if len(queryString) > 0 {
+		return nil, errors.New(`invalid template`)
+	}
+
+	return queries, nil
+}
+
+func getCassandraQueriesFromQueryStrings(session cassandra.Session, queries []string) []cassandra.Query {
+	var casQueries []cassandra.Query
+
+	for _, query := range queries {
+		casQueries = append(casQueries, session.Query(query))
+	}
+
+	return casQueries
+}
+
+func contructSchemaQueries(session cassandra.Session, cfg *Schema) ([]cassandra.Query, error) {
+	params := constructTemplateParams(*cfg)
+
+	queryFile, err := getQueryFileAsBytes(`v004-go-tmpl.cql.tmpl`, params)
+	if err != nil {
+		return nil, err
+	}
+
+	queryStrings, err := getQueriesFromBytes(queryFile)
+	if err != nil {
+		return nil, err
+	}
+
+	casQueries := getCassandraQueriesFromQueryStrings(session, queryStrings)
+
+	return casQueries, nil
+}
+
+func generateSchemaIfNotPresent(session cassandra.Session, cfg *Schema) error {
+	casQueries, err := contructSchemaQueries(session, cfg)
+	if err != nil {
+		return err
+	}
+
+	for _, query := range casQueries {
+		if err := query.Exec(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

pkg/cassandra/config/schema_test.go

@@ -0,0 +1,55 @@
+package config
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestQueryGenerationFromBytes(t *testing.T) {
+	queriesAsString := `
+query1 -- comment (this should be removed)
+query1-continue
+query1-finished; --
+
+
+query2;
+query-3 query-3-continue query-3-finished;
+`
+	expGeneratedQueries := []string{
+		`query1
+query1-continue
+query1-finished;
+`,
+		`query2;
+`,
+		`query-3 query-3-continue query-3-finished;
+`,
+	}
+
+	queriesAsBytes := []byte(queriesAsString)
+	queries, err := getQueriesFromBytes(queriesAsBytes)
+	require.NoError(t, err)
+
+	require.Equal(t, len(expGeneratedQueries), len(queries))
+
+	for i := range len(expGeneratedQueries) {
+		require.Equal(t, expGeneratedQueries[i], queries[i])
+	}
+}
+
+func TestInvalidQueryTemplate(t *testing.T) {
+	queriesAsString := `
+	query1 -- comment (this should be removed)
+	query1-continue
+	query1-finished; --
+
+
+	query2;
+	query-3 query-3-continue query-3-finished -- missing semicolon
+	`
+
+	queriesAsBytes := []byte(queriesAsString)
+	_, err := getQueriesFromBytes(queriesAsBytes)
+	require.ErrorContains(t, err, "invalid template")
+}