diff --git a/runtime/ai/developer_agent.go b/runtime/ai/developer_agent.go index cdd35105efb..f5f3c96c4cf 100644 --- a/runtime/ai/developer_agent.go +++ b/runtime/ai/developer_agent.go @@ -68,11 +68,11 @@ func (t *DeveloperAgent) Handler(ctx context.Context, args *DeveloperAgentArgs) return nil, err } if args.CurrentFilePath != "" { - _, err := s.CallTool(ctx, RoleAssistant, ReadFileName, nil, &ReadFileArgs{ + _, _ = s.CallTool(ctx, RoleAssistant, ReadFileName, nil, &ReadFileArgs{ Path: args.CurrentFilePath, }) - if err != nil { - return nil, err + if ctx.Err() != nil { // Ignore tool error since the file may not exist + return nil, ctx.Err() } } diff --git a/runtime/ai/instructions/data/development.md b/runtime/ai/instructions/data/development.md index ac28a0297cc..cfe856a8a74 100644 --- a/runtime/ai/instructions/data/development.md +++ b/runtime/ai/instructions/data/development.md @@ -244,7 +244,7 @@ The following tools are typically available for project development: - `project_status` for checking resource names and their current status (idle, running, error) - `query_sql` for running SQL against a connector; use `SELECT` statements with `LIMIT` clauses and low timeouts, and be mindful of performance or making too many queries - `query_metrics_view` for querying a metrics view; useful for answering data questions and validating dashboard behavior -- `list_tables` and `get_table` for accessing the information schema of a database connector +- `list_tables` and `show_table` for accessing the information schema of a database connector - `list_buckets` and `list_bucket_files` for exploring files in object stores like S3 or GCS; to preview file contents, load one file into a table using a model and query it with `query_sql` {% if .external %} diff --git a/runtime/ai/show_table.go b/runtime/ai/show_table.go index 76b2b12f9cd..d0654880142 100644 --- a/runtime/ai/show_table.go +++ b/runtime/ai/show_table.go @@ -19,9 +19,9 @@ var _ Tool[*ShowTableArgs, *ShowTableResult] = (*ShowTable)(nil) type ShowTableArgs struct { Connector string `json:"connector,omitempty" jsonschema:"Optional OLAP connector name. Defaults to the instance's default OLAP connector."` - Table string `json:"table" jsonschema:"The name of the table to describe."` - Database string `json:"database,omitempty" jsonschema:"Optional database name for connectors that support multiple databases."` - DatabaseSchema string `json:"database_schema,omitempty" jsonschema:"Optional database schema name."` + Table string `json:"table" jsonschema:"Name of the table to describe. Must be a simple table name; database/schema names should be provided using the separate fields."` + Database string `json:"database,omitempty" jsonschema:"Database that contains the table (defaults to the connector's default database if applicable)."` + DatabaseSchema string `json:"database_schema,omitempty" jsonschema:"Database schema that contains the table (defaults to the connector's default schema if applicable)."` } type ShowTableResult struct { @@ -29,18 +29,19 @@ type ShowTableResult struct { IsView bool `json:"is_view"` Columns []ColumnInfo `json:"columns"` PhysicalSizeBytes int64 `json:"physical_size_bytes,omitempty" jsonschema:"The physical size of the table in bytes. If 0 or omitted, size information is not available."` + DDL string `json:"ddl,omitempty" jsonschema:"The SQL DDL statement (CREATE TABLE/VIEW) for this table, if available."` } type ColumnInfo struct { - Name string `json:"name"` - Type string `json:"type"` + Name string `json:"name" jsonschema:"The name of the column."` + Type string `json:"type" jsonschema:"The data type of the column. This is a generic type code and does not exactly match the underlying SQL type."` } func (t *ShowTable) Spec() *mcp.Tool { return &mcp.Tool{ Name: ShowTableName, Title: "Show Table", - Description: "Show schema and column information for a table in an OLAP connector.", + Description: "Show schema and column information for a table in an OLAP connector. Note: Table, schema and database names passed to this tool are case sensitive; if you get an error and you're working with a database that folds unquoted identifiers (e.g Snowflake folds to uppercase), you may need to retry with the casing adjusted accordingly.", Meta: map[string]any{ "openai/toolInvocation/invoking": "Getting table schema...", "openai/toolInvocation/invoked": "Got table schema", @@ -75,11 +76,15 @@ func (t *ShowTable) Handler(ctx context.Context, args *ShowTableArgs) (*ShowTabl // Load physical size _ = olap.InformationSchema().LoadPhysicalSize(ctx, []*drivers.OlapTable{table}) + // Load DDL + _ = olap.InformationSchema().LoadDDL(ctx, table) + // Build result result := &ShowTableResult{ Name: table.Name, IsView: table.View, PhysicalSizeBytes: table.PhysicalSizeBytes, + DDL: table.DDL, Columns: make([]ColumnInfo, 0), } diff --git a/runtime/drivers/athena/olap.go b/runtime/drivers/athena/olap.go index a7e79baf17d..8398516705f 100644 --- a/runtime/drivers/athena/olap.go +++ b/runtime/drivers/athena/olap.go @@ -92,6 +92,11 @@ func (c *Connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (c *Connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + return nil // Not implemented +} + // Lookup implements drivers.OLAPInformationSchema. func (c *Connection) Lookup(ctx context.Context, db, schema, name string) (*drivers.OlapTable, error) { meta, err := c.GetTable(ctx, db, schema, name) diff --git a/runtime/drivers/bigquery/olap.go b/runtime/drivers/bigquery/olap.go index f2773a09cd9..12ec76ed2b6 100644 --- a/runtime/drivers/bigquery/olap.go +++ b/runtime/drivers/bigquery/olap.go @@ -133,6 +133,35 @@ func (c *Connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (c *Connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + client, err := c.getClient(ctx) + if err != nil { + return err + } + + q := fmt.Sprintf("SELECT ddl FROM `%s.%s.INFORMATION_SCHEMA.TABLES` WHERE table_name = @name", table.Database, table.DatabaseSchema) + cq := client.Query(q) + cq.Parameters = []bigquery.QueryParameter{ + {Name: "name", Value: table.Name}, + } + + it, err := cq.Read(ctx) + if err != nil { + return err + } + + var row struct { + DDL string `bigquery:"ddl"` + } + err = it.Next(&row) + if err != nil { + return err + } + table.DDL = row.DDL + return nil +} + // Lookup implements drivers.OLAPInformationSchema. func (c *Connection) Lookup(ctx context.Context, db, schema, name string) (*drivers.OlapTable, error) { client, err := c.getClient(ctx) diff --git a/runtime/drivers/bigquery/olap_test.go b/runtime/drivers/bigquery/olap_test.go index 9849ac6b6e5..243902c0bc2 100644 --- a/runtime/drivers/bigquery/olap_test.go +++ b/runtime/drivers/bigquery/olap_test.go @@ -151,6 +151,18 @@ func TestExec(t *testing.T) { require.NoError(t, err) } +func TestLoadDDL(t *testing.T) { + testmode.Expensive(t) + _, olap := acquireTestBigQuery(t) + + table, err := olap.InformationSchema().Lookup(t.Context(), "rilldata", "integration_test", "all_datatypes") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(t.Context(), table) + require.NoError(t, err) + require.NotEmpty(t, table.DDL) + require.Contains(t, table.DDL, "all_datatypes") +} + func TestScan(t *testing.T) { testmode.Expensive(t) _, olap := acquireTestBigQuery(t) diff --git a/runtime/drivers/clickhouse/information_schema.go b/runtime/drivers/clickhouse/information_schema.go index 1b450c083f8..6d105afbf8a 100644 --- a/runtime/drivers/clickhouse/information_schema.go +++ b/runtime/drivers/clickhouse/information_schema.go @@ -407,6 +407,27 @@ func (c *Connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return err } +func (c *Connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + conn, release, err := c.acquireMetaConn(ctx) + if err != nil { + return err + } + defer func() { _ = release() }() + + schema := table.DatabaseSchema + if schema == "" { + schema = c.config.Database // In Clickhouse, this is actually like a schema + } + + var ddl string + err = conn.QueryRowxContext(ctx, fmt.Sprintf("SHOW CREATE TABLE %s.%s", safeSQLName(schema), safeSQLName(table.Name))).Scan(&ddl) + if err != nil { + return err + } + table.DDL = ddl + return nil +} + func scanTables(rows *sqlx.Rows) ([]*drivers.OlapTable, error) { var res []*drivers.OlapTable diff --git a/runtime/drivers/clickhouse/information_schema_test.go b/runtime/drivers/clickhouse/information_schema_test.go index 5ab62cbb683..e695696d7f8 100644 --- a/runtime/drivers/clickhouse/information_schema_test.go +++ b/runtime/drivers/clickhouse/information_schema_test.go @@ -39,6 +39,7 @@ func TestInformationSchema(t *testing.T) { t.Run("testInformationSchemaGetTable", func(t *testing.T) { testInformationSchemaGetTable(t, infoSchema) }) t.Run("testInformationSchemaListDatabaseSchemasPagination", func(t *testing.T) { testInformationSchemaListDatabaseSchemasPagination(t, infoSchema) }) t.Run("testInformationSchemaListTablesPagination", func(t *testing.T) { testInformationSchemaListTablesPagination(t, infoSchema) }) + t.Run("testLoadDDL", func(t *testing.T) { testLoadDDL(t, conn) }) } func testInformationSchemaAll(t *testing.T, conn drivers.Handle) { @@ -374,3 +375,24 @@ func prepareConn(t *testing.T, conn drivers.Handle) { }) require.NoError(t, err) } + +func testLoadDDL(t *testing.T, conn drivers.Handle) { + olap, _ := conn.AsOLAP("") + ctx := context.Background() + + // Test DDL for a table + table, err := olap.InformationSchema().Lookup(ctx, "", "", "foo") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(ctx, table) + require.NoError(t, err) + require.Contains(t, table.DDL, "CREATE TABLE") + require.Contains(t, table.DDL, "foo") + + // Test DDL for a view + view, err := olap.InformationSchema().Lookup(ctx, "", "", "model") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(ctx, view) + require.NoError(t, err) + require.Contains(t, view.DDL, "CREATE VIEW") + require.Contains(t, view.DDL, "model") +} diff --git a/runtime/drivers/druid/information_schema.go b/runtime/drivers/druid/information_schema.go index ef182eb24b6..f17c836aaa1 100644 --- a/runtime/drivers/druid/information_schema.go +++ b/runtime/drivers/druid/information_schema.go @@ -252,6 +252,10 @@ func (c *connection) Lookup(ctx context.Context, db, schema, name string) (*driv return tables[0], nil } +func (c *connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + return nil // Not implemented +} + func (c *connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.OlapTable) error { q := `SELECT datasource, diff --git a/runtime/drivers/duckdb/information_schema.go b/runtime/drivers/duckdb/information_schema.go index ebd061f4d9b..7d44a8589ee 100644 --- a/runtime/drivers/duckdb/information_schema.go +++ b/runtime/drivers/duckdb/information_schema.go @@ -273,6 +273,21 @@ func (c *connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +func (c *connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + db, release, err := c.acquireDB() + if err != nil { + return err + } + defer func() { _ = release() }() + + ddl, err := db.DDL(ctx, table.Database, table.DatabaseSchema, table.Name) + if err != nil { + return c.checkErr(err) + } + table.DDL = ddl + return nil +} + func scanTables(rows []*rduckdb.Table) ([]*drivers.OlapTable, error) { var res []*drivers.OlapTable diff --git a/runtime/drivers/duckdb/information_schema_test.go b/runtime/drivers/duckdb/information_schema_test.go index 46013522bb8..ff1747758db 100644 --- a/runtime/drivers/duckdb/information_schema_test.go +++ b/runtime/drivers/duckdb/information_schema_test.go @@ -41,6 +41,7 @@ func TestInformationSchema(t *testing.T) { t.Run("testInformationSchemaListTables", func(t *testing.T) { testInformationSchemaListTables(t, infoSchema, database, databaseSchema) }) t.Run("testInformationSchemaGetTable", func(t *testing.T) { testInformationSchemaGetTable(t, infoSchema, database, databaseSchema) }) t.Run("testInformationSchemaListTablesPagination", func(t *testing.T) { testInformationSchemaListTablesPagination(t, infoSchema, database, databaseSchema) }) + t.Run("testLoadDDL", func(t *testing.T) { testLoadDDL(t, olap) }) } func TestInformationSchemaMotherduck(t *testing.T) { @@ -283,3 +284,23 @@ func testInformationSchemaListTablesPagination(t *testing.T, infoSchema drivers. require.Equal(t, 6, len(tables)) require.Empty(t, nextToken) } + +func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { + ctx := context.Background() + + // Test DDL for a materialized table + table, err := olap.InformationSchema().Lookup(ctx, "", "", "bar") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(ctx, table) + require.NoError(t, err) + require.Contains(t, table.DDL, "CREATE TABLE") + require.Contains(t, table.DDL, "bar") + + // Test DDL for a view + view, err := olap.InformationSchema().Lookup(ctx, "", "", "model") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(ctx, view) + require.NoError(t, err) + require.Contains(t, view.DDL, "CREATE VIEW") + require.Contains(t, view.DDL, "model") +} diff --git a/runtime/drivers/mysql/olap.go b/runtime/drivers/mysql/olap.go index d7ca9c31810..e157c652e34 100644 --- a/runtime/drivers/mysql/olap.go +++ b/runtime/drivers/mysql/olap.go @@ -118,6 +118,40 @@ func (c *connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (c *connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + db, err := c.getDB(ctx) + if err != nil { + return err + } + + // SHOW CREATE TABLE works for both tables and views in MySQL. + // For tables it returns columns: [Table, Create Table]. + // For views it returns columns: [View, Create View, character_set_client, collation_connection]. + // We extract the DDL by column name to avoid depending on column order or count. + rows, err := db.QueryxContext(ctx, fmt.Sprintf("SHOW CREATE TABLE %s", drivers.DialectMySQL.EscapeTable(table.Database, table.DatabaseSchema, table.Name))) + if err != nil { + return err + } + defer rows.Close() + + if rows.Next() { + res := make(map[string]any) + if err := rows.MapScan(res); err != nil { + return err + } + for _, key := range []string{"Create Table", "Create View"} { + if v, ok := res[key]; ok && v != nil { + if b, ok := v.([]byte); ok { + table.DDL = string(b) + } + break + } + } + } + return rows.Err() +} + // Lookup implements drivers.OLAPInformationSchema. func (c *connection) Lookup(ctx context.Context, db, schema, name string) (*drivers.OlapTable, error) { meta, err := c.GetTable(ctx, db, schema, name) diff --git a/runtime/drivers/mysql/olap_test.go b/runtime/drivers/mysql/olap_test.go index 282fd012795..5f2f4190931 100644 --- a/runtime/drivers/mysql/olap_test.go +++ b/runtime/drivers/mysql/olap_test.go @@ -33,7 +33,9 @@ func TestOLAP(t *testing.T) { t.Run("Test Scan Full Table", func(t *testing.T) { testFullTableScan(t, olap) }) - + t.Run("Test LoadDDL", func(t *testing.T) { + testLoadDDL(t, olap) + }) } func testMapScan(t *testing.T, olap drivers.OLAPStore) { @@ -486,6 +488,33 @@ func testFullTableScan(t *testing.T, olap drivers.OLAPStore) { require.Equal(t, count, 3) } +func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { + // Test DDL for a table + table, err := olap.InformationSchema().Lookup(t.Context(), "", "", "all_datatypes") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(t.Context(), table) + require.NoError(t, err) + require.Contains(t, table.DDL, "CREATE TABLE") + require.Contains(t, table.DDL, "all_datatypes") + + // Create a view and test DDL for it + err = olap.Exec(t.Context(), &drivers.Statement{Query: "CREATE OR REPLACE VIEW test_ddl_view AS SELECT int_col, varchar_col FROM all_datatypes"}) + require.NoError(t, err) + t.Cleanup(func() { + _ = olap.Exec(t.Context(), &drivers.Statement{Query: "DROP VIEW IF EXISTS test_ddl_view"}) + }) + + view, err := olap.InformationSchema().Lookup(t.Context(), "", "", "test_ddl_view") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(t.Context(), view) + require.NoError(t, err) + // MySQL's DDL output for views is wack, so splitting the pieces out like this. Not worth fixing as it does contain the essential information. + require.Contains(t, strings.ToLower(view.DDL), "create") + require.Contains(t, strings.ToLower(view.DDL), "view") + require.Contains(t, strings.ToLower(view.DDL), "test_ddl_view") + require.Contains(t, strings.ToLower(view.DDL), "as select") +} + func acquireTestMySQL(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "mysql") conn, err := drivers.Open("mysql", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) diff --git a/runtime/drivers/olap.go b/runtime/drivers/olap.go index c1177bd54c6..4490036436e 100644 --- a/runtime/drivers/olap.go +++ b/runtime/drivers/olap.go @@ -176,6 +176,9 @@ type OLAPInformationSchema interface { // LoadPhysicalSize populates the PhysicalSizeBytes field of table metadata. // It should be called after All or Lookup and not on manually created tables. LoadPhysicalSize(ctx context.Context, tables []*OlapTable) error + // LoadDDL populates the DDL field of a single table's metadata. + // Drivers that don't support DDL retrieval should return nil (leaving DDL empty). + LoadDDL(ctx context.Context, table *OlapTable) error } // OlapTable represents a table in an information schema. @@ -190,6 +193,7 @@ type OlapTable struct { Schema *runtimev1.StructType UnsupportedCols map[string]string PhysicalSizeBytes int64 + DDL string } // Dialect enumerates OLAP query languages. diff --git a/runtime/drivers/pinot/information_schema.go b/runtime/drivers/pinot/information_schema.go index d2c73e72bf8..cc4ee26ee7d 100644 --- a/runtime/drivers/pinot/information_schema.go +++ b/runtime/drivers/pinot/information_schema.go @@ -359,6 +359,11 @@ func (c *connection) Lookup(ctx context.Context, db, schema, name string) (*driv return table, nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (c *connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + return nil // Not implemented +} + // LoadPhysicalSize populates the PhysicalSizeBytes field of the tables. // This was not tested when implemented so should be tested when pinot becomes a fairly used connector. func (c *connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.OlapTable) error { diff --git a/runtime/drivers/postgres/olap.go b/runtime/drivers/postgres/olap.go index 2728912dca9..6d595bbf10a 100644 --- a/runtime/drivers/postgres/olap.go +++ b/runtime/drivers/postgres/olap.go @@ -98,6 +98,62 @@ func (c *connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +// Note: table.Database is not used; in Postgres, the database is determined by the connection. +func (c *connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + db, err := c.getDB(ctx) + if err != nil { + return err + } + + schema := table.DatabaseSchema + if schema == "" { + if err := db.QueryRowContext(ctx, "SELECT current_schema()").Scan(&schema); err != nil { + return err + } + } + + if table.View { + // For views: use pg_get_viewdef + var ddl string + q := ` + SELECT 'CREATE VIEW ' || quote_ident(n.nspname) || '.' || quote_ident(c.relname) || ' AS ' || pg_get_viewdef(c.oid, true) + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('v', 'm') + ` + err = db.QueryRowContext(ctx, q, schema, table.Name).Scan(&ddl) + if err != nil { + return err + } + table.DDL = ddl + return nil + } + + // Postgres does not have a built-in way to get the DDL for a table, so we reconstruct a basic CREATE TABLE statement from the available metadata (won't include indexes, constraints, etc.). + q := ` + SELECT + 'CREATE TABLE ' || quote_ident(n.nspname) || '.' || quote_ident(c.relname) || ' (' || + string_agg( + quote_ident(a.attname) || ' ' || format_type(a.atttypid, a.atttypmod) || + CASE WHEN a.attnotnull THEN ' NOT NULL' ELSE '' END, + ', ' ORDER BY a.attnum + ) || ')' + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + JOIN pg_attribute a ON a.attrelid = c.oid + WHERE n.nspname = $1 AND c.relname = $2 AND a.attnum > 0 AND NOT a.attisdropped + GROUP BY n.nspname, c.relname + ` + var ddl string + err = db.QueryRowContext(ctx, q, schema, table.Name).Scan(&ddl) + if err != nil { + return err + } + table.DDL = ddl + return nil +} + // Lookup implements drivers.OLAPInformationSchema. func (c *connection) Lookup(ctx context.Context, db, schema, name string) (*drivers.OlapTable, error) { meta, err := c.GetTable(ctx, db, schema, name) diff --git a/runtime/drivers/postgres/olap_test.go b/runtime/drivers/postgres/olap_test.go index 43382814ed0..a1cd1c9d8a5 100644 --- a/runtime/drivers/postgres/olap_test.go +++ b/runtime/drivers/postgres/olap_test.go @@ -64,6 +64,10 @@ func TestPgxOLAP(t *testing.T) { t.Run("test exec", func(t *testing.T) { testExec(t, olap) }) + + t.Run("test LoadDDL", func(t *testing.T) { + testLoadDDL(t, olap) + }) } func testOLAP(t *testing.T, olap drivers.OLAPStore) { @@ -394,6 +398,31 @@ func testExec(t *testing.T, olap drivers.OLAPStore) { require.Equal(t, "test", res["name"]) } +func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { + // Test DDL for a table + table, err := olap.InformationSchema().Lookup(t.Context(), "", "", "all_datatypes") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(t.Context(), table) + require.NoError(t, err) + require.Contains(t, table.DDL, "CREATE TABLE") + require.Contains(t, table.DDL, "all_datatypes") + + // Create a view and test DDL for it + tableName := fmt.Sprintf("test_ddl_view_%d", time.Now().UnixNano()) + err = olap.Exec(t.Context(), &drivers.Statement{Query: fmt.Sprintf("CREATE VIEW %s AS SELECT id, name FROM all_datatypes", tableName)}) + require.NoError(t, err) + t.Cleanup(func() { + _ = olap.Exec(t.Context(), &drivers.Statement{Query: fmt.Sprintf("DROP VIEW IF EXISTS %s", tableName)}) + }) + + view, err := olap.InformationSchema().Lookup(t.Context(), "", "", tableName) + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(t.Context(), view) + require.NoError(t, err) + require.Contains(t, view.DDL, "CREATE VIEW") + require.Contains(t, view.DDL, tableName) +} + func acquireTestPostgres(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "postgres") conn, err := drivers.Open("postgres", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) diff --git a/runtime/drivers/redshift/olap.go b/runtime/drivers/redshift/olap.go index f9f4d33b7d7..380bfd907d6 100644 --- a/runtime/drivers/redshift/olap.go +++ b/runtime/drivers/redshift/olap.go @@ -114,6 +114,11 @@ func (c *Connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (c *Connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + return nil // Not implemented +} + // Lookup implements drivers.OLAPInformationSchema. func (c *Connection) Lookup(ctx context.Context, db, schema, name string) (*drivers.OlapTable, error) { meta, err := c.GetTable(ctx, db, schema, name) diff --git a/runtime/drivers/snowflake/olap.go b/runtime/drivers/snowflake/olap.go index 1192034f3c9..082943d25ca 100644 --- a/runtime/drivers/snowflake/olap.go +++ b/runtime/drivers/snowflake/olap.go @@ -3,6 +3,7 @@ package snowflake import ( "context" "fmt" + "strings" "github.com/jmoiron/sqlx" runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1" @@ -98,6 +99,31 @@ func (c *connection) LoadPhysicalSize(ctx context.Context, tables []*drivers.Ola return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (c *connection) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + db, err := c.getDB(ctx) + if err != nil { + return err + } + + // HACK: Since All and Lookup don't always return the correct casing, we uppercase the table name here as that's usually necessary in Snowflake. + // This is a workaround until we return correct casing from All and Lookup. + fqn := drivers.DialectSnowflake.EscapeTable(strings.ToUpper(table.Database), strings.ToUpper(table.DatabaseSchema), strings.ToUpper(table.Name)) + + objectType := "TABLE" + if table.View { + objectType = "VIEW" + } + + var ddl string + err = db.QueryRowContext(ctx, fmt.Sprintf("SELECT GET_DDL('%s', ?)", objectType), fqn).Scan(&ddl) + if err != nil { + return err + } + table.DDL = ddl + return nil +} + // Lookup implements drivers.OLAPInformationSchema. func (c *connection) Lookup(ctx context.Context, db, schema, name string) (*drivers.OlapTable, error) { meta, err := c.GetTable(ctx, db, schema, name) diff --git a/runtime/drivers/snowflake/olap_test.go b/runtime/drivers/snowflake/olap_test.go index 1324f7e9af0..3092a7ed5b9 100644 --- a/runtime/drivers/snowflake/olap_test.go +++ b/runtime/drivers/snowflake/olap_test.go @@ -2,6 +2,7 @@ package snowflake_test import ( "encoding/json" + "strings" "testing" "time" @@ -153,6 +154,18 @@ func TestComplexTypes(t *testing.T) { require.NoError(t, rows.Err()) } +func TestLoadDDL(t *testing.T) { + testmode.Expensive(t) + _, olap := acquireTestSnowflake(t) + + table, err := olap.InformationSchema().Lookup(t.Context(), "", "", "all_datatypes") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(t.Context(), table) + require.NoError(t, err) + require.Contains(t, table.DDL, "CREATE") + require.Contains(t, strings.ToUpper(table.DDL), "ALL_DATATYPES") +} + func TestDryRun(t *testing.T) { testmode.Expensive(t) diff --git a/runtime/drivers/snowflake/snowflake.go b/runtime/drivers/snowflake/snowflake.go index d7b0f37df7a..57666acedb6 100644 --- a/runtime/drivers/snowflake/snowflake.go +++ b/runtime/drivers/snowflake/snowflake.go @@ -8,6 +8,7 @@ import ( "encoding/pem" "errors" "fmt" + "io" "strings" "time" @@ -24,6 +25,9 @@ import ( func init() { drivers.Register("snowflake", driver{}) drivers.RegisterAsConnector("snowflake", driver{}) + + // Naughty Snowflake does logging inside the library using a global. + gosnowflake.GetLogger().SetOutput(io.Discard) } var spec = drivers.Spec{ diff --git a/runtime/drivers/starrocks/information_schema.go b/runtime/drivers/starrocks/information_schema.go index 3db4bb24be5..993a0b6aff8 100644 --- a/runtime/drivers/starrocks/information_schema.go +++ b/runtime/drivers/starrocks/information_schema.go @@ -200,6 +200,29 @@ func (i *informationSchema) LoadPhysicalSize(ctx context.Context, tables []*driv return nil } +// LoadDDL implements drivers.OLAPInformationSchema. +func (i *informationSchema) LoadDDL(ctx context.Context, table *drivers.OlapTable) error { + db := i.c.db + + catalog := table.Database + if catalog == "" { + catalog = i.c.configProp.Catalog + } + schema := table.DatabaseSchema + if schema == "" { + schema = i.c.configProp.Database + } + + q := fmt.Sprintf("SHOW CREATE TABLE %s.%s.%s", safeSQLName(catalog), safeSQLName(schema), safeSQLName(table.Name)) + var name, ddl string + err := db.QueryRowxContext(ctx, q).Scan(&name, &ddl) + if err != nil { + return err + } + table.DDL = ddl + return nil +} + // InformationSchema interface implementation for drivers.InformationSchema var _ drivers.InformationSchema = (*informationSchemaImpl)(nil) diff --git a/runtime/drivers/starrocks/olap_test.go b/runtime/drivers/starrocks/olap_test.go index 9cf07fe83a7..b9812e0ed66 100644 --- a/runtime/drivers/starrocks/olap_test.go +++ b/runtime/drivers/starrocks/olap_test.go @@ -143,6 +143,10 @@ func TestStarRocksOLAP(t *testing.T) { t.Run("DecimalPrecision", func(t *testing.T) { testDecimalPrecision(t, olap) }) + + t.Run("LoadDDL", func(t *testing.T) { + testLoadDDL(t, conn) + }) } func testVarcharNotBinary(t *testing.T, olap drivers.OLAPStore) { @@ -1299,6 +1303,19 @@ func testAllTypesOutput(t *testing.T, olap drivers.OLAPStore) { t.Log("================================================================================") } +func testLoadDDL(t *testing.T, conn drivers.Handle) { + olap, _ := conn.AsOLAP("default") + ctx := context.Background() + + // Test DDL for the all_types table + table, err := olap.InformationSchema().Lookup(ctx, "", "test_db", "all_types") + require.NoError(t, err) + err = olap.InformationSchema().LoadDDL(ctx, table) + require.NoError(t, err) + require.Contains(t, table.DDL, "CREATE TABLE") + require.Contains(t, table.DDL, "all_types") +} + // formatValue formats a value for display, truncating long strings func formatValue(val any) string { if val == nil { diff --git a/runtime/pkg/rduckdb/db.go b/runtime/pkg/rduckdb/db.go index ca849f1d2ba..e9b97a02b69 100644 --- a/runtime/pkg/rduckdb/db.go +++ b/runtime/pkg/rduckdb/db.go @@ -68,6 +68,10 @@ type DB interface { // Schema returns the schema of the database. Schema(ctx context.Context, ilike, name string, pageSize uint32, pageToken string) ([]*Table, string, error) + + // DDL returns the DDL (CREATE statement) for the named table or view. Returns "" if unavailable. + // The database and schema parameters are optional; when non-empty they qualify the lookup. + DDL(ctx context.Context, database, schema, name string) (string, error) } type DBOptions struct { diff --git a/runtime/pkg/rduckdb/generic.go b/runtime/pkg/rduckdb/generic.go index 72d17e2ffe7..05cf9053e49 100644 --- a/runtime/pkg/rduckdb/generic.go +++ b/runtime/pkg/rduckdb/generic.go @@ -400,6 +400,41 @@ func (m *generic) RenameTable(ctx context.Context, oldName, newName string) (res return err } +// DDL implements DB. +func (m *generic) DDL(ctx context.Context, database, schema, name string) (string, error) { + conn, err := m.acquireConn(ctx) + if err != nil { + return "", err + } + defer func() { _ = conn.Close() }() + + var q string + var args []any + if database != "" && schema != "" { + q = ` + SELECT sql FROM duckdb_tables() WHERE database_name = ? AND schema_name = ? AND table_name = ? + UNION ALL + SELECT sql FROM duckdb_views() WHERE database_name = ? AND schema_name = ? AND view_name = ? + ` + args = []any{database, schema, name, database, schema, name} + } else { + // Fall back to current_database()/current_schema() to avoid collisions across attached databases. + q = ` + SELECT sql FROM duckdb_tables() WHERE database_name = current_database() AND schema_name = current_schema() AND table_name = ? + UNION ALL + SELECT sql FROM duckdb_views() WHERE database_name = current_database() AND schema_name = current_schema() AND view_name = ? + ` + args = []any{name, name} + } + + var sqlStr *string + err = conn.QueryRowxContext(ctx, q, args...).Scan(&sqlStr) + if err != nil || sqlStr == nil { + return "", nil + } + return *sqlStr, nil +} + // Schema implements DB. func (m *generic) Schema(ctx context.Context, ilike, name string, pageSize uint32, pageToken string) ([]*Table, string, error) { conn, err := m.acquireConn(ctx) diff --git a/runtime/pkg/rduckdb/information_schema.go b/runtime/pkg/rduckdb/information_schema.go index 04862153661..d73e878be5e 100644 --- a/runtime/pkg/rduckdb/information_schema.go +++ b/runtime/pkg/rduckdb/information_schema.go @@ -19,6 +19,28 @@ type Table struct { SizeBytes int64 `db:"-"` } +func (d *db) DDL(ctx context.Context, database, schema, name string) (string, error) { + connx, release, err := d.AcquireReadConnection(ctx) + if err != nil { + return "", err + } + defer func() { _ = release() }() + + // We disregard database and schema since they're not applicable here due to how we use them for table versioning (see treatment in Schema(...) below). + q := ` + SELECT sql FROM duckdb_tables() WHERE table_name = ? + UNION ALL + SELECT sql FROM duckdb_views() WHERE view_name = ? + ` + + var sqlStr *string + err = connx.QueryRowxContext(ctx, q, name, name).Scan(&sqlStr) + if err != nil || sqlStr == nil { + return "", nil + } + return *sqlStr, nil +} + func (d *db) Schema(ctx context.Context, ilike, name string, pageSize uint32, pageToken string) ([]*Table, string, error) { if ilike != "" && name != "" { return nil, "", fmt.Errorf("cannot specify both `ilike` and `name`")