diff --git a/mc2mc/internal/client/odps.go b/mc2mc/internal/client/odps.go index 788ee00..698b040 100644 --- a/mc2mc/internal/client/odps.go +++ b/mc2mc/internal/client/odps.go @@ -39,8 +39,7 @@ func NewODPSClient(logger *slog.Logger, client *odps.Odps) *odpsClient { // when context is cancelled. func (c *odpsClient) ExecSQL(ctx context.Context, query string, additionalHints map[string]string) error { if c.isDryRun { - c.logger.Info("dry run mode, skipping execution") - return nil + c.logger.Info("[DRY-RUN] running query in dry-run mode with EXPLAIN.") } hints := addHints(additionalHints, query) diff --git a/mc2mc/mc2mc.go b/mc2mc/mc2mc.go index ca9cf78..79b134e 100644 --- a/mc2mc/mc2mc.go +++ b/mc2mc/mc2mc.go @@ -17,7 +17,7 @@ import ( "github.com/goto/transformers/mc2mc/internal/client" "github.com/goto/transformers/mc2mc/internal/config" "github.com/goto/transformers/mc2mc/internal/logger" - "github.com/goto/transformers/mc2mc/internal/query" + "github.com/goto/transformers/mc2mc/pkg/query" ) func mc2mc(envs []string) error { @@ -84,6 +84,7 @@ func mc2mc(envs []string) error { query.WithPartitionValue(cfg.DevEnablePartitionValue == "true"), query.WithCostAttributionLabel(cfg.CostAttributionTeam), query.WithColumnOrder(), + query.WithDryRun(cfg.DryRun), ).Build() if err != nil { return errors.WithStack(err) @@ -100,6 +101,7 @@ func mc2mc(envs []string) error { query.WithPartitionValue(cfg.DevEnablePartitionValue == "true"), query.WithCostAttributionLabel(cfg.CostAttributionTeam), query.WithColumnOrder(), + query.WithDryRun(cfg.DryRun), ) // -- TODO(START): refactor this part -- @@ -155,6 +157,7 @@ func mc2mc(envs []string) error { query.WithQuery(string(raw)), query.WithCostAttributionLabel(cfg.CostAttributionTeam), query.WithMethod(query.MERGE), + query.WithDryRun(cfg.DryRun), ).Build() if err != nil { return errors.WithStack(err) diff --git a/mc2mc/internal/query/builder.go b/mc2mc/pkg/query/builder.go similarity index 89% rename from mc2mc/internal/query/builder.go rename to mc2mc/pkg/query/builder.go index 1e95df1..06710e9 100644 --- a/mc2mc/internal/query/builder.go +++ b/mc2mc/pkg/query/builder.go @@ -30,6 +30,7 @@ type Builder struct { enableAutoPartition bool enablePartitionValue bool + enableDryRun bool } // NewBuilder creates a new query builder with the given options @@ -68,12 +69,6 @@ func (b *Builder) Build() (string, error) { if b.method == MERGE { // split query components hrs, vars, queries := SplitQueryComponents(b.query) - if len(queries) <= 1 { - if b.costAttributionTeam != "" { - return fmt.Sprintf("%s\n%s", b.query, getCostAttributionComment(b.costAttributionTeam)), nil - } - return b.query, nil - } query := b.constructMergeQuery(hrs, vars, queries) return query, nil } @@ -141,18 +136,30 @@ func (b *Builder) Build() (string, error) { } } + if b.enableDryRun { + // append explain to the query for dry run + query = fmt.Sprintf("EXPLAIN\n%s", query) + } + // construct final query with headers, drops, variables and udfs if hr != "" { hr += "\n" } - if drops != "" { - drops += "\n" + dropsStr := "" + if len(drops) > 0 { + if b.enableDryRun { + for i, drop := range drops { + drops[i] = fmt.Sprintf("EXPLAIN\n%s", drop) + } + } + dropsStr = strings.Join(drops, "\n;\n") + dropsStr += "\n;\n" } if varsAndUDFs != "" { varsAndUDFs += "\n" } - query = fmt.Sprintf("%s%s%s%s", hr, drops, varsAndUDFs, query) + query = fmt.Sprintf("%s%s%s%s", hr, dropsStr, varsAndUDFs, query) if b.costAttributionTeam != "" { query = fmt.Sprintf("%s\n%s\n", query, getCostAttributionComment(b.costAttributionTeam)) } @@ -201,6 +208,14 @@ func (b *Builder) constructOverridedValues(query string) (string, error) { // constructMergeQueries constructs merge queries with headers and variables func (b *Builder) constructMergeQuery(hrs, vars, queries []string) string { + if !b.enableDryRun && len(queries) == 1 { + // maintaining existing logic for single query + if b.costAttributionTeam != "" { + return fmt.Sprintf("%s\n%s", b.query, getCostAttributionComment(b.costAttributionTeam)) + } + return b.query + } + builder := strings.Builder{} for i, q := range queries { q = strings.TrimSpace(q) @@ -215,6 +230,10 @@ func (b *Builder) constructMergeQuery(hrs, vars, queries []string) string { if variables != "" && !IsDDL(q) { // skip variables if it's ddl builder.WriteString(fmt.Sprintf("%s\n", variables)) } + if b.enableDryRun { + // append explain to the DDL query part + builder.WriteString("EXPLAIN\n") + } builder.WriteString(fmt.Sprintf("%s\n;", q)) if b.costAttributionTeam != "" { builder.WriteString(fmt.Sprintf("\n%s\n", getCostAttributionComment(b.costAttributionTeam))) diff --git a/mc2mc/internal/query/builder_test.go b/mc2mc/pkg/query/builder_test.go similarity index 78% rename from mc2mc/internal/query/builder_test.go rename to mc2mc/pkg/query/builder_test.go index 6dc8a94..64292df 100644 --- a/mc2mc/internal/query/builder_test.go +++ b/mc2mc/pkg/query/builder_test.go @@ -5,7 +5,7 @@ import ( "testing" "github.com/goto/transformers/mc2mc/internal/logger" - "github.com/goto/transformers/mc2mc/internal/query" + "github.com/goto/transformers/mc2mc/pkg/query" "github.com/stretchr/testify/assert" ) @@ -994,6 +994,309 @@ SET append_test.id = 2 }) } +func TestBuilder_BuildWithDryRun(t *testing.T) { + t.Run("returns dry run query for merge load method", func(t *testing.T) { + queryToExecute := `DELETE FROM project.playground.table WHERE id = 1; +INSERT INTO project.playground.table SELECT * FROM project.playground.table_source;` + odspClient := &mockOdpsClient{} + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.MERGE), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Equal(t, `EXPLAIN +DELETE FROM project.playground.table WHERE id = 1 +; +--*--optimus-break-marker--*-- +EXPLAIN +INSERT INTO project.playground.table SELECT * FROM project.playground.table_source +;`, query) + }) + + t.Run("returns dry run query for APPEND load method", func(t *testing.T) { + queryToExecute := `select * from project.playground.table;` + odspClient := &mockOdpsClient{ + orderedColumns: func() ([]string, error) { + return []string{"col1", "col2", "_partitiontime"}, nil + }, + partitionResult: func() ([]string, error) { + return []string{}, nil + }, + } + destinationTableID := "project.playground.table_destination" + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.APPEND), + query.WithDestination(destinationTableID), + query.WithColumnOrder(), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Equal(t, `EXPLAIN +INSERT INTO TABLE project.playground.table_destination +SELECT col1, col2, _partitiontime FROM ( +select * from project.playground.table +) +;`, query) + }) + + t.Run("returns dry run query for replace load method", func(t *testing.T) { + queryToExecute := `select * from project.playground.table;` + odspClient := &mockOdpsClient{ + orderedColumns: func() ([]string, error) { + return []string{"col1", "col2", "_partitiontime"}, nil + }, + partitionResult: func() ([]string, error) { + return []string{}, nil + }, + } + destinationTableID := "project.playground.table_destination" + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.REPLACE), + query.WithDestination(destinationTableID), + query.WithColumnOrder(), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Equal(t, `EXPLAIN +INSERT OVERWRITE TABLE project.playground.table_destination +SELECT col1, col2, _partitiontime FROM ( +select * from project.playground.table +) +;`, query) + }) + + t.Run("returns dry run query for append load method with partition", func(t *testing.T) { + queryToExecute := `select * from project.playground.table;` + odspClient := &mockOdpsClient{ + orderedColumns: func() ([]string, error) { + return []string{"col1", "col2", "_partitiontime"}, nil + }, + partitionResult: func() ([]string, error) { + return []string{"col3"}, nil + }, + } + destinationTableID := "project.playground.table_destination" + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.APPEND), + query.WithDestination(destinationTableID), + query.WithColumnOrder(), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Equal(t, `EXPLAIN +INSERT INTO TABLE project.playground.table_destination PARTITION (col3) +SELECT col1, col2, _partitiontime FROM ( +select * from project.playground.table +) +;`, query) + }) + + t.Run("returns dry run query with drops and variables", func(t *testing.T) { + queryToExecute := `SET odps.sql.decimal.odps2=true; +DROP TABLE IF EXISTS project.playground.temp_table; +DROP TABLE IF EXISTS project.playground.temp_table_2; +@variable := SELECT 1 as value; +SELECT * FROM project.playground.table;` + odspClient := &mockOdpsClient{ + orderedColumns: func() ([]string, error) { + return []string{"col1", "col2"}, nil + }, + partitionResult: func() ([]string, error) { + return []string{}, nil + }, + } + destinationTableID := "project.playground.table_destination" + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.APPEND), + query.WithDestination(destinationTableID), + query.WithColumnOrder(), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Equal(t, `SET odps.sql.decimal.odps2=true +; +EXPLAIN +DROP TABLE IF EXISTS project.playground.temp_table +; +EXPLAIN +DROP TABLE IF EXISTS project.playground.temp_table_2 +; +@variable := SELECT 1 as value +; +EXPLAIN +INSERT INTO TABLE project.playground.table_destination +SELECT col1, col2 FROM ( +SELECT * FROM project.playground.table +) +;`, query) + }) + + t.Run("returns dry run query with overridden values", func(t *testing.T) { + queryToExecute := `select * from project.playground.table;` + odspClient := &mockOdpsClient{ + orderedColumns: func() ([]string, error) { + return []string{"col1", "col2", "_partitiontime"}, nil + }, + partitionResult: func() ([]string, error) { + return []string{}, nil + }, + } + destinationTableID := "project.playground.table_destination" + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.APPEND), + query.WithDestination(destinationTableID), + query.WithOverridedValue("_partitiontime", "TIMESTAMP('2021-01-01')"), + query.WithColumnOrder(), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Equal(t, `EXPLAIN +INSERT INTO TABLE project.playground.table_destination +SELECT col1, col2, _partitiontime FROM ( +SELECT col1, col2, TIMESTAMP('2021-01-01') as _partitiontime FROM ( +select * from project.playground.table +) +) +;`, query) + }) + + t.Run("returns dry run query with cost attribution", func(t *testing.T) { + queryToExecute := `select * from project.playground.table;` + odspClient := &mockOdpsClient{ + orderedColumns: func() ([]string, error) { + return []string{"col1", "col2"}, nil + }, + partitionResult: func() ([]string, error) { + return []string{}, nil + }, + } + destinationTableID := "project.playground.table_destination" + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.APPEND), + query.WithDestination(destinationTableID), + query.WithCostAttributionLabel("test-team"), + query.WithColumnOrder(), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.Contains(t, query, "EXPLAIN") + assert.Contains(t, query, "--cost_attribution_team=test-team") + }) + + t.Run("returns error for empty query in dry run mode", func(t *testing.T) { + odspClient := &mockOdpsClient{} + + queryToExecute, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithDryRun(true), + ).Build() + + assert.Error(t, err) + assert.Empty(t, queryToExecute) + assert.Contains(t, err.Error(), "query is required") + }) + + t.Run("returns dry run query for merge method with complex query components", func(t *testing.T) { + queryToExecute := `SET odps.table.append2.enable=true; +@src := SELECT 1 id; +@src2 := SELECT id FROM append_tmp; +CREATE TABLE append_tmp AS SELECT * FROM @src; +CREATE TABLE append_tmp2(id bigint); +MERGE INTO append_test +USING (SELECT * FROM @src2) source +on append_test.id = source.id +WHEN MATCHED THEN UPDATE +SET append_test.id = 2;` + odspClient := &mockOdpsClient{} + + query, err := query.NewBuilder( + logger.NewDefaultLogger(), + odspClient, + query.WithQuery(queryToExecute), + query.WithMethod(query.MERGE), + query.WithCostAttributionLabel("test-team"), + query.WithDryRun(true), + ).Build() + + assert.NoError(t, err) + assert.NoError(t, err) + assert.Equal(t, `SET odps.table.append2.enable=true +; +@src := SELECT 1 id +; +@src2 := SELECT id FROM append_tmp +; +EXPLAIN +CREATE TABLE append_tmp AS SELECT * FROM @src +; +--cost_attribution_team=test-team + + +--*--optimus-break-marker--*-- +SET odps.table.append2.enable=true +; +EXPLAIN +CREATE TABLE append_tmp2(id bigint) +; +--cost_attribution_team=test-team + + +--*--optimus-break-marker--*-- +SET odps.table.append2.enable=true +; +@src := SELECT 1 id +; +@src2 := SELECT id FROM append_tmp +; +EXPLAIN +MERGE INTO append_test +USING (SELECT * FROM @src2) source +on append_test.id = source.id +WHEN MATCHED THEN UPDATE +SET append_test.id = 2 +; +--cost_attribution_team=test-team + +`, query) + }) +} + type mockOdpsClient struct { partitionResult func() ([]string, error) execSQLResult func() error diff --git a/mc2mc/internal/query/helper.go b/mc2mc/pkg/query/helper.go similarity index 96% rename from mc2mc/internal/query/helper.go rename to mc2mc/pkg/query/helper.go index 050ecc2..bf6cceb 100644 --- a/mc2mc/internal/query/helper.go +++ b/mc2mc/pkg/query/helper.go @@ -165,7 +165,7 @@ func SeparateVariablesUDFsAndQuery(query string) (string, string) { return variableUDFStr, queryStr } -func SeparateDropsAndQuery(query string) (string, string) { +func SeparateDropsAndQuery(query string) ([]string, string) { drops := []string{} query = strings.TrimSpace(query) remainingQueries := []string{} @@ -179,7 +179,7 @@ func SeparateDropsAndQuery(query string) (string, string) { } stmtWithoutComment := RemoveComments(stmt) if dropPattern.MatchString(strings.TrimSpace(stmtWithoutComment)) { - drops = append(drops, stmt) + drops = append(drops, strings.TrimSpace(stmt)) } else if strings.TrimSpace(stmtWithoutComment) == "" { // if the statement is empty, it's a comment, then omit it // since it doesn't make sense to execute this statement @@ -188,19 +188,10 @@ func SeparateDropsAndQuery(query string) (string, string) { } } - dropStr := "" - if len(drops) > 0 { - for i, drop := range drops { - drops[i] = strings.TrimSpace(drop) - } - dropStr = strings.Join(drops, "\n;\n") - dropStr += "\n;" - } - // join the remaining queries back together queryStr := strings.Join(remainingQueries, "\n;\n") - return dropStr, queryStr + return drops, queryStr } func RemoveComments(query string) string { diff --git a/mc2mc/internal/query/helper_test.go b/mc2mc/pkg/query/helper_test.go similarity index 99% rename from mc2mc/internal/query/helper_test.go rename to mc2mc/pkg/query/helper_test.go index 655bf8c..ec6e16d 100644 --- a/mc2mc/internal/query/helper_test.go +++ b/mc2mc/pkg/query/helper_test.go @@ -5,7 +5,7 @@ import ( "github.com/stretchr/testify/assert" - "github.com/goto/transformers/mc2mc/internal/query" + "github.com/goto/transformers/mc2mc/pkg/query" ) func TestSplitQueryComponents(t *testing.T) { @@ -375,7 +375,7 @@ on append_test.id = source.id WHEN MATCHED THEN UPDATE SET append_test.id = 2;` drops, query := query.SeparateDropsAndQuery(q1) - assert.Equal(t, "DROP TABLE IF EXISTS append_test\n;", drops) + assert.Equal(t, []string{"DROP TABLE IF EXISTS append_test"}, drops) assert.Equal(t, `MERGE INTO append_test USING (SELECT * FROM @src) source on append_test.id = source.id diff --git a/mc2mc/internal/query/method.go b/mc2mc/pkg/query/method.go similarity index 100% rename from mc2mc/internal/query/method.go rename to mc2mc/pkg/query/method.go diff --git a/mc2mc/internal/query/option.go b/mc2mc/pkg/query/option.go similarity index 91% rename from mc2mc/internal/query/option.go rename to mc2mc/pkg/query/option.go index 3c03d4c..f1f5ab8 100644 --- a/mc2mc/internal/query/option.go +++ b/mc2mc/pkg/query/option.go @@ -52,3 +52,9 @@ func WithCostAttributionLabel(teamName string) Option { b.costAttributionTeam = teamName } } + +func WithDryRun(enable bool) Option { + return func(b *Builder) { + b.enableDryRun = enable + } +}