Merge pull request #15 from PostHog/add-if-not-exists-option

PostHog · Oct 22, 2024 · 9e690c3 · 9e690c3
2 parents f4f5123 + 5cd8b5f
commit 9e690c3
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -53,6 +53,9 @@ Here's a quick rundown of the commands available:
 # Dump database schema to file _without_ kafka or materialized view tables
 ./synch dump-schema --no-kafkas --no-mat-views <clickhouse_url> <file> <database>
 
+# Dump database schema to file _with_ IF NOT EXISTS in CREATE TABLE statements
+./synch dump-schema --if-not-exists <clickhouse_url> <file> <database>
+
 # <clickhouse_url> here looks like `"clickhouse://user:password@host:port"`
 
 # Synchronize a table across clusters

diff --git a/main.go b/main.go
@@ -82,6 +82,7 @@ func main() {
 		noMatViews   = false
 		onlyKafkas   = false
 		onlyMatViews = false
+		ifNotExists  = false
 	)
 
 	dumpSchemaCmd := &cobra.Command{
@@ -110,6 +111,7 @@ func main() {
 				NoMatViews:   noMatViews,
 				OnlyKafkas:   onlyKafkas,
 				OnlyMatViews: onlyMatViews,
+				IfNotExists:  ifNotExists,
 			}
 
 			err = Write(&opts)
@@ -128,6 +130,7 @@ func main() {
 	dumpSchemaCmd.Flags().BoolVar(&noMatViews, "no-mat-views", false, "Don't dump materialized views")
 	dumpSchemaCmd.Flags().BoolVar(&onlyKafkas, "only-kafkas", false, "Dump only Kafka tables")
 	dumpSchemaCmd.Flags().BoolVar(&onlyMatViews, "only-mat-views", false, "Dump only materialized views")
+	dumpSchemaCmd.Flags().BoolVar(&ifNotExists, "if-not-exists", false, "Add IF NOT EXISTS to CREATE TABLE statements")
 	cmd.AddCommand(dumpSchemaCmd)
 
 	var (

diff --git a/schemas.go b/schemas.go
@@ -20,6 +20,7 @@ type Options struct {
 	OnlyKafkas     bool
 	OnlyMatViews   bool
 	Apply          bool
+	IfNotExists    bool
 }
 
 var (
@@ -65,7 +66,7 @@ func Compare(opts *Options) error {
 			if !includes(tables2, tableName) {
 				fmt.Printf("-- Table '%s.%s' is missing in the destination\n", dbName, tableName)
 				if !opts.TableNamesOnly {
-					tableCreateStmt, err := fetchTableCreateStmt(opts.DB, dbName, tableName)
+					tableCreateStmt, err := fetchTableCreateStmt(opts.DB, dbName, tableName, opts.IfNotExists)
 					if err != nil {
 						return err
 					}
@@ -147,7 +148,7 @@ func Write(opts *Options) error {
 			tables = append(tables, newTables...)
 		}
 		for _, tableName := range tables {
-			tableCreateStmt, err := fetchTableCreateStmt(opts.DB, dbName, tableName)
+			tableCreateStmt, err := fetchTableCreateStmt(opts.DB, dbName, tableName, opts.IfNotExists)
 			if err != nil {
 				return err
 			}
@@ -243,14 +244,21 @@ func dbCreateStmt(db *sql.DB, dbName string) (string, error) {
 	return strings.Replace(createStmt, "CREATE DATABASE", "CREATE DATABASE IF NOT EXISTS", 1), nil
 }
 
-func fetchTableCreateStmt(db *sql.DB, dbName string, tableName string) (string, error) {
+func fetchTableCreateStmt(db *sql.DB, dbName string, tableName string, ifNotExists bool) (string, error) {
 	var createStmt string
 	queryStmt := fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s` FORMAT PrettySpaceNoEscapes;", dbName, tableName)
 	err := db.QueryRow(queryStmt).Scan(&createStmt)
 	if err != nil {
 		return "", fmt.Errorf("getting table '%s.%s' statement: %v", dbName, tableName, err)
 	}
 
+	if ifNotExists {
+		createStmt = strings.Replace(createStmt, "CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)
+		createStmt = strings.Replace(createStmt, "CREATE VIEW", "CREATE VIEW IF NOT EXISTS", 1)
+		createStmt = strings.Replace(createStmt, "CREATE DICTIONARY", "CREATE DICTIONARY IF NOT EXISTS", 1)
+		createStmt = strings.Replace(createStmt, "CREATE MATERIALIZED VIEW", "CREATE MATERIALIZED VIEW IF NOT EXISTS", 1)
+	}
+
 	return createStmt, nil
 }