From 3b5bc2459c04f4454710602e24ef573547257194 Mon Sep 17 00:00:00 2001 From: Fan Yang Date: Tue, 7 Jan 2025 19:19:09 +0800 Subject: [PATCH] feat: support MySQL CTAS --- backend/executor.go | 46 ++++++++++++++++++++++++++++++--------------- main_test.go | 36 +++++++++++++++++------------------ 2 files changed, 49 insertions(+), 33 deletions(-) diff --git a/backend/executor.go b/backend/executor.go index 6f7ec30..9953065 100644 --- a/backend/executor.go +++ b/backend/executor.go @@ -61,10 +61,13 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (sql.Row } n := root - ctx.GetLogger().WithFields(logrus.Fields{ - "Query": ctx.Query(), - "NodeType": fmt.Sprintf("%T", n), - }).Traceln("Building node:", n) + + if log := ctx.GetLogger(); log.Logger.IsLevelEnabled(logrus.TraceLevel) { + log.WithFields(logrus.Fields{ + "Query": ctx.Query(), + "NodeType": fmt.Sprintf("%T", n), + }).Traceln("Building node:", n) + } // TODO; find a better way to fallback to the base builder switch n.(type) { @@ -114,7 +117,13 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (sql.Row } // Fallback to the base builder if the plan contains system/user variables or is not a pure data query. - if containsVariable(n) || !IsPureDataQuery(n) { + tree := n + switch n := n.(type) { + case *plan.TableCopier: + tree = n.Source + } + if containsVariable(tree) || !IsPureDataQuery(tree) { + ctx.GetLogger().Traceln("Falling back to the base builder") return b.base.Build(ctx, root, r) } @@ -133,11 +142,14 @@ func (b *DuckBuilder) Build(ctx *sql.Context, root sql.Node, r sql.Row) (sql.Row return nil, err } return b.base.Build(ctx, root, r) - // SubqueryAlias is for select * from view + // ResolvedTable is for `SELECT * FROM table` and `TABLE table` + // SubqueryAlias is for `SELECT * FROM view`` case *plan.ResolvedTable, *plan.SubqueryAlias, *plan.TableAlias: return b.executeQuery(ctx, node, conn) case *plan.Distinct, *plan.OrderedDistinct: return b.executeQuery(ctx, node, conn) + case *plan.TableCopier: + return b.executeDML(ctx, node, conn) case sql.Expressioner: return b.executeExpressioner(ctx, node, conn) case *plan.DeleteFrom: @@ -174,7 +186,7 @@ func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, conn *stdsql.Co case *plan.ShowTables: duckSQL = ctx.Query() case *plan.ResolvedTable: - // SQLGlot cannot translate MySQL's `TABLE t` into DuckDB's `FROM t` - it produces `"table" AS t` instead. + // SQLGlot cannot translate MySQL's `TABLE t` into DuckDB's `FROM t` - it produces `"table" AS t` instead. duckSQL = `FROM ` + catalog.ConnectIdentifiersANSI(n.Database().Name(), n.Name()) default: duckSQL, err = transpiler.TranslateWithSQLGlot(ctx.Query()) @@ -183,10 +195,12 @@ func (b *DuckBuilder) executeQuery(ctx *sql.Context, n sql.Node, conn *stdsql.Co return nil, catalog.ErrTranspiler.New(err) } - ctx.GetLogger().WithFields(logrus.Fields{ - "Query": ctx.Query(), - "DuckSQL": duckSQL, - }).Trace("Executing Query...") + if log := ctx.GetLogger(); log.Logger.IsLevelEnabled(logrus.TraceLevel) { + log.WithFields(logrus.Fields{ + "Query": ctx.Query(), + "DuckSQL": duckSQL, + }).Trace("Executing Query...") + } // Execute the DuckDB query rows, err := conn.QueryContext(ctx.Context, duckSQL) @@ -204,10 +218,12 @@ func (b *DuckBuilder) executeDML(ctx *sql.Context, n sql.Node, conn *stdsql.Conn return nil, catalog.ErrTranspiler.New(err) } - ctx.GetLogger().WithFields(logrus.Fields{ - "Query": ctx.Query(), - "DuckSQL": duckSQL, - }).Trace("Executing DML...") + if log := ctx.GetLogger(); log.Logger.IsLevelEnabled(logrus.TraceLevel) { + log.WithFields(logrus.Fields{ + "Query": ctx.Query(), + "DuckSQL": duckSQL, + }).Trace("Executing DML...") + } // Execute the DuckDB query result, err := conn.ExecContext(ctx.Context, duckSQL) diff --git a/main_test.go b/main_test.go index 178ebe3..600bc8d 100644 --- a/main_test.go +++ b/main_test.go @@ -1134,24 +1134,24 @@ func TestCreateTable(t *testing.T) { "CREATE_TABLE_tt_(pk_int_primary_key,_d_datetime(6)_default_current_timestamp(6))", "Identifier_lengths", "table_charset_options", - "show_create_table_t3", - "show_create_table_t4", - "create_table_with_select_preserves_default", - "create_table_t1_select_*_from_a;", - "create_table_t2_select_j_from_a;", - "create_table_t3_select_j_as_i_from_a;", - "create_table_t4_select_j_+_1_from_a;", - "create_table_t5_select_a.j_from_a;", - "create_table_t6_select_sqa.j_from_(select_i,_j_from_a)_sqa;", - "show_create_table_t7;", - "create_table_t8_select_*_from_(select_*_from_a)_a_join_(select_*_from_b)_b;", - "show_create_table_t9;", - "create_table_t11_select_sum(j)_over()_as_jj_from_a;", - "create_table_t12_select_j_from_a_group_by_j;", - "create_table_t13_select_*_from_c;", - "event_contains_CREATE_TABLE_AS", - "CREATE_EVENT_foo_ON_SCHEDULE_EVERY_1_YEAR_DO_CREATE_TABLE_bar_AS_SELECT_1;", - "trigger_contains_CREATE_TABLE_AS", + // "show_create_table_t3", + // "show_create_table_t4", + // "create_table_with_select_preserves_default", + // "create_table_t1_select_*_from_a;", + // "create_table_t2_select_j_from_a;", + // "create_table_t3_select_j_as_i_from_a;", + // "create_table_t4_select_j_+_1_from_a;", + // "create_table_t5_select_a.j_from_a;", + // "create_table_t6_select_sqa.j_from_(select_i,_j_from_a)_sqa;", + // "show_create_table_t7;", + // "create_table_t8_select_*_from_(select_*_from_a)_a_join_(select_*_from_b)_b;", + // "show_create_table_t9;", + // "create_table_t11_select_sum(j)_over()_as_jj_from_a;", + // "create_table_t12_select_j_from_a_group_by_j;", + // "create_table_t13_select_*_from_c;", + // "event_contains_CREATE_TABLE_AS", + // "CREATE_EVENT_foo_ON_SCHEDULE_EVERY_1_YEAR_DO_CREATE_TABLE_bar_AS_SELECT_1;", + // "trigger_contains_CREATE_TABLE_AS", "CREATE_TRIGGER_foo_AFTER_UPDATE_ON_t_FOR_EACH_ROW_BEGIN_CREATE_TABLE_bar_AS_SELECT_1;_END;", }