Merge pull request #33 from PostHog/fix/ducklake-concurrent-attach-and-ddl-rewrite

Fix DuckLake concurrent attachment and DDL rewriting for ETL tools

Author: EDsCODE

server/catalog.go

@@ -396,6 +396,8 @@ var (
 	tableCheckRegex = regexp.MustCompile(`(?i),?\s*CHECK\s*\((?:[^()]*|\([^()]*\))*\)`)
 	// CONSTRAINT name prefix (for named constraints)
 	constraintNameRegex = regexp.MustCompile(`(?i),?\s*CONSTRAINT\s+\w+\s+(?:PRIMARY\s+KEY|UNIQUE|FOREIGN\s+KEY|CHECK)\s*\([^)]+\)(?:\s+REFERENCES\s+\w+(?:\.\w+)?\s*\([^)]+\))?(?:\s+ON\s+(?:DELETE|UPDATE)\s+(?:CASCADE|SET\s+NULL|SET\s+DEFAULT|RESTRICT|NO\s+ACTION))*`)
+	// Detect CREATE TABLE statements (handles multiple spaces like "CREATE  TABLE")
+	createTableDetectRegex = regexp.MustCompile(`(?i)^CREATE\s+(?:TEMPORARY\s+|TEMP\s+|UNLOGGED\s+)?TABLE\b`)
 )
 
 // PostgreSQL-specific SET parameters that DuckDB doesn't support.
@@ -603,11 +605,22 @@ func rewritePgCatalogQuery(query string) string {
 func rewriteForDuckLake(query string) string {
 	upperQuery := strings.ToUpper(strings.TrimSpace(query))
 
+	// Strip leading SQL comments (e.g., /*Fivetran*/) for detection
+	// This handles tools that prefix queries with comments
+	commentStripped := upperQuery
+	for strings.HasPrefix(commentStripped, "/*") {
+		endIdx := strings.Index(commentStripped, "*/")
+		if endIdx == -1 {
+			break
+		}
+		commentStripped = strings.TrimSpace(commentStripped[endIdx+2:])
+	}
+
 	// Only rewrite CREATE TABLE statements
-	if !strings.HasPrefix(upperQuery, "CREATE TABLE") &&
-		!strings.HasPrefix(upperQuery, "CREATE TEMPORARY TABLE") &&
-		!strings.HasPrefix(upperQuery, "CREATE TEMP TABLE") &&
-		!strings.HasPrefix(upperQuery, "CREATE UNLOGGED TABLE") {
+	// Use regex to handle multiple spaces (e.g., "CREATE  TABLE")
+	isCreateTable := createTableDetectRegex.MatchString(commentStripped)
+
+	if !isCreateTable {
 		return query
 	}
 

server/ducklake_test.go

@@ -150,6 +150,17 @@ func TestRewriteForDuckLake(t *testing.T) {
 				created_at TIMESTAMP
 			)`,
 		},
+		// ETL tool queries with comment prefix and double spaces
+		{
+			name:     "comment prefix with PRIMARY KEY",
+			input:    `/*ETL*/CREATE  TABLE "schema"."table" ( "id" CHARACTER VARYING(256) , "created_at" TIMESTAMP WITH TIME ZONE , PRIMARY KEY ("id") )`,
+			expected: `/*ETL*/CREATE  TABLE "schema"."table" ( "id" CHARACTER VARYING(256) , "created_at" TIMESTAMP WITH TIME ZONE  )`,
+		},
+		{
+			name:     "multiple comment prefixes",
+			input:    `/* comment 1 */ /* comment 2 */ CREATE TABLE test (id INTEGER PRIMARY KEY)`,
+			expected: `/* comment 1 */ /* comment 2 */ CREATE TABLE test (id INTEGER)`,
+		},
 	}
 
 	for _, tt := range tests {

server/server.go

@@ -308,6 +308,23 @@ func (s *Server) attachDuckLake(db *sql.DB) error {
 		return nil // DuckLake not configured
 	}
 
+	// Serialize DuckLake attachment to avoid race conditions
+	// Multiple connections trying to attach simultaneously can cause
+	// "database with name '__ducklake_metadata_ducklake' already exists" errors
+	s.duckLakeMu.Lock()
+	defer s.duckLakeMu.Unlock()
+
+	// Check if DuckLake catalog is already attached
+	var count int
+	err := db.QueryRow("SELECT COUNT(*) FROM duckdb_databases() WHERE database_name = 'ducklake'").Scan(&count)
+	if err == nil && count > 0 {
+		// Already attached, just set as default
+		if _, err := db.Exec("USE ducklake"); err != nil {
+			return fmt.Errorf("failed to set DuckLake as default catalog: %w", err)
+		}
+		return nil
+	}
+
 	// Create S3 secret if using object store
 	// - With explicit credentials (S3AccessKey set) or custom endpoint
 	// - With credential_chain provider (for AWS S3)
@@ -359,13 +376,9 @@ func (s *Server) attachDuckLake(db *sql.DB) error {
 //   - "config": explicit credentials (for MinIO or when you have access keys)
 //   - "credential_chain": AWS SDK credential chain (env vars, config files, instance metadata, etc.)
 //
+// Note: Caller must hold duckLakeMu to avoid race conditions.
 // See: https://duckdb.org/docs/stable/core_extensions/httpfs/s3api
 func (s *Server) createS3Secret(db *sql.DB) error {
-	// Serialize secret creation to avoid DuckDB write-write conflicts
-	// when multiple connections try to create the same secret simultaneously
-	s.duckLakeMu.Lock()
-	defer s.duckLakeMu.Unlock()
-
 	// Check if secret already exists to avoid unnecessary creation
 	var count int
 	err := db.QueryRow("SELECT COUNT(*) FROM duckdb_secrets() WHERE name = 'ducklake_s3'").Scan(&count)