Merge pull request #42 from PostHog/feature/ducklake-integration-tests

fuziontech · web-flow · commit 056bcbc74243 · 2025-12-19T13:09:09.000-08:00
Run integration tests on DuckLake by default
diff --git a/tests/integration/docker-compose.yml b/tests/integration/docker-compose.yml
@@ -1,6 +1,5 @@
-version: '3.8'
-
 services:
+  # PostgreSQL for comparison testing (existing)
   postgres:
     image: postgres:16-alpine
     container_name: duckgres-test-postgres
@@ -20,3 +19,55 @@ services:
       retries: 10
     tmpfs:
       - /var/lib/postgresql/data
+
+  # PostgreSQL for DuckLake metadata storage
+  ducklake-metadata:
+    image: postgres:16-alpine
+    container_name: duckgres-test-ducklake-metadata
+    environment:
+      POSTGRES_USER: ducklake
+      POSTGRES_PASSWORD: ducklake
+      POSTGRES_DB: ducklake
+    ports:
+      - "35433:5432"
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U ducklake -d ducklake"]
+      interval: 2s
+      timeout: 5s
+      retries: 10
+    tmpfs:
+      - /var/lib/postgresql/data
+
+  # MinIO for DuckLake object storage
+  minio:
+    image: minio/minio:latest
+    container_name: duckgres-test-minio
+    command: server /data --console-address ":9001"
+    environment:
+      MINIO_ROOT_USER: minioadmin
+      MINIO_ROOT_PASSWORD: minioadmin
+    ports:
+      - "39000:9000"   # S3 API
+      - "39001:9001"   # Web console
+    healthcheck:
+      test: ["CMD", "mc", "ready", "local"]
+      interval: 2s
+      timeout: 5s
+      retries: 10
+    tmpfs:
+      - /data
+
+  # Creates the ducklake bucket on startup
+  minio-init:
+    image: minio/mc:latest
+    container_name: duckgres-test-minio-init
+    depends_on:
+      minio:
+        condition: service_healthy
+    entrypoint: >
+      /bin/sh -c "
+      mc alias set minio http://minio:9000 minioadmin minioadmin;
+      mc mb minio/ducklake --ignore-existing;
+      mc anonymous set download minio/ducklake;
+      echo 'Bucket ducklake created successfully';
+      "
diff --git a/tests/integration/harness.go b/tests/integration/harness.go
@@ -25,6 +25,7 @@ type TestHarness struct {
 	tmpDir      string
 	pgPort      int
 	dgPort      int
+	useDuckLake bool
 	mu          sync.Mutex
 }
 
@@ -36,21 +37,33 @@ type HarnessConfig struct {
 	SkipPostgres bool
 	// Verbose enables verbose logging
 	Verbose bool
+	// UseDuckLake enables DuckLake mode (requires ducklake-metadata and minio)
+	UseDuckLake bool
+	// DuckLakeMetadataPort is the port for the DuckLake metadata PostgreSQL (default: 35433)
+	DuckLakeMetadataPort int
+	// MinIOPort is the port for MinIO S3 API (default: 39000)
+	MinIOPort int
 }
 
 // DefaultConfig returns the default harness configuration
 func DefaultConfig() HarnessConfig {
+	// Default to DuckLake mode unless DUCKGRES_TEST_NO_DUCKLAKE is set
+	useDuckLake := os.Getenv("DUCKGRES_TEST_NO_DUCKLAKE") == ""
 	return HarnessConfig{
-		PostgresPort: 35432,
-		SkipPostgres: false,
-		Verbose:      os.Getenv("DUCKGRES_TEST_VERBOSE") != "",
+		PostgresPort:         35432,
+		SkipPostgres:         false,
+		Verbose:              os.Getenv("DUCKGRES_TEST_VERBOSE") != "",
+		UseDuckLake:          useDuckLake,
+		DuckLakeMetadataPort: 35433,
+		MinIOPort:            39000,
 	}
 }
 
 // NewTestHarness creates a new test harness
 func NewTestHarness(cfg HarnessConfig) (*TestHarness, error) {
 	h := &TestHarness{
-		pgPort: cfg.PostgresPort,
+		pgPort:      cfg.PostgresPort,
+		useDuckLake: cfg.UseDuckLake,
 	}
 
 	// Create temp directory for Duckgres
@@ -61,7 +74,7 @@ func NewTestHarness(cfg HarnessConfig) (*TestHarness, error) {
 	h.tmpDir = tmpDir
 
 	// Start Duckgres server
-	if err := h.startDuckgres(); err != nil {
+	if err := h.startDuckgres(cfg); err != nil {
 		os.RemoveAll(tmpDir)
 		return nil, fmt.Errorf("failed to start Duckgres: %w", err)
 	}
@@ -90,7 +103,7 @@ func NewTestHarness(cfg HarnessConfig) (*TestHarness, error) {
 }
 
 // startDuckgres starts the Duckgres server
-func (h *TestHarness) startDuckgres() error {
+func (h *TestHarness) startDuckgres(harnessCfg HarnessConfig) error {
 	port := findAvailablePort()
 	h.dgPort = port
 
@@ -111,6 +124,22 @@ func (h *TestHarness) startDuckgres() error {
 		Users: map[string]string{
 			"testuser": "testpass",
 		},
+		Extensions: []string{"ducklake"},
+	}
+
+	// Configure DuckLake if enabled
+	if harnessCfg.UseDuckLake {
+		cfg.DuckLake = server.DuckLakeConfig{
+			MetadataStore: fmt.Sprintf("postgres:host=127.0.0.1 port=%d user=ducklake password=ducklake dbname=ducklake", harnessCfg.DuckLakeMetadataPort),
+			ObjectStore:   "s3://ducklake/data/",
+			S3Provider:    "config",
+			S3Endpoint:    fmt.Sprintf("127.0.0.1:%d", harnessCfg.MinIOPort),
+			S3AccessKey:   "minioadmin",
+			S3SecretKey:   "minioadmin",
+			S3Region:      "us-east-1",
+			S3UseSSL:      false,
+			S3URLStyle:    "path",
+		}
 	}
 
 	srv, err := server.New(cfg)
@@ -185,7 +214,17 @@ func (h *TestHarness) connectDuckgres() error {
 }
 
 // loadFixtures loads the test schema and data into Duckgres
+// In DuckLake mode, tables are automatically created in ducklake.main
+// because the server runs "USE ducklake" to set the default catalog
 func (h *TestHarness) loadFixtures() error {
+	// In DuckLake mode, drop existing tables first since metadata persists
+	if h.useDuckLake {
+		if err := h.cleanupDuckLakeTables(); err != nil {
+			// Log but don't fail - tables might not exist
+			fmt.Printf("Warning: cleanup failed (may be OK): %v\n", err)
+		}
+	}
+
 	// Read and execute schema
 	schemaPath := filepath.Join(getTestDir(), "fixtures", "schema.sql")
 	schemaSQL, err := os.ReadFile(schemaPath)
@@ -226,6 +265,45 @@ func (h *TestHarness) loadFixtures() error {
 	return nil
 }
 
+// cleanupDuckLakeTables drops existing tables in DuckLake before loading fixtures
+func (h *TestHarness) cleanupDuckLakeTables() error {
+	// Drop views first (they depend on tables)
+	views := []string{"order_details", "user_stats", "active_users"}
+	for _, v := range views {
+		h.DuckgresDB.Exec(fmt.Sprintf("DROP VIEW IF EXISTS %s", v))
+	}
+
+	// Drop tables in reverse dependency order
+	tables := []string{
+		"test_schema.schema_test",
+		"array_test",
+		"documents",
+		"metrics",
+		"empty_table",
+		"nullable_test",
+		"json_data",
+		"events",
+		"sales",
+		"categories",
+		"order_items",
+		"orders",
+		"products",
+		"users",
+		"types_test",
+	}
+
+	for _, t := range tables {
+		if _, err := h.DuckgresDB.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", t)); err != nil {
+			// Ignore errors - table might not exist or schema might not exist
+		}
+	}
+
+	// Drop test schema
+	h.DuckgresDB.Exec("DROP SCHEMA IF EXISTS test_schema")
+
+	return nil
+}
+
 // Close shuts down the test harness
 func (h *TestHarness) Close() error {
 	h.mu.Lock()
@@ -297,6 +375,39 @@ func IsPostgresRunning(port int) bool {
 	return true
 }
 
+// IsDuckLakeInfraRunning checks if the DuckLake infrastructure (metadata postgres + minio) is running
+func IsDuckLakeInfraRunning(metadataPort, minioPort int) bool {
+	// Check DuckLake metadata PostgreSQL
+	metaConn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", metadataPort), time.Second)
+	if err != nil {
+		return false
+	}
+	metaConn.Close()
+
+	// Check MinIO
+	minioConn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", minioPort), time.Second)
+	if err != nil {
+		return false
+	}
+	minioConn.Close()
+
+	return true
+}
+
+// WaitForDuckLakeInfra waits for DuckLake infrastructure to be ready
+func WaitForDuckLakeInfra(metadataPort, minioPort int, timeout time.Duration) error {
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if IsDuckLakeInfraRunning(metadataPort, minioPort) {
+			// Give MinIO a bit more time to initialize the bucket
+			time.Sleep(500 * time.Millisecond)
+			return nil
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+	return fmt.Errorf("timeout waiting for DuckLake infrastructure (metadata:%d, minio:%d)", metadataPort, minioPort)
+}
+
 // Helper functions
 
 func findAvailablePort() int {
diff --git a/tests/integration/setup_test.go b/tests/integration/setup_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"testing"
+	"time"
 )
 
 var (
@@ -16,8 +17,10 @@ var (
 
 // TestMain sets up and tears down the test environment
 func TestMain(m *testing.M) {
-	// Check if PostgreSQL is running
-	pgPort := 35432
+	cfg := DefaultConfig()
+
+	// Check if PostgreSQL (for comparison) is running
+	pgPort := cfg.PostgresPort
 	if !IsPostgresRunning(pgPort) {
 		fmt.Println("PostgreSQL container not running. Starting it...")
 		if err := StartPostgresContainer(); err != nil {
@@ -27,8 +30,20 @@ func TestMain(m *testing.M) {
 		}
 	}
 
-	// Create test harness
-	cfg := DefaultConfig()
+	// Check and wait for DuckLake infrastructure if DuckLake mode is enabled
+	if cfg.UseDuckLake {
+		if !IsDuckLakeInfraRunning(cfg.DuckLakeMetadataPort, cfg.MinIOPort) {
+			fmt.Println("DuckLake infrastructure not running. Waiting for it...")
+			if err := WaitForDuckLakeInfra(cfg.DuckLakeMetadataPort, cfg.MinIOPort, 30*time.Second); err != nil {
+				fmt.Printf("DuckLake infrastructure not available: %v\n", err)
+				fmt.Println("Falling back to vanilla DuckDB mode (set DUCKGRES_TEST_NO_DUCKLAKE=1 to suppress this)")
+				cfg.UseDuckLake = false
+			} else {
+				fmt.Println("DuckLake infrastructure is ready")
+			}
+		}
+	}
+
 	cfg.SkipPostgres = skipPostgresCompare
 
 	var err error
diff --git a/transpiler/transform/ddl.go b/transpiler/transform/ddl.go
@@ -263,36 +263,42 @@ func (t *DDLTransform) isUnsupportedColumnConstraint(c *pg_query.Constraint) boo
 	return false
 }
 
-// isUnsupportedDefault checks if a DEFAULT expression is unsupported (e.g., now(), current_timestamp)
+// isUnsupportedDefault checks if a DEFAULT expression is unsupported by DuckLake.
+// DuckLake only supports simple numeric and string literals as defaults.
+// Returns true if the default should be stripped.
 func (t *DDLTransform) isUnsupportedDefault(expr *pg_query.Node) bool {
 	if expr == nil {
 		return false
 	}
 
-	// Check for function calls to now() or current_timestamp
+	// Check for function calls (e.g., now(), current_timestamp)
 	if funcCall := expr.GetFuncCall(); funcCall != nil {
-		if len(funcCall.Funcname) == 1 {
-			if name := funcCall.Funcname[0].GetString_(); name != nil {
-				funcName := strings.ToLower(name.Sval)
-				if funcName == "now" || funcName == "current_timestamp" {
-					return true
-				}
-			}
-		}
+		return true // All function calls are unsupported
 	}
 
 	// Check for SQLValueFunction (CURRENT_TIMESTAMP, CURRENT_DATE, etc.)
 	if svf := expr.GetSqlvalueFunction(); svf != nil {
-		switch svf.Op {
-		case pg_query.SQLValueFunctionOp_SVFOP_CURRENT_TIMESTAMP,
-			pg_query.SQLValueFunctionOp_SVFOP_CURRENT_TIMESTAMP_N,
-			pg_query.SQLValueFunctionOp_SVFOP_LOCALTIMESTAMP,
-			pg_query.SQLValueFunctionOp_SVFOP_LOCALTIMESTAMP_N:
-			return true
+		return true // All SQL value functions are unsupported
+	}
+
+	// Check for boolean constants (DEFAULT true/false)
+	// DuckLake only supports numeric and string literals
+	if typeCast := expr.GetTypeCast(); typeCast != nil {
+		return t.isUnsupportedDefault(typeCast.Arg)
+	}
+
+	// Check for A_Const nodes - only allow Integer and String
+	if aconst := expr.GetAConst(); aconst != nil {
+		switch aconst.Val.(type) {
+		case *pg_query.A_Const_Ival, *pg_query.A_Const_Fval, *pg_query.A_Const_Sval:
+			return false // These are supported
+		default:
+			return true // Booleans, NULLs, etc. are not supported
 		}
 	}
 
-	return false
+	// Column references, expressions, etc. are not supported
+	return true
 }
 
 // isConstraintCommand checks if an ALTER TABLE command is adding a constraint
