Tech Debt: Merge MergeConstraints and ApplyShardedRecommendations tra…

…nsformation on exported schema (#2366)

- now the parsing of sql file happens once and then one by one all transformations are applied then Deparse the final schema file

* Added unit tests
* Shifted tests from exportSchema_test.go to transformer_test.go

yb-voyager/cmd/common.go

@@ -1210,6 +1210,10 @@ type TargetSizingRecommendations struct {
 //====== AssesmentReport struct methods ======//
 
 func ParseJSONToAssessmentReport(reportPath string) (*AssessmentReport, error) {
+	if !utils.FileOrFolderExists(reportPath) {
+		return nil, fmt.Errorf("report file %q does not exist", reportPath)
+	}
+
 	var report AssessmentReport
 	err := jsonfile.NewJsonFile[AssessmentReport](reportPath).Load(&report)
 	if err != nil {

yb-voyager/cmd/constants.go

@@ -61,7 +61,6 @@ const (
 	ROW_UPDATE_STATUS_NOT_STARTED   = 0
 	ROW_UPDATE_STATUS_IN_PROGRESS   = 1
 	ROW_UPDATE_STATUS_COMPLETED     = 3
-	COLOCATION_CLAUSE               = "colocation"
 	//phase names used in call-home payload
 	ANALYZE_PHASE                    = "analyze-schema"
 	EXPORT_SCHEMA_PHASE              = "export-schema"

yb-voyager/src/adaptiveparallelism/adaptive_parallelism_test.go

@@ -18,9 +18,11 @@ limitations under the License.
 package adaptiveparallelism
 
 import (
+	"os"
 	"strconv"
 	"testing"
 
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/yugabyte/yb-voyager/yb-voyager/src/tgtdb"
 )
@@ -112,6 +114,14 @@ func (d *dummyTargetYugabyteDB) UpdateNumConnectionsInPool(delta int) error {
 	return nil
 }
 
+func TestMain(m *testing.M) {
+	// to avoid info level logs flooding the unit test output
+	log.SetLevel(log.WarnLevel)
+
+	exitCode := m.Run()
+	os.Exit(exitCode)
+}
+
 func TestMaxCpuUsage(t *testing.T) {
 	yb := &dummyTargetYugabyteDB{
 		size:          3,

yb-voyager/src/constants/constants.go

@@ -49,5 +49,6 @@ const (
 )
 
 const (
-	OBFUSCATE_STRING = "XXXXX"
+	OBFUSCATE_STRING  = "XXXXX"
+	COLOCATION_CLAUSE = "colocation"
 )

yb-voyager/src/query/queryparser/query_parser.go

@@ -54,7 +54,7 @@ func ParsePLPGSQLToJson(query string) (string, error) {
 	return jsonString, err
 }
 
-func ParseSqlFile(filePath string) (*pg_query.ParseResult, error) {
+func ParseSqlFile(filePath string) ([]*pg_query.RawStmt, error) {
 	log.Debugf("parsing the file [%s]", filePath)
 	bytes, err := os.ReadFile(filePath)
 	if err != nil {
@@ -67,7 +67,7 @@ func ParseSqlFile(filePath string) (*pg_query.ParseResult, error) {
 	}
 	log.Debugf("sql file contents: %s\n", string(bytes))
 	log.Debugf("parse tree: %v\n", tree)
-	return tree, nil
+	return tree.Stmts, nil
 }
 
 func ProcessDDL(parseTree *pg_query.ParseResult) (DDLObject, error) {

yb-voyager/src/query/sqltransformer/helpers.go

@@ -0,0 +1,67 @@
+/*
+Copyright (c) YugabyteDB, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package sqltransformer
+
+import (
+	pg_query "github.com/pganalyze/pg_query_go/v6"
+	"github.com/yugabyte/yb-voyager/yb-voyager/src/constants"
+)
+
+// addColocatedOption adds "WITH (colocated = true)" to the "Options" array
+// for either a CreateStmt or CreateMatViewStmt.
+func addColocationOptionToCreateTable(createStmt *pg_query.CreateStmt) {
+	if createStmt == nil {
+		return
+	}
+	// If Options slice is nil, initialize it
+	if createStmt.Options == nil {
+		createStmt.Options = []*pg_query.Node{}
+	}
+
+	// Build DefElem: defname = "colocated", arg = "false"
+	defElemNode := &pg_query.Node{
+		Node: &pg_query.Node_DefElem{
+			DefElem: &pg_query.DefElem{
+				Defname: constants.COLOCATION_CLAUSE,
+				Arg:     pg_query.MakeStrNode("false"),
+			},
+		},
+	}
+
+	createStmt.Options = append(createStmt.Options, defElemNode)
+}
+
+func addColocationOptionToCreateMaterializedView(createMatViewStmt *pg_query.CreateTableAsStmt) {
+	if createMatViewStmt == nil {
+		return
+	}
+	// If Options slice is nil, initialize it
+	if createMatViewStmt.Into.Options == nil {
+		createMatViewStmt.Into.Options = []*pg_query.Node{}
+	}
+
+	// Build DefElem: defname = "colocated", arg = "false"
+	defElemNode := &pg_query.Node{
+		Node: &pg_query.Node_DefElem{
+			DefElem: &pg_query.DefElem{
+				Defname: constants.COLOCATION_CLAUSE,
+				Arg:     pg_query.MakeStrNode("false"),
+			},
+		},
+	}
+
+	createMatViewStmt.Into.Options = append(createMatViewStmt.Into.Options, defElemNode)
+}

yb-voyager/src/query/sqltransformer/transformer.go

@@ -20,6 +20,7 @@ import (
 	"slices"
 
 	pg_query "github.com/pganalyze/pg_query_go/v6"
+	log "github.com/sirupsen/logrus"
 	"github.com/yugabyte/yb-voyager/yb-voyager/src/query/queryparser"
 )
 
@@ -63,7 +64,9 @@ Note: Need to keep the relative ordering of statements(tables) intact.
 Because there can be cases like Foreign Key constraints that depend on the order of tables.
 */
 func (t *Transformer) MergeConstraints(stmts []*pg_query.RawStmt) ([]*pg_query.RawStmt, error) {
-	// TODO: Ensure removing all the ALTER stmts which are merged into CREATE. No duplicates.
+	if len(stmts) == 0 {
+		return stmts, nil
+	}
 
 	createStmtMap := make(map[string]*pg_query.RawStmt)
 	for _, stmt := range stmts {
@@ -153,3 +156,39 @@ func (t *Transformer) MergeConstraints(stmts []*pg_query.RawStmt) ([]*pg_query.R
 
 	return result, nil
 }
+
+// write a tranformation function which converts the given tables into Sharded table by adding clause WITH (colocated = true)
+func (t *Transformer) ConvertToShardedTables(stmts []*pg_query.RawStmt, isObjectSharded func(objectName string) bool) ([]*pg_query.RawStmt, error) {
+	if len(stmts) == 0 {
+		return stmts, nil
+	}
+
+	var result []*pg_query.RawStmt
+	for _, stmt := range stmts {
+		stmtType := queryparser.GetStatementType(stmt.Stmt.ProtoReflect())
+
+		switch stmtType {
+		case queryparser.PG_QUERY_CREATE_STMT: // CREATE TABLE case
+			objectName := queryparser.GetObjectNameFromRangeVar(stmt.Stmt.GetCreateStmt().Relation)
+			if isObjectSharded(objectName) {
+				log.Infof("adding colocation option to CREATE TABLE for object %v", objectName)
+				addColocationOptionToCreateTable(stmt.Stmt.GetCreateStmt())
+			}
+
+			result = append(result, stmt)
+		case queryparser.PG_QUERY_CREATE_TABLE_AS_STMT: // CREATE MATERIALIZED VIEW case
+			objectName := queryparser.GetObjectNameFromRangeVar(stmt.Stmt.GetCreateTableAsStmt().Into.Rel)
+			if isObjectSharded(objectName) {
+				log.Infof("adding colocation option to CREATE MATERIALIZED VIEW for object %v", objectName)
+				addColocationOptionToCreateMaterializedView(stmt.Stmt.GetCreateTableAsStmt())
+			}
+
+			result = append(result, stmt)
+		default:
+			result = append(result, stmt)
+		}
+
+	}
+
+	return result, nil
+}