feat: persist table storage selection

Author: wei

backend/engine.go

@@ -22,6 +22,7 @@ import (
 	sqle "github.com/dolthub/go-mysql-server"
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/analyzer"
+	"github.com/dolthub/go-mysql-server/sql/plan"
 	"github.com/dolthub/go-mysql-server/sql/types"
 	"github.com/dolthub/vitess/go/vt/sqlparser"
 )
@@ -33,13 +34,37 @@ func NewEngine(provider *catalog.DatabaseProvider) (*sqle.Engine, *DuckBuilder)
 	parser := &mysqlParser{Parser: sql.NewMysqlParser()}
 	overrides := sql.EngineOverrides{
 		Builder: sql.BuilderOverrides{Parser: parser},
+		Hooks: sql.ExecutionHooks{
+			CreateTable: sql.CreateTable{
+				PreSQLExecution: prepareMySQLCreateTableStorage,
+			},
+		},
 	}
 	engine := sqle.New(analyzer.NewBuilder(provider).AddOverrides(overrides).Build(), nil)
 	builder := NewDuckBuilder(engine.Analyzer.ExecBuilder, provider)
 	engine.Analyzer.ExecBuilder.PriorityBuilder = builder
 	return engine, builder
 }
 
+// prepareMySQLCreateTableStorage bridges the planner's table-option map to
+// the catalog's request-scoped storage selector. The catalog consumes the
+// selector while creating the table and persists it in the managed comment;
+// no credentials, endpoints, or object paths are accepted from SQL.
+func prepareMySQLCreateTableStorage(ctx *sql.Context, _ sql.StatementRunner, node sql.Node) (sql.Node, error) {
+	create, ok := node.(*plan.CreateTable)
+	if !ok {
+		return node, nil
+	}
+	selection, err := catalog.ResolveMySQLTableStorage(create.TableOpts)
+	if err != nil {
+		return nil, err
+	}
+	if err := catalog.SetTableStorageSelection(ctx, selection); err != nil {
+		return nil, err
+	}
+	return create, nil
+}
+
 // registerMySQLCompatibilitySystemVariables keeps MyDuck's advertised SQL
 // compatibility level stable across GMS upgrades. Clients such as MySQL Shell
 // branch on @@version and otherwise probe newer variables MyDuck does not
@@ -67,13 +92,21 @@ type mysqlParser struct {
 func (p *mysqlParser) ParseSimple(query string) (sqlparser.Statement, error) {
 	compat := rewriteMySQLCompatibility(query)
 	stmt, err := p.Parser.ParseSimple(compat.query)
-	return normalizeMySQLStatement(stmt, compat.replacements), err
+	stmt = normalizeMySQLStatement(stmt, compat.replacements)
+	if err == nil {
+		err = validateMySQLTableStorageStatement(stmt)
+	}
+	return stmt, err
 }
 
 func (p *mysqlParser) Parse(ctx *sql.Context, query string, multi bool) (sqlparser.Statement, string, string, error) {
 	compat := rewriteMySQLCompatibility(query)
 	stmt, parsed, remainder, err := p.Parser.Parse(ctx, compat.query, multi)
-	return normalizeMySQLStatement(stmt, compat.replacements), compat.restoreParsedQuery(parsed), remainder, err
+	stmt = normalizeMySQLStatement(stmt, compat.replacements)
+	if err == nil {
+		err = validateMySQLTableStorageStatement(stmt)
+	}
+	return stmt, compat.restoreParsedQuery(parsed), remainder, err
 }
 
 func (p *mysqlParser) ParseWithOptions(
@@ -85,7 +118,11 @@ func (p *mysqlParser) ParseWithOptions(
 ) (sqlparser.Statement, string, string, error) {
 	compat := rewriteMySQLCompatibility(query)
 	stmt, parsed, remainder, err := p.Parser.ParseWithOptions(ctx, compat.query, delimiter, multi, options)
-	return normalizeMySQLStatement(stmt, compat.replacements), compat.restoreParsedQuery(parsed), remainder, err
+	stmt = normalizeMySQLStatement(stmt, compat.replacements)
+	if err == nil {
+		err = validateMySQLTableStorageStatement(stmt)
+	}
+	return stmt, compat.restoreParsedQuery(parsed), remainder, err
 }
 
 func (p *mysqlParser) ParseOneWithOptions(
@@ -95,7 +132,33 @@ func (p *mysqlParser) ParseOneWithOptions(
 ) (sqlparser.Statement, int, error) {
 	compat := rewriteMySQLCompatibility(query)
 	stmt, index, err := p.Parser.ParseOneWithOptions(ctx, compat.query, options)
-	return normalizeMySQLStatement(stmt, compat.replacements), compat.originalOffset(index), err
+	stmt = normalizeMySQLStatement(stmt, compat.replacements)
+	if err == nil {
+		err = validateMySQLTableStorageStatement(stmt)
+	}
+	return stmt, compat.originalOffset(index), err
+}
+
+// validateMySQLTableStorageStatement runs before the planner turns table
+// options into a map. That preserves duplicate ENGINE/myduck_storage
+// declarations, which would otherwise be silently overwritten by the map.
+func validateMySQLTableStorageStatement(stmt sqlparser.Statement) error {
+	ddl, ok := stmt.(*sqlparser.DDL)
+	if !ok || ddl.TableSpec == nil || len(ddl.TableSpec.TableOpts) == 0 {
+		return nil
+	}
+	options := make([]catalog.TableStorageOption, 0, len(ddl.TableSpec.TableOpts))
+	for _, option := range ddl.TableSpec.TableOpts {
+		if option == nil {
+			continue
+		}
+		options = append(options, catalog.TableStorageOption{
+			Name:  option.Name,
+			Value: option.Value,
+		})
+	}
+	_, err := catalog.NormalizeTableStorageOptions(options)
+	return err
 }
 
 type mysqlOptionReplacement struct {

backend/engine_test.go

@@ -18,6 +18,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/apecloud/myduckserver/catalog"
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/vitess/go/vt/sqlparser"
 	"github.com/stretchr/testify/require"
@@ -154,3 +155,36 @@ func TestMySQLParserRestoresDatabaseFiltersInMultiQuery(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, strings.Index(query, " SELECT 1"), index)
 }
+
+func TestMySQLParserTableStorageOptions(t *testing.T) {
+	parser := &mysqlParser{Parser: sql.NewMysqlParser()}
+	for _, query := range []string{
+		"CREATE TABLE object_table (id INT) ENGINE=DUCKLAKE",
+		"CREATE TABLE local_table (id INT) ENGINE=InnoDB",
+	} {
+		stmt, _, _, err := parser.ParseWithOptions(context.Background(), query, ';', false, sqlparser.ParserOptions{})
+		require.NoError(t, err, query)
+		ddl, ok := stmt.(*sqlparser.DDL)
+		require.True(t, ok, query)
+		require.NotNil(t, ddl.TableSpec, query)
+		require.NotEmpty(t, ddl.TableSpec.TableOpts, query)
+		for _, option := range ddl.TableSpec.TableOpts {
+			t.Logf("%s => name=%q value=%q", query, option.Name, option.Value)
+		}
+	}
+}
+
+func TestMySQLParserRejectsConflictingTableStorageOptions(t *testing.T) {
+	parser := &mysqlParser{Parser: sql.NewMysqlParser()}
+	for _, test := range []struct {
+		query string
+		want  error
+	}{
+		{query: "CREATE TABLE duplicate_engine (id INT) ENGINE=DUCKLAKE ENGINE=DUCKLAKE", want: catalog.ErrTableStorageDuplicate},
+		{query: "CREATE TABLE conflicting_engine (id INT) ENGINE=DUCKLAKE ENGINE=LOCAL", want: catalog.ErrTableStorageConflict},
+	} {
+		_, _, _, err := parser.ParseWithOptions(context.Background(), test.query, ';', false, sqlparser.ParserOptions{})
+		require.Error(t, err, test.query)
+		require.ErrorIs(t, err, test.want, test.query)
+	}
+}

catalog/database.go

@@ -146,7 +146,16 @@ func (d *Database) Name() string {
 	return d.name
 }
 
-func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, temporary bool) error {
+func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, storage TableStorageSelection, temporary bool) error {
+	if err := storage.Validate(); err != nil {
+		return err
+	}
+	if temporary && storage.IsObjectStorage() {
+		return fmt.Errorf("%w: temporary tables cannot use object storage", ErrInvalidTableStorage)
+	}
+	if storage.Kind == "" {
+		storage = DefaultTableStorageSelection()
+	}
 	var columns []string
 	var columnCommentSQLs []string
 	var fullTableName string
@@ -269,7 +278,13 @@ func (d *Database) createAllTable(ctx *sql.Context, name string, schema sql.Prim
 	b.WriteString(")")
 
 	// Add comment to the table
-	info := ExtraTableInfo{schema.PkOrdinals, withoutIndex, fullSequenceName, nil}
+	info := ExtraTableInfo{
+		PkOrdinals: schema.PkOrdinals,
+		Replicated: withoutIndex,
+		Sequence:   fullSequenceName,
+		Checks:     nil,
+		Storage:    storage.Kind,
+	}
 	b.WriteString(fmt.Sprintf(
 		"; COMMENT ON TABLE %s IS '%s'",
 		fullTableName,
@@ -322,14 +337,88 @@ func isIndexCreationDisabled(ctx *sql.Context) bool {
 func (d *Database) CreateTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string) error {
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	return d.createAllTable(ctx, name, schema, collation, comment, false)
+	storage := DefaultTableStorageSelection()
+	if selected, ok := TableStorageSelectionFromContext(ctx); ok {
+		storage = selected
+	}
+	return d.createAllTable(ctx, name, schema, collation, comment, storage, false)
+}
+
+// CreateTableWithStorage is the explicit catalog boundary for protocol
+// adapters that already normalized a table selector. It is intentionally
+// limited to selection propagation and metadata; object-table physical
+// routing is owned by the follow-up storage implementation.
+func (d *Database) CreateTableWithStorage(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID, comment string, storage TableStorageSelection) error {
+	if err := storage.Validate(); err != nil {
+		return err
+	}
+	if err := SetTableStorageSelection(ctx, storage); err != nil {
+		return err
+	}
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	return d.createAllTable(ctx, name, schema, collation, comment, storage, false)
+}
+
+// RecordTableStorageSelection updates the managed table metadata for a table
+// created by a protocol path that bypasses sql.TableCreator (currently the
+// PostgreSQL handler). It preserves the user-visible table comment and makes
+// the selection available after a fresh catalog reload.
+func (d *Database) RecordTableStorageSelection(ctx *sql.Context, name string, storage TableStorageSelection) error {
+	if err := storage.Validate(); err != nil {
+		return err
+	}
+	if d.catalog == "temp" && storage.IsObjectStorage() {
+		return fmt.Errorf("%w: temporary tables cannot use object storage", ErrInvalidTableStorage)
+	}
+
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	rows, err := adapter.QueryCatalog(ctx, `
+		SELECT comment
+		FROM duckdb_tables()
+		WHERE database_name = ? AND schema_name = ? AND table_name = ?
+	`, d.catalog, d.name, name)
+	if err != nil {
+		return ErrDuckDB.New(err)
+	}
+	defer rows.Close()
+	if !rows.Next() {
+		if err := rows.Err(); err != nil {
+			return ErrDuckDB.New(err)
+		}
+		return sql.ErrTableNotFound.New(name)
+	}
+
+	var rawComment stdsql.NullString
+	if err := rows.Scan(&rawComment); err != nil {
+		_ = rows.Close()
+		return ErrDuckDB.New(err)
+	}
+	if err := rows.Close(); err != nil {
+		return ErrDuckDB.New(err)
+	}
+	comment := DecodeComment[ExtraTableInfo](rawComment.String)
+	info := comment.Meta
+	info.Storage = storage.Kind
+	encoded := NewCommentWithMeta(comment.Text, info).Encode()
+	_, err = adapter.Exec(ctx, fmt.Sprintf(`COMMENT ON TABLE %s IS '%s'`, FullTableName(d.catalog, d.name, name), encoded))
+	if err != nil {
+		return ErrDuckDB.New(err)
+	}
+	return nil
 }
 
 // CreateTemporaryTable implements sql.CreateTemporaryTable.
 func (d *Database) CreateTemporaryTable(ctx *sql.Context, name string, schema sql.PrimaryKeySchema, collation sql.CollationID) error {
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	return d.createAllTable(ctx, name, schema, collation, "", true)
+	storage := DefaultTableStorageSelection()
+	if selected, ok := TableStorageSelectionFromContext(ctx); ok {
+		storage = selected
+	}
+	return d.createAllTable(ctx, name, schema, collation, "", storage, true)
 }
 
 // DropTable implements sql.TableDropper.

catalog/table.go

@@ -32,6 +32,26 @@ type ExtraTableInfo struct {
 	Replicated bool
 	Sequence   string
 	Checks     []sql.CheckDefinition
+	// Storage records the table's durable storage class. An empty value is
+	// treated as local for metadata written before table-level storage
+	// selection existed; newly-created tables always write the explicit value.
+	Storage TableStorageKind `json:"storage,omitempty"`
+}
+
+// StorageKind returns the effective storage class for this metadata. Missing
+// storage metadata is the backwards-compatible local-table behavior.
+func (info ExtraTableInfo) StorageKind() TableStorageKind {
+	if info.Storage == TableStorageObject {
+		return TableStorageObject
+	}
+	return TableStorageLocal
+}
+
+func (info *ExtraTableInfo) normalizeStorage() {
+	if info == nil {
+		return
+	}
+	info.Storage = info.StorageKind()
 }
 
 type ColumnInfo struct {
@@ -75,6 +95,10 @@ func NewTable(db *Database, name string, hasPrimaryKey bool) *Table {
 }
 
 func (t *Table) withComment(comment *Comment[ExtraTableInfo]) *Table {
+	if comment == nil {
+		comment = NewComment[ExtraTableInfo]("")
+	}
+	comment.Meta.normalizeStorage()
 	t.comment = comment
 	return t
 }
@@ -100,7 +124,12 @@ func (t *Table) withSchema(ctx *sql.Context) error {
 }
 
 func (t *Table) ExtraTableInfo() ExtraTableInfo {
-	return t.comment.Meta
+	if t.comment == nil {
+		return ExtraTableInfo{Storage: TableStorageLocal}
+	}
+	info := t.comment.Meta
+	info.normalizeStorage()
+	return info
 }
 
 func (t *Table) HasPrimaryKey() bool {