Implement CREATE SCHEMA

.gitignore

@@ -8,6 +8,9 @@
 # Test binary, built with `go test -c`
 *.test
 
+# Test DBs
+*.sqlite3
+
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
 

exec_test.go

@@ -4,6 +4,7 @@
 	"context"
 	"database/sql"
 	"fmt"
+	"os"
 	"reflect"
 	"testing"
 	"time"
@@ -13,20 +14,30 @@
 	zetasqlite "github.com/goccy/go-zetasqlite"
 )
 
+func openTestDB(t *testing.T) *sql.DB {
+	dbPath := ":memory:"
+	if path := os.Getenv("ZETASQLITE_TEST_DB"); path != "" {
+		dbPath = path
+	}
+	db, err := sql.Open("zetasqlite", dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return db
+}
+
 func TestExec(t *testing.T) {
 	now := time.Now()
 	ctx := context.Background()
 	ctx = zetasqlite.WithCurrentTime(ctx, now)
-	db, err := sql.Open("zetasqlite", ":memory:")
-	if err != nil {
-		t.Fatal(err)
-	}
+	db := openTestDB(t)
 	defer db.Close()
-	for _, test := range []struct {
-		name        string
-		query       string
-		args        []interface{}
-		expectedErr bool
+	tests := []struct {
+		name     string
+		query    string
+		args     []interface{}
+		wantErr  bool
+		validate func(t *testing.T, db *sql.DB)
 	}{
 		{
 			name: "create table with all types",
@@ -64,6 +75,12 @@
 DROP TABLE recreate_table;
 CREATE TABLE recreate_table ( b string );
 INSERT recreate_table (b) VALUES ('hello');
+`,
+		},
+		{
+			name: "create schema",
+			query: `
+CREATE SCHEMA new_schema;
 `,
 		},
 		{
@@ -134,11 +151,47 @@
 COMMIT TRANSACTION;
 `,
 		},
-	} {
+		{
+			name:    "create schema",
+			query:   "CREATE SCHEMA test_schema",
+			wantErr: false,
+			validate: func(t *testing.T, db *sql.DB) {
+				// Create a table in the schema
+				_, err := db.Exec("CREATE TABLE test_schema.test_table (id INT64)")
+				if err != nil {
+					t.Errorf("failed to create table in schema: %v", err)
+				}
+				// Query the table with schema qualification
+				rows, err := db.Query("SELECT * FROM test_schema.test_table")
+				if err != nil {
+					t.Errorf("failed to query table in schema: %v", err)
+				}
+				defer rows.Close()
+			},
+		},
+		{
+			name:    "create schema if not exists",
+			query:   "CREATE SCHEMA IF NOT EXISTS test_schema",
+			wantErr: false,
+		},
+		{
+			name:    "create schema - already exists",
+			query:   "CREATE SCHEMA test_schema",
+			wantErr: true,
+		},
+	}
+	for _, test := range tests {
 		test := test
 		t.Run(test.name, func(t *testing.T) {
 			if _, err := db.ExecContext(ctx, test.query); err != nil {
-				t.Fatal(err)
+				if !test.wantErr {
+					t.Fatal(err)
+				}
+			} else if test.wantErr {
+				t.Fatal("expected error")
+			}
+			if test.validate != nil {
+				test.validate(t, db)
 			}
 		})
 	}
@@ -148,10 +201,7 @@
 	now := time.Now()
 	ctx := context.Background()
 	ctx = zetasqlite.WithCurrentTime(ctx, now)
-	db, err := sql.Open("zetasqlite", ":memory:")
-	if err != nil {
-		t.Fatal(err)
-	}
+	db := openTestDB(t)
 	defer db.Close()
 	if _, err := db.ExecContext(ctx, `
 CREATE TABLE table (
@@ -213,10 +263,7 @@
 	now := time.Now()
 	ctx := context.Background()
 	ctx = zetasqlite.WithCurrentTime(ctx, now)
-	db, err := sql.Open("zetasqlite", ":memory:")
-	if err != nil {
-		t.Fatal(err)
-	}
+	db := openTestDB(t)
 	defer db.Close()
 	if _, err := db.ExecContext(ctx, "CREATE TEMP TABLE tmp_table (id INT64)"); err != nil {
 		t.Fatal(err)
@@ -234,10 +281,7 @@
 
 func TestWildcardTable(t *testing.T) {
 	ctx := context.Background()
-	db, err := sql.Open("zetasqlite", ":memory:")
-	if err != nil {
-		t.Fatal(err)
-	}
+	db := openTestDB(t)
 	defer db.Close()
 	if _, err := db.ExecContext(
 		ctx,
@@ -341,10 +385,7 @@
 
 func TestTemplatedArgFunc(t *testing.T) {
 	ctx := context.Background()
-	db, err := sql.Open("zetasqlite", ":memory:")
-	if err != nil {
-		t.Fatal(err)
-	}
+	db := openTestDB(t)
 	defer db.Close()
 	t.Run("simple any arguments", func(t *testing.T) {
 		if _, err := db.ExecContext(
@@ -438,10 +479,7 @@
 
 func TestJavaScriptUDF(t *testing.T) {
 	ctx := context.Background()
-	db, err := sql.Open("zetasqlite", ":memory:")
-	if err != nil {
-		t.Fatal(err)
-	}
+	db := openTestDB(t)
 	defer db.Close()
 	t.Run("operation", func(t *testing.T) {
 		if _, err := db.ExecContext(

internal/analyzer.go

@@ -93,6 +93,7 @@ func newAnalyzerOptions() (*zetasql.AnalyzerOptions, error) {
 		ast.CreateFunctionStmt,
 		ast.CreateTableFunctionStmt,
 		ast.CreateViewStmt,
+		ast.CreateSchemaStmt,
 		ast.DropFunctionStmt,
 	})
 	// Enable QUALIFY without WHERE
@@ -263,6 +264,8 @@ func (a *Analyzer) analyzeTemplatedFunctionWithRuntimeArgument(ctx context.Conte
 
 func (a *Analyzer) newStmtAction(ctx context.Context, query string, args []driver.NamedValue, node ast.StatementNode) (StmtAction, error) {
 	switch node.Kind() {
+	case ast.CreateSchemaStmt:
+		return a.newCreateSchemaStmtAction(ctx, query, args, node.(*ast.CreateSchemaStmtNode))
 	case ast.CreateTableStmt:
 		return a.newCreateTableStmtAction(ctx, query, args, node.(*ast.CreateTableStmtNode))
 	case ast.CreateTableAsSelectStmt:
@@ -745,3 +748,16 @@ func getArgsFromParams(values []driver.NamedValue, params []*ast.ParameterNode)
 	}
 	return args, nil
 }
+
+func (a *Analyzer) newCreateSchemaStmtAction(_ context.Context, query string, _ []driver.NamedValue, node *ast.CreateSchemaStmtNode) (*CreateSchemaStmtAction, error) {
+	schemaName := node.NamePath()[0]
+	if _, exists := a.catalog.schemaMap[schemaName]; exists && node.CreateMode() != ast.CreateOrReplaceMode {
+		return nil, fmt.Errorf("schema %s already exists", schemaName)
+	}
+	spec := NewSchemaSpec(schemaName)
+	return &CreateSchemaStmtAction{
+		query:   query,
+		spec:    spec,
+		catalog: a.catalog,
+	}, nil
+}

internal/catalog.go

@@ -51,6 +51,7 @@ const (
 	TableSpecKind    CatalogSpecKind = "table"
 	ViewSpecKind     CatalogSpecKind = "view"
 	FunctionSpecKind CatalogSpecKind = "function"
+	SchemaSpecKind   CatalogSpecKind = "schema"
 	catalogName                      = "zetasqlite"
 )
 
@@ -61,8 +62,10 @@ type Catalog struct {
 	tables       []*TableSpec
 	functions    []*FunctionSpec
 	catalog      *types.SimpleCatalog
+	schemas      []*SchemaSpec
 	tableMap     map[string]*TableSpec
 	funcMap      map[string]*FunctionSpec
+	schemaMap    map[string]*SchemaSpec
 }
 
 func newSimpleCatalog(name string) *types.SimpleCatalog {
@@ -73,10 +76,11 @@ func newSimpleCatalog(name string) *types.SimpleCatalog {
 
 func NewCatalog(db *sql.DB) *Catalog {
 	return &Catalog{
-		db:       db,
-		catalog:  newSimpleCatalog(catalogName),
-		tableMap: map[string]*TableSpec{},
-		funcMap:  map[string]*FunctionSpec{},
+		db:        db,
+		catalog:   newSimpleCatalog(catalogName),
+		tableMap:  map[string]*TableSpec{},
+		funcMap:   map[string]*FunctionSpec{},
+		schemaMap: map[string]*SchemaSpec{},
 	}
 }
 
@@ -209,6 +213,10 @@ func (c *Catalog) Sync(ctx context.Context, conn *Conn) error {
 			if err := c.loadFunctionSpec(spec); err != nil {
 				return fmt.Errorf("failed to load function spec: %w", err)
 			}
+		case SchemaSpecKind:
+			if err := c.loadSchemaSpec(spec); err != nil {
+				return fmt.Errorf("failed to load schema spec: %w", err)
+			}
 		default:
 			return fmt.Errorf("unknown catalog spec kind %s", kind)
 		}
@@ -247,6 +255,18 @@ func (c *Catalog) AddNewFunctionSpec(ctx context.Context, conn *Conn, spec *Func
 	return nil
 }
 
+func (c *Catalog) AddNewSchemaSpec(ctx context.Context, conn *Conn, spec *SchemaSpec) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if err := c.saveSchemaSpec(ctx, conn, spec); err != nil {
+		return fmt.Errorf("failed to save schema spec: %w", err)
+	}
+	c.schemas = append(c.schemas, spec)
+	c.schemaMap[spec.Name] = spec
+	return nil
+}
+
 func (c *Catalog) DeleteTableSpec(ctx context.Context, conn *Conn, name string) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -374,6 +394,26 @@ func (c *Catalog) saveFunctionSpec(ctx context.Context, conn *Conn, spec *Functi
 	return nil
 }
 
+func (c *Catalog) saveSchemaSpec(ctx context.Context, conn *Conn, spec *SchemaSpec) error {
+	encoded, err := json.Marshal(spec)
+	if err != nil {
+		return fmt.Errorf("failed to encode schema spec: %w", err)
+	}
+	now := time.Now()
+	if _, err := conn.ExecContext(
+		ctx,
+		upsertCatalogQuery,
+		sql.Named("name", spec.Name),
+		sql.Named("kind", string(SchemaSpecKind)),
+		sql.Named("spec", string(encoded)),
+		sql.Named("updatedAt", now),
+		sql.Named("createdAt", now),
+	); err != nil {
+		return fmt.Errorf("failed to save schema spec: %w", err)
+	}
+	return nil
+}
+
 func (c *Catalog) createCatalogTablesIfNotExists(ctx context.Context, conn *Conn) error {
 	if _, err := conn.ExecContext(ctx, createCatalogTableQuery); err != nil {
 		return fmt.Errorf("failed to create catalog table: %w", err)
@@ -403,6 +443,16 @@ func (c *Catalog) loadFunctionSpec(spec string) error {
 	return nil
 }
 
+func (c *Catalog) loadSchemaSpec(spec string) error {
+	var schemaSpec SchemaSpec
+	if err := json.Unmarshal([]byte(spec), &schemaSpec); err != nil {
+		return fmt.Errorf("failed to decode schema spec: %w", err)
+	}
+	c.schemas = append(c.schemas, &schemaSpec)
+	c.schemaMap[schemaSpec.Name] = &schemaSpec
+	return nil
+}
+
 func (c *Catalog) trimmedLastPath(path []string) []string {
 	if len(path) == 0 {
 		return path

internal/schema_spec.go

@@ -0,0 +1,22 @@
+package internal
+
+import (
+	"time"
+)
+
+// SchemaSpec represents a schema in the database
+type SchemaSpec struct {
+	Name      string    `json:"name"`
+	UpdatedAt time.Time `json:"updatedAt"`
+	CreatedAt time.Time `json:"createdAt"`
+}
+
+// NewSchemaSpec creates a new schema specification
+func NewSchemaSpec(name string) *SchemaSpec {
+	now := time.Now()
+	return &SchemaSpec{
+		Name:      name,
+		UpdatedAt: now,
+		CreatedAt: now,
+	}
+}

internal/spec.go

@@ -124,7 +124,18 @@ func (s *TableSpec) Column(name string) *ColumnSpec {
 }
 
 func (s *TableSpec) TableName() string {
-	return formatPath(s.NamePath)
+	if len(s.NamePath) > 1 {
+		// First element is schema, rest is table name
+		return fmt.Sprintf("%s.%s", s.NamePath[0], strings.Join(s.NamePath[1:], "_"))
+	}
+	return strings.Join(s.NamePath, "_")
+}
+
+func (s *TableSpec) GetSchema() string {
+	if len(s.NamePath) > 1 {
+		return s.NamePath[0]
+	}
+	return ""
 }
 
 func (s *TableSpec) SQLiteSchema() string {