feat: support MySQL CTAS

apecloud · Jan 7, 2025 · 3b5bc24 · 3b5bc24
1 parent 052e742
commit 3b5bc24
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 33 deletions.
diff --git 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
@@ -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;",
 	}