Merge pull request #31 from PostHog/feat/ducklake-compatibility

Add DuckLake compatibility layer for PostgreSQL DDL

Author: EDsCODE

server/catalog.go

@@ -368,6 +368,34 @@ var (
 	// Two patterns: one for version() with existing AS alias, one for version() without
 	versionFuncWithAliasRegex = regexp.MustCompile(`(?i)\bversion\s*\(\s*\)(\s+AS\s+)`)
 	versionFuncRegex          = regexp.MustCompile(`(?i)\bversion\s*\(\s*\)`)
+
+	// DuckLake compatibility: Strip unsupported constraints from CREATE TABLE
+	// PRIMARY KEY constraint (inline or table-level)
+	primaryKeyInlineRegex = regexp.MustCompile(`(?i)\s+PRIMARY\s+KEY`)
+	// UNIQUE constraint (inline)
+	uniqueInlineRegex = regexp.MustCompile(`(?i)\s+UNIQUE`)
+	// REFERENCES / FOREIGN KEY (inline) with optional ON DELETE/UPDATE clauses
+	referencesRegex = regexp.MustCompile(`(?i)\s+REFERENCES\s+\w+(?:\.\w+)?(?:\s*\([^)]+\))?(?:\s+ON\s+(?:DELETE|UPDATE)\s+(?:CASCADE|SET\s+NULL|SET\s+DEFAULT|RESTRICT|NO\s+ACTION))*`)
+	// CHECK constraint (inline) - handles one level of nested parentheses
+	checkConstraintRegex = regexp.MustCompile(`(?i)\s+CHECK\s*\((?:[^()]*|\([^()]*\))*\)`)
+	// SERIAL types -> INTEGER types
+	serialRegex      = regexp.MustCompile(`(?i)\bSERIAL\b`)
+	bigserialRegex   = regexp.MustCompile(`(?i)\bBIGSERIAL\b`)
+	smallserialRegex = regexp.MustCompile(`(?i)\bSMALLSERIAL\b`)
+	// DEFAULT now()/current_timestamp (not supported in DuckLake)
+	defaultNowRegex = regexp.MustCompile(`(?i)\s+DEFAULT\s+(?:now\s*\(\s*\)|current_timestamp|CURRENT_TIMESTAMP)`)
+	// GENERATED columns
+	generatedRegex = regexp.MustCompile(`(?i)\s+GENERATED\s+(?:ALWAYS|BY\s+DEFAULT)\s+AS\s+(?:IDENTITY(?:\s*\([^)]*\))?|[^,)]+)`)
+	// Table-level PRIMARY KEY constraint: PRIMARY KEY (col1, col2)
+	tablePrimaryKeyRegex = regexp.MustCompile(`(?i),?\s*PRIMARY\s+KEY\s*\([^)]+\)`)
+	// Table-level UNIQUE constraint: UNIQUE (col1, col2)
+	tableUniqueRegex = regexp.MustCompile(`(?i),?\s*UNIQUE\s*\([^)]+\)`)
+	// Table-level FOREIGN KEY constraint
+	tableForeignKeyRegex = regexp.MustCompile(`(?i),?\s*FOREIGN\s+KEY\s*\([^)]+\)\s*REFERENCES\s+\w+(?:\.\w+)?\s*\([^)]+\)(?:\s+ON\s+(?:DELETE|UPDATE)\s+(?:CASCADE|SET\s+NULL|SET\s+DEFAULT|RESTRICT|NO\s+ACTION))*`)
+	// Table-level CHECK constraint - handles one level of nested parentheses
+	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))*`)
 )
 
 // PostgreSQL-specific SET parameters that DuckDB doesn't support.
@@ -568,3 +596,73 @@ func rewritePgCatalogQuery(query string) string {
 	return query
 }
 
+// rewriteForDuckLake rewrites PostgreSQL DDL to be compatible with DuckLake limitations.
+// DuckLake does not support: PRIMARY KEY, UNIQUE, FOREIGN KEY, CHECK constraints,
+// SERIAL types, DEFAULT now(), GENERATED columns, or indexes.
+// This function strips these unsupported features so DDL can execute.
+func rewriteForDuckLake(query string) string {
+	upperQuery := strings.ToUpper(strings.TrimSpace(query))
+
+	// 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") {
+		return query
+	}
+
+	// Strip named constraints first (CONSTRAINT name PRIMARY KEY/UNIQUE/FOREIGN KEY/CHECK)
+	query = constraintNameRegex.ReplaceAllString(query, "")
+
+	// Strip table-level constraints
+	query = tablePrimaryKeyRegex.ReplaceAllString(query, "")
+	query = tableUniqueRegex.ReplaceAllString(query, "")
+	query = tableForeignKeyRegex.ReplaceAllString(query, "")
+	query = tableCheckRegex.ReplaceAllString(query, "")
+
+	// Strip inline constraints
+	query = primaryKeyInlineRegex.ReplaceAllString(query, "")
+	query = uniqueInlineRegex.ReplaceAllString(query, "")
+	query = referencesRegex.ReplaceAllString(query, "")
+	query = checkConstraintRegex.ReplaceAllString(query, "")
+
+	// Convert SERIAL types to INTEGER types
+	query = smallserialRegex.ReplaceAllString(query, "SMALLINT")
+	query = bigserialRegex.ReplaceAllString(query, "BIGINT")
+	query = serialRegex.ReplaceAllString(query, "INTEGER")
+
+	// Strip DEFAULT now()/current_timestamp (DuckLake only allows literal defaults)
+	query = defaultNowRegex.ReplaceAllString(query, "")
+
+	// Strip GENERATED columns
+	query = generatedRegex.ReplaceAllString(query, "")
+
+	return query
+}
+
+// isNoOpCommand returns true if the command should be acknowledged but not executed.
+// These are PostgreSQL features that DuckLake doesn't support.
+func isNoOpCommand(cmdType string) bool {
+	switch cmdType {
+	case "CREATE INDEX", "DROP INDEX", "REINDEX",
+		"CLUSTER", "VACUUM", "ANALYZE",
+		"GRANT", "REVOKE", "COMMENT",
+		"REFRESH", // REFRESH MATERIALIZED VIEW
+		"ALTER TABLE ADD CONSTRAINT": // Constraints not supported in DuckLake
+		return true
+	default:
+		return false
+	}
+}
+
+// getNoOpCommandTag returns the command tag to send for a no-op command.
+// Some internal command types need to be mapped to standard PostgreSQL tags.
+func getNoOpCommandTag(cmdType string) string {
+	switch cmdType {
+	case "ALTER TABLE ADD CONSTRAINT":
+		return "ALTER TABLE" // PostgreSQL returns "ALTER TABLE" for constraint operations
+	default:
+		return cmdType
+	}
+}
+

server/conn.go

@@ -292,6 +292,9 @@ func (c *clientConn) handleQuery(body []byte) error {
 	// Rewrite pg_catalog function calls for compatibility
 	query = rewritePgCatalogQuery(query)
 
+	// Rewrite CREATE TABLE for DuckLake compatibility (strip unsupported constraints)
+	query = rewriteForDuckLake(query)
+
 	// Determine command type for proper response
 	upperQuery := strings.ToUpper(query)
 	cmdType := c.getCommandType(upperQuery)
@@ -301,6 +304,15 @@ func (c *clientConn) handleQuery(body []byte) error {
 		return c.handleCopy(query, upperQuery)
 	}
 
+	// Handle no-op commands (CREATE INDEX, VACUUM, etc.) - DuckLake doesn't support these
+	if isNoOpCommand(cmdType) {
+		log.Printf("[%s] No-op command (DuckLake limitation): %s", c.username, query)
+		writeCommandComplete(c.writer, getNoOpCommandTag(cmdType))
+		writeReadyForQuery(c.writer, c.txStatus)
+		c.writer.Flush()
+		return nil
+	}
+
 	// For non-SELECT queries, use Exec
 	if cmdType != "SELECT" {
 		// Handle nested BEGIN: PostgreSQL issues a warning but continues,
@@ -505,6 +517,12 @@ func (c *clientConn) getCommandType(upperQuery string) string {
 		return "DROP SCHEMA"
 	case strings.HasPrefix(upperQuery, "DROP"):
 		return "DROP"
+	case strings.Contains(upperQuery, "ADD CONSTRAINT") ||
+		strings.Contains(upperQuery, "ADD PRIMARY KEY") ||
+		strings.Contains(upperQuery, "ADD UNIQUE") ||
+		strings.Contains(upperQuery, "ADD FOREIGN KEY") ||
+		strings.Contains(upperQuery, "ADD CHECK"):
+		return "ALTER TABLE ADD CONSTRAINT"
 	case strings.HasPrefix(upperQuery, "ALTER"):
 		return "ALTER TABLE"
 	case strings.HasPrefix(upperQuery, "TRUNCATE"):
@@ -1027,6 +1045,9 @@ func (c *clientConn) handleParse(body []byte) {
 	// Rewrite pg_catalog function calls for compatibility (same as simple query protocol)
 	rewrittenQuery := rewritePgCatalogQuery(query)
 
+	// Rewrite CREATE TABLE for DuckLake compatibility (strip unsupported constraints)
+	rewrittenQuery = rewriteForDuckLake(rewrittenQuery)
+
 	// Convert PostgreSQL $1, $2 placeholders to ? for database/sql
 	convertedQuery, numParams := convertPlaceholders(rewrittenQuery)
 
@@ -1321,6 +1342,13 @@ func (c *clientConn) handleExecute(body []byte) {
 		return
 	}
 
+	// Handle no-op commands (CREATE INDEX, VACUUM, etc.) - DuckLake doesn't support these
+	if isNoOpCommand(cmdType) {
+		log.Printf("[%s] No-op command (DuckLake limitation): %s", c.username, p.stmt.query)
+		writeCommandComplete(c.writer, getNoOpCommandTag(cmdType))
+		return
+	}
+
 	if !returnsResults {
 		// Handle nested BEGIN: PostgreSQL issues a warning but continues,
 		// while DuckDB throws an error. Match PostgreSQL behavior.