Create Cassandra db schema on session initialization

.github/workflows/ci-e2e-cassandra.yml

@@ -22,14 +22,19 @@ jobs:
       fail-fast: false
       matrix:
         jaeger-version: [v1, v2]
+        skip-apply-schema: [true, false]
         version:
         - distribution: cassandra
           major: 4.x
           schema: v004
         - distribution: cassandra
           major: 5.x
           schema: v004
-    name: ${{ matrix.version.distribution }} ${{ matrix.version.major }} ${{ matrix.jaeger-version }}
+        exclude:
+        # Exclude v1 as create schema on fly is available for v2 only
+        - jaeger-version: v1
+          skip-apply-schema: true
+    name: ${{ matrix.version.distribution }} ${{ matrix.version.major }} ${{ matrix.jaeger-version }} ${{ matrix.skip-apply-schema == true && 'auto' || 'manual' }}
     steps:
     - name: Harden Runner
       uses: step-security/harden-runner@91182cccc01eb5e619899d80e4e971d6181294a7 # v2.10.1
@@ -45,6 +50,8 @@ jobs:
     - name: Run cassandra integration tests
       id: test-execution
       run: bash scripts/cassandra-integration-test.sh ${{ matrix.version.major }} ${{ matrix.version.schema }} ${{ matrix.jaeger-version }}
+      env:
+        CASSANDRA_CREATE_SCHEMA: ${{ matrix.skip-apply-schema }}
 
     - name: Upload coverage to codecov
       uses: ./.github/actions/upload-codecov

cmd/jaeger/config-cassandra.yaml

@@ -34,6 +34,7 @@ extensions:
         cassandra:
           schema: 
             keyspace: "jaeger_v1_dc1"
+            create: "${env:CASSANDRA_CREATE_SCHEMA:-false}"
           connection:
             auth: 
               basic:
@@ -44,7 +45,8 @@ extensions:
       another_storage:
         cassandra:
           schema: 
-            keyspace: "jaeger_v1_dc1"
+            keyspace: "jaeger_v1_dc1_archive"
+            create: "${env:CASSANDRA_CREATE_SCHEMA:-false}"
           connection: 
             auth:
               basic: 

cmd/jaeger/internal/integration/cassandra_test.go

@@ -11,6 +11,7 @@ import (
 
 func TestCassandraStorage(t *testing.T) {
 	integration.SkipUnlessEnv(t, "cassandra")
+
 	s := &E2EStorageIntegration{
 		ConfigFile: "../../config-cassandra.yaml",
 		StorageIntegration: integration.StorageIntegration{

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,52 @@
 	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 {
+	if !c.Schema.CreateSchema {
+		return nil
+	}
+
+	cfg := *c // clone because we need to connect without specifying a keyspace
+	cfg.Schema.Keyspace = ""
+
+	session, err := createSession(&cfg)
+	if err != nil {
+		return err
+	}
+
+	sc := schemaCreator{
+		session: session,
+		schema:  c.Schema,
+	}
+
+	return sc.createSchemaIfNotPresent()
+}
+
+// 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 +283,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,145 @@
+// 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
+}
+
+type schemaCreator struct {
+	session cassandra.Session
+	schema  Schema
+}
+
+func (sc *schemaCreator) constructTemplateParams() templateParams {
+	return templateParams{
+		Keyspace:                  sc.schema.Keyspace,
+		Replication:               fmt.Sprintf("{'class': 'NetworkTopologyStrategy', 'replication_factor': '%v' }", sc.schema.ReplicationFactor),
+		CompactionWindowInMinutes: int64(sc.schema.CompactionWindow / time.Minute),
+		TraceTTLInSeconds:         int64(sc.schema.TraceTTL / time.Second),
+		DependenciesTTLInSeconds:  int64(sc.schema.DependenciesTTL / time.Second),
+	}
+}
+
+func (*schemaCreator) 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 (*schemaCreator) 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]
+		}
+
+		trimmedLine := bytes.TrimSpace(line)
+
+		if len(trimmedLine) == 0 {
+			continue
+		}
+
+		extractedLines = append(extractedLines, trimmedLine)
+	}
+
+	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 (sc *schemaCreator) getCassandraQueriesFromQueryStrings(queries []string) []cassandra.Query {
+	var casQueries []cassandra.Query
+
+	for _, query := range queries {
+		casQueries = append(casQueries, sc.session.Query(query))
+	}
+
+	return casQueries
+}
+
+func (sc *schemaCreator) contructSchemaQueries() ([]cassandra.Query, error) {
+	params := sc.constructTemplateParams()
+
+	queryFile, err := sc.getQueryFileAsBytes(`v004-go-tmpl.cql.tmpl`, params)
+	if err != nil {
+		return nil, err
+	}
+
+	queryStrings, err := sc.getQueriesFromBytes(queryFile)
+	if err != nil {
+		return nil, err
+	}
+
+	casQueries := sc.getCassandraQueriesFromQueryStrings(queryStrings)
+
+	return casQueries, nil
+}
+
+func (sc *schemaCreator) createSchemaIfNotPresent() error {
+	casQueries, err := sc.contructSchemaQueries()
+	if err != nil {
+		return err
+	}
+
+	for _, query := range casQueries {
+		if err := query.Exec(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}