Merge pull request #268 from intelops/migrate-with-env

modified migration with ttl with ttl interval as configurable

Author: vijeyash1

cmd/cli/commands/sql.go

@@ -5,8 +5,6 @@ import (
 	"log"
 	"os"
 
-	"github.com/golang-migrate/migrate/v4"
-	cm "github.com/golang-migrate/migrate/v4/database/clickhouse"
 	_ "github.com/golang-migrate/migrate/v4/source/file"
 	"github.com/intelops/kubviz/cmd/cli/config"
 	"github.com/spf13/cobra"
@@ -47,28 +45,15 @@ migration sql -e --no`,
 
 		if executeFlag {
 			if yesFlag {
-				db, cfg, err := config.OpenClickHouseConn()
+				cfg, err := config.New()
 				if err != nil {
-					log.Fatalf("Failed to open ClickHouse connection: %v", err)
+					log.Fatalf("failed to parse the env : %v", err.Error())
+					return
 				}
-				defer db.Close()
-				driver, err := cm.WithInstance(db, &cm.Config{})
-				if err != nil {
-					log.Fatalf("Failed to create migrate driver: %v", err)
-				}
-
-				m, err := migrate.NewWithDatabaseInstance(
-					fmt.Sprintf("file://%s", cfg.SchemaPath),
-					"clickhouse",
-					driver,
-				)
-				if err != nil {
-					log.Fatalf("Clickhouse Migration initialization failed: %v", err)
-				}
-				if err := m.Up(); err != nil && err != migrate.ErrNoChange {
-					log.Fatalf("Migration failed: %v", err)
+				if err := cfg.Migrate(); err != nil {
+					log.Fatalf("failed to migrate : %v", err.Error())
+					return
 				}
-				fmt.Println("Clickhouse Migrations applied successfully!")
 			} else {
 				fmt.Println("Clickhouse Migration skipped due to --no flag.")
 			}

cmd/cli/config/config.go

@@ -1,10 +1,6 @@
 package config
 
 import (
-	"database/sql"
-	"fmt"
-
-	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/kelseyhightower/envconfig"
 )
 
@@ -14,40 +10,15 @@ type Config struct {
 	ClickHouseUsername string `envconfig:"CLICKHOUSE_USERNAME"`
 	ClickHousePassword string `envconfig:"CLICKHOUSE_PASSWORD"`
 	SchemaPath         string `envconfig:"SCHEMA_PATH" default:"/sql"`
+	TtlInterval        string `envconfig:"TTL_INTERVAL" default:"1"`
+	TtlUnit            string `envconfig:"TTL_UNIT" default:"MONTH"`
 }
 
-func OpenClickHouseConn() (*sql.DB, *Config, error) {
+func New() (*Config, error) {
 	var cfg Config
 	err := envconfig.Process("", &cfg)
 	if err != nil {
-		return nil, nil, err
-	}
-	var options clickhouse.Options
-
-	if cfg.ClickHouseUsername != "" && cfg.ClickHousePassword != "" {
-		fmt.Println("Using provided username and password")
-		options = clickhouse.Options{
-			Addr:  []string{fmt.Sprintf("%s:%d", cfg.DBAddress, cfg.DbPort)},
-			Debug: true,
-			Auth: clickhouse.Auth{
-				Username: cfg.ClickHouseUsername,
-				Password: cfg.ClickHousePassword,
-			},
-		}
-	} else {
-		fmt.Println("Using connection without username and password")
-		options = clickhouse.Options{
-			Addr: []string{fmt.Sprintf("%s:%d", cfg.DBAddress, cfg.DbPort)},
-		}
-	}
-
-	conn := clickhouse.OpenDB(&options)
-	if err := conn.Ping(); err != nil {
-		if exception, ok := err.(*clickhouse.Exception); ok {
-			return nil, nil, fmt.Errorf("[%d] %s %s", exception.Code, exception.Message, exception.StackTrace)
-		} else {
-			return nil, nil, err
-		}
+		return nil, err
 	}
-	return conn, &cfg, nil
+	return &cfg, nil
 }

cmd/cli/config/utils.go

@@ -0,0 +1,123 @@
+package config
+
+import (
+	"bytes"
+	"database/sql"
+	"fmt"
+	"html/template"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/ClickHouse/clickhouse-go/v2"
+	"github.com/golang-migrate/migrate/v4"
+	ch "github.com/golang-migrate/migrate/v4/database/clickhouse"
+	_ "github.com/golang-migrate/migrate/v4/source/file"
+)
+
+func (cfg *Config) openClickHouseConn() (*sql.DB, error) {
+
+	var options clickhouse.Options
+	if cfg.ClickHouseUsername != "" && cfg.ClickHousePassword != "" {
+		fmt.Println("Using provided username and password")
+		options = clickhouse.Options{
+			Addr:  []string{fmt.Sprintf("%s:%d", cfg.DBAddress, cfg.DbPort)},
+			Debug: true,
+			Auth: clickhouse.Auth{
+				Username: cfg.ClickHouseUsername,
+				Password: cfg.ClickHousePassword,
+			},
+		}
+
+	} else {
+		fmt.Println("Using connection without username and password")
+		options = clickhouse.Options{
+			Addr: []string{fmt.Sprintf("%s:%d", cfg.DBAddress, cfg.DbPort)},
+		}
+	}
+
+	conn := clickhouse.OpenDB(&options)
+	if err := conn.Ping(); err != nil {
+		if exception, ok := err.(*clickhouse.Exception); ok {
+			return nil, fmt.Errorf("[%d] %s %s", exception.Code, exception.Message, exception.StackTrace)
+		} else {
+			return nil, err
+		}
+	}
+	return conn, nil
+}
+
+func (cfg *Config) processSQLTemplate(filePath string) (string, error) {
+
+	file, err := os.Open(filePath)
+	if err != nil {
+		return "", err
+	}
+	data, err := io.ReadAll(file)
+	if err != nil {
+		return "", err
+	}
+	tmpl, err := template.New("sql").Parse(string(data))
+	if err != nil {
+		return "", err
+	}
+
+	params := map[string]string{
+		"TTLValue": cfg.TtlInterval,
+		"TTLUnit":  cfg.TtlUnit,
+	}
+
+	var buf bytes.Buffer
+	err = tmpl.Execute(&buf, params)
+	if err != nil {
+		return "", err
+	}
+	return buf.String(), nil
+}
+
+func (cfg *Config) Migrate() error {
+	dir := cfg.SchemaPath
+	files, err := os.ReadDir(dir)
+	if err != nil {
+		return fmt.Errorf("failed to read directory %w", err)
+	}
+
+	for _, file := range files {
+		if strings.HasSuffix(file.Name(), ".up.sql") {
+			fullPath := filepath.Join(dir, file.Name())
+			processedSQL, err := cfg.processSQLTemplate(fullPath)
+			if err != nil {
+				return fmt.Errorf("failed to process the sql template for file %s : %w", file.Name(), err)
+			}
+
+			err = os.WriteFile(fullPath, []byte(processedSQL), 0644)
+			if err != nil {
+				return fmt.Errorf("failed to write to file %s : %w", fullPath, err)
+			}
+		}
+	}
+
+	conn, err := cfg.openClickHouseConn()
+	if err != nil {
+		return fmt.Errorf("unable to create a clickhouse conection %w", err)
+	}
+
+	driver, err := ch.WithInstance(conn, &ch.Config{})
+	if err != nil {
+		return fmt.Errorf("failed to create migrate driver %w", err)
+	}
+	m, err := migrate.NewWithDatabaseInstance(
+		fmt.Sprintf("file://%s", dir),
+		"clickhouse",
+		driver,
+	)
+	if err != nil {
+		return fmt.Errorf("clickhouse migration initialization failed %w", err)
+	}
+	if err := m.Up(); err != nil && err != migrate.ErrNoChange {
+		return fmt.Errorf("migration failed %w", err)
+	}
+	fmt.Println("Clickhouse Migration applied successfully!")
+	return nil
+}

dockerfiles/migration/Dockerfile

@@ -16,7 +16,7 @@ WORKDIR /
 COPY --from=builder /workspace/migration .
 COPY --from=builder /workspace/sql /sql
 COPY --from=builder /workspace/script /script
-
+RUN chmod -R 777 /sql
 USER 65532:65532
 
 ENTRYPOINT ["/migration"]

sql/20230915101844_trivy_misconfig.down.sql renamed to sql/0000010_trivy_misconfig.down.sql

@@ -15,7 +15,7 @@ CREATE TABLE IF NOT EXISTS trivy_misconfig (
 	misconfig_severity String,
 	misconfig_status   String,
     EventTime          DateTime('UTC'),
-	ExpiryDate DateTime DEFAULT now() + INTERVAL 6 MONTH
+	ExpiryDate DateTime DEFAULT now() + INTERVAL {{.TTLValue}} {{.TTLUnit}}
 ) ENGINE = MergeTree() 
 ORDER BY ExpiryDate 
 TTL ExpiryDate;

sql/20230915101844_trivy_misconfig.up.sql renamed to sql/0000010_trivy_misconfig.up.sql

@@ -11,7 +11,7 @@ CREATE TABLE IF NOT EXISTS trivyimage (
 	vul_severity        String,
 	vul_published_date  DateTime('UTC'),
 	vul_last_modified_date DateTime('UTC'),
-	ExpiryDate DateTime DEFAULT now() + INTERVAL 6 MONTH
+	ExpiryDate DateTime DEFAULT now() + INTERVAL {{.TTLValue}} {{.TTLUnit}}
 ) ENGINE = MergeTree() 
 ORDER BY ExpiryDate 
 TTL ExpiryDate;

sql/20230915101910_trivyimage.down.sql renamed to sql/0000011_trivyimage.down.sql

@@ -6,7 +6,7 @@ CREATE TABLE IF NOT EXISTS dockerhubbuild (
 	Owner         String,
 	Event         String,
     EventTime     DateTime('UTC'),
-	ExpiryDate DateTime DEFAULT now() + INTERVAL 6 MONTH
+	ExpiryDate DateTime DEFAULT now() + INTERVAL {{.TTLValue}} {{.TTLUnit}}
 ) ENGINE = MergeTree() 
 ORDER BY ExpiryDate 
 TTL ExpiryDate;