Create Cassandra db schema on session initialization

pkg/cassandra/config/config.go

@@ -58,6 +58,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,cassandraTTLValidation"`
+	// DependenciesTTL is Time To Live (TTL) for dependencies data. Should at least be 1 second
+	DependenciesTTL time.Duration `mapstructure:"dependencies_ttl" valid:"optional,cassandraTTLValidation"`
+	// 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,isDurationGreaterThanOrEqualMinute"`
 }
 
 type Query struct {
@@ -86,7 +98,13 @@
 func DefaultConfiguration() Configuration {
 	return Configuration{
 		Schema: Schema{
-			Keyspace: "jaeger_v1_test",
+			CreateSchema:      false,
+			Keyspace:          "jaeger_v1_dc1",
+			Datacenter:        "test",
+			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 +124,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,8 +173,31 @@
 	NewSession() (cassandra.Session, error)
 }
 
+func (c *Configuration) newSessionPrerequisites() error {
+	if !c.Schema.CreateSchema {
+		return nil
+	}
+	cluster, err := c.NewCluster()
+
+	cluster.Keyspace = ""
+
+	session, err := cluster.CreateSession()
+	if err != nil {
+		return err
+	}
+
+	wSession := gocqlw.WrapCQLSession(session)
+	defer wSession.Close()
+
+	return generateSchemaIfNotPresent(wSession, &c.Schema)
+}
+
 // NewSession creates a new Cassandra session
 func (c *Configuration) NewSession() (cassandra.Session, error) {
+	if err := c.newSessionPrerequisites(); err != nil {
+		return nil, err
+	}
+
 	cluster, err := c.NewCluster()
 	if err != nil {
 		return nil, err
@@ -211,6 +273,23 @@
 }
 
 func (c *Configuration) Validate() error {
+	govalidator.CustomTypeTagMap.Set("cassandraTTLValidation", func(i interface{}, _ interface{}) bool {
+		duration, ok := i.(time.Duration)
+		if !ok {
+			return false
+		}
+		// Either set to no-expiry using 0 or if set should be greater than or equal to second
+		return duration == 0 || duration >= time.Second
+	})
+
+	govalidator.CustomTypeTagMap.Set("isDurationGreaterThanOrEqualMinute", func(i interface{}, _ interface{}) bool {
+		duration, ok := i.(time.Duration)
+		if !ok {
+			return false
+		}
+		return duration >= time.Minute
+	})
+
 	_, err := govalidator.ValidateStruct(c)
 	return err
 }

pkg/cassandra/config/config_test.go

@@ -5,6 +5,7 @@ package config
 
 import (
 	"testing"
+	"time"
 
 	"github.com/gocql/gocql"
 	"github.com/stretchr/testify/assert"
@@ -94,3 +95,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,143 @@
+// 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 {
+	Schema
+	// 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, error) {
+	params := TemplateParams{
+		Schema: cfg,
+	}
+
+	params.Replication = fmt.Sprintf("{'class': 'NetworkTopologyStrategy', 'replication_factor': '%v' }", params.ReplicationFactor)
+	params.CompactionWindowInMinutes = int64(params.CompactionWindow / time.Minute)
+	params.TraceTTLInSeconds = int64(params.TraceTTL / time.Second)
+	params.DependenciesTTLInSeconds = int64(params.DependenciesTTL / time.Second)
+
+	return params, nil
+}
+
+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, err := constructTemplateParams(*cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	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
+}