Pulling the GenerateSqlPatchSchemaStatements into the sqlfmt package …

…(and cleaning up pkg import cycle)
dolthub · May 17, 2024 · 34b6539 · 34b6539
1 parent 02b3834
commit 34b6539
Show file tree

Hide file tree

Showing 4 changed files with 163 additions and 162 deletions.
diff --git a/go/libraries/doltcore/diff/table_deltas.go b/go/libraries/doltcore/diff/table_deltas.go
@@ -24,7 +24,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
-	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/sqlfmt"
 	"github.com/dolthub/dolt/go/libraries/utils/set"
 	"github.com/dolthub/dolt/go/store/prolly/tree"
 	"github.com/dolthub/dolt/go/store/types"
@@ -620,36 +619,6 @@ func (td TableDelta) GetRowData(ctx context.Context) (from, to durable.Index, er
 	return from, to, nil
 }
 
-// GetDataDiffStatement returns any data diff in SQL statements for given table including INSERT, UPDATE and DELETE row statements.
-func GetDataDiffStatement(tableName string, sch schema.Schema, row sql.Row, rowDiffType ChangeType, colDiffTypes []ChangeType) (string, error) {
-	if len(row) != len(colDiffTypes) {
-		return "", fmt.Errorf("expected the same size for columns and diff types, got %d and %d", len(row), len(colDiffTypes))
-	}
-
-	switch rowDiffType {
-	case Added:
-		return sqlfmt.SqlRowAsInsertStmt(row, tableName, sch)
-	case Removed:
-		return sqlfmt.SqlRowAsDeleteStmt(row, tableName, sch, 0)
-	case ModifiedNew:
-		updatedCols := set.NewEmptyStrSet()
-		for i, diffType := range colDiffTypes {
-			if diffType != None {
-				updatedCols.Add(sch.GetAllCols().GetByIndex(i).Name)
-			}
-		}
-		if updatedCols.Size() == 0 {
-			return "", nil
-		}
-		return sqlfmt.SqlRowAsUpdateStmt(row, tableName, sch, updatedCols)
-	case ModifiedOld:
-		// do nothing, we only issue UPDATE for ModifiedNew
-		return "", nil
-	default:
-		return "", fmt.Errorf("unexpected row diff type: %v", rowDiffType)
-	}
-}
-
 // WorkingSetContainsOnlyIgnoredTables returns true if all changes in working set are ignored tables.
 // Otherwise, if there are any non-ignored changes, returns false.
 // Note that only unstaged tables are subject to dolt_ignore (this is consistent with what git does.)

diff --git a/go/libraries/doltcore/sqle/dolt_patch_table_function.go b/go/libraries/doltcore/sqle/dolt_patch_table_function.go
@@ -28,9 +28,7 @@ import (
 	"github.com/dolthub/vitess/go/mysql"
 	"golang.org/x/exp/slices"
 
-	"github.com/dolthub/dolt/go/cmd/dolt/errhand"
 	"github.com/dolthub/dolt/go/libraries/doltcore/diff"
-	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
@@ -514,7 +512,7 @@ func getPatchNodes(ctx *sql.Context, dbData env.DbData, tableDeltas []diff.Table
 		// Get SCHEMA DIFF
 		var schemaStmts []string
 		if includeSchemaDiff {
-			schemaStmts, err = GenerateSqlPatchSchemaStatements(ctx, toRefDetails.root, td)
+			schemaStmts, err = sqlfmt.GenerateSqlPatchSchemaStatements(ctx, toRefDetails.root, td)
 			if err != nil {
 				return nil, err
 			}
@@ -595,14 +593,14 @@ func getDataSqlPatchResults(ctx *sql.Context, diffQuerySch, targetSch sql.Schema
 
 		var stmt string
 		if oldRow.Row != nil {
-			stmt, err = diff.GetDataDiffStatement(tn, tsch, oldRow.Row, oldRow.RowDiff, oldRow.ColDiffs)
+			stmt, err = sqlfmt.GenerateDataDiffStatement(tn, tsch, oldRow.Row, oldRow.RowDiff, oldRow.ColDiffs)
 			if err != nil {
 				return nil, err
 			}
 		}
 
 		if newRow.Row != nil {
-			stmt, err = diff.GetDataDiffStatement(tn, tsch, newRow.Row, newRow.RowDiff, newRow.ColDiffs)
+			stmt, err = sqlfmt.GenerateDataDiffStatement(tn, tsch, newRow.Row, newRow.RowDiff, newRow.ColDiffs)
 			if err != nil {
 				return nil, err
 			}
@@ -614,131 +612,6 @@ func getDataSqlPatchResults(ctx *sql.Context, diffQuerySch, targetSch sql.Schema
 	}
 }
 
-// GenerateSqlPatchSchemaStatements examines the table schema changes in the specified TableDelta |td| and returns
-// a slice of SQL path statements that represent the equivalent SQL DDL statements for those schema changes. The
-// specified RootValue, |toRoot|, must be the RootValue that was used as the "To" root when computing the specified
-// TableDelta.
-func GenerateSqlPatchSchemaStatements(ctx *sql.Context, toRoot doltdb.RootValue, td diff.TableDelta) ([]string, error) {
-	toSchemas, err := doltdb.GetAllSchemas(ctx, toRoot)
-	if err != nil {
-		return nil, fmt.Errorf("could not read schemas from toRoot, cause: %s", err.Error())
-	}
-
-	fromSch, toSch, err := td.GetSchemas(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("cannot retrieve schema for table %s, cause: %s", td.ToName, err.Error())
-	}
-
-	var ddlStatements []string
-	if td.IsDrop() {
-		ddlStatements = append(ddlStatements, sqlfmt.DropTableStmt(td.FromName))
-	} else if td.IsAdd() {
-		stmt, err := sqlfmt.GenerateCreateTableStatement(td.ToName, td.ToSch, td.ToFks, td.ToFksParentSch)
-		if err != nil {
-			return nil, errhand.VerboseErrorFromError(err)
-		}
-		ddlStatements = append(ddlStatements, stmt)
-	} else {
-		stmts, err := GetNonCreateNonDropTableSqlSchemaDiff(td, toSchemas, fromSch, toSch)
-		if err != nil {
-			return nil, err
-		}
-		ddlStatements = append(ddlStatements, stmts...)
-	}
-
-	return ddlStatements, nil
-}
-
-// GetNonCreateNonDropTableSqlSchemaDiff returns any schema diff in SQL statements that is NEITHER 'CREATE TABLE' NOR 'DROP TABLE' statements.
-func GetNonCreateNonDropTableSqlSchemaDiff(td diff.TableDelta, toSchemas map[string]schema.Schema, fromSch, toSch schema.Schema) ([]string, error) {
-	if td.IsAdd() || td.IsDrop() {
-		// use add and drop specific methods
-		return nil, nil
-	}
-
-	var ddlStatements []string
-	if td.FromName != td.ToName {
-		ddlStatements = append(ddlStatements, sqlfmt.RenameTableStmt(td.FromName, td.ToName))
-	}
-
-	eq := schema.SchemasAreEqual(fromSch, toSch)
-	if eq && !td.HasFKChanges() {
-		return ddlStatements, nil
-	}
-
-	colDiffs, unionTags := diff.DiffSchColumns(fromSch, toSch)
-	for _, tag := range unionTags {
-		cd := colDiffs[tag]
-		switch cd.DiffType {
-		case diff.SchDiffNone:
-		case diff.SchDiffAdded:
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableAddColStmt(td.ToName, sqlfmt.GenerateCreateTableColumnDefinition(*cd.New, sql.CollationID(td.ToSch.GetCollation()))))
-		case diff.SchDiffRemoved:
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableDropColStmt(td.ToName, cd.Old.Name))
-		case diff.SchDiffModified:
-			// Ignore any primary key set changes here
-			if cd.Old.IsPartOfPK != cd.New.IsPartOfPK {
-				continue
-			}
-			if cd.Old.Name != cd.New.Name {
-				ddlStatements = append(ddlStatements, sqlfmt.AlterTableRenameColStmt(td.ToName, cd.Old.Name, cd.New.Name))
-			}
-			if !cd.Old.TypeInfo.Equals(cd.New.TypeInfo) {
-				ddlStatements = append(ddlStatements, sqlfmt.AlterTableModifyColStmt(td.ToName,
-					sqlfmt.GenerateCreateTableColumnDefinition(*cd.New, sql.CollationID(td.ToSch.GetCollation()))))
-			}
-		}
-	}
-
-	// Print changes between a primary key set change. It contains an ALTER TABLE DROP and an ALTER TABLE ADD
-	if !schema.ColCollsAreEqual(fromSch.GetPKCols(), toSch.GetPKCols()) {
-		ddlStatements = append(ddlStatements, sqlfmt.AlterTableDropPks(td.ToName))
-		if toSch.GetPKCols().Size() > 0 {
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableAddPrimaryKeys(td.ToName, toSch.GetPKCols().GetColumnNames()))
-		}
-	}
-
-	for _, idxDiff := range diff.DiffSchIndexes(fromSch, toSch) {
-		switch idxDiff.DiffType {
-		case diff.SchDiffNone:
-		case diff.SchDiffAdded:
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableAddIndexStmt(td.ToName, idxDiff.To))
-		case diff.SchDiffRemoved:
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableDropIndexStmt(td.FromName, idxDiff.From))
-		case diff.SchDiffModified:
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableDropIndexStmt(td.FromName, idxDiff.From))
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableAddIndexStmt(td.ToName, idxDiff.To))
-		}
-	}
-
-	for _, fkDiff := range diff.DiffForeignKeys(td.FromFks, td.ToFks) {
-		switch fkDiff.DiffType {
-		case diff.SchDiffNone:
-		case diff.SchDiffAdded:
-			parentSch := toSchemas[fkDiff.To.ReferencedTableName]
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableAddForeignKeyStmt(fkDiff.To, toSch, parentSch))
-		case diff.SchDiffRemoved:
-			from := fkDiff.From
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableDropForeignKeyStmt(from.TableName, from.Name))
-		case diff.SchDiffModified:
-			from := fkDiff.From
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableDropForeignKeyStmt(from.TableName, from.Name))
-
-			parentSch := toSchemas[fkDiff.To.ReferencedTableName]
-			ddlStatements = append(ddlStatements, sqlfmt.AlterTableAddForeignKeyStmt(fkDiff.To, toSch, parentSch))
-		}
-	}
-
-	// Handle charset/collation changes
-	toCollation := toSch.GetCollation()
-	fromCollation := fromSch.GetCollation()
-	if toCollation != fromCollation {
-		ddlStatements = append(ddlStatements, sqlfmt.AlterTableCollateStmt(td.ToName, fromCollation, toCollation))
-	}
-
-	return ddlStatements, nil
-}
-
 // getDiffQuery returns diff schema for specified columns and array of sql.Expression as projection to be used
 // on diff table function row iter. This function attempts to imitate running a query
 // fmt.Sprintf("select %s, %s from dolt_diff('%s', '%s', '%s')", columnsWithDiff, "diff_type", fromRef, toRef, tableName)

diff --git a/go/libraries/doltcore/sqle/sqlfmt/schema_fmt.go b/go/libraries/doltcore/sqle/sqlfmt/schema_fmt.go
@@ -20,10 +20,168 @@ import (
 
 	"github.com/dolthub/go-mysql-server/sql"
 
+	"github.com/dolthub/dolt/go/cmd/dolt/errhand"
+	"github.com/dolthub/dolt/go/libraries/doltcore/diff"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/utils/set"
 )
 
+// GenerateDataDiffStatement returns any data diff in SQL statements for given table including INSERT, UPDATE and DELETE row statements.
+func GenerateDataDiffStatement(tableName string, sch schema.Schema, row sql.Row, rowDiffType diff.ChangeType, colDiffTypes []diff.ChangeType) (string, error) {
+	if len(row) != len(colDiffTypes) {
+		return "", fmt.Errorf("expected the same size for columns and diff types, got %d and %d", len(row), len(colDiffTypes))
+	}
+
+	switch rowDiffType {
+	case diff.Added:
+		return SqlRowAsInsertStmt(row, tableName, sch)
+	case diff.Removed:
+		return SqlRowAsDeleteStmt(row, tableName, sch, 0)
+	case diff.ModifiedNew:
+		updatedCols := set.NewEmptyStrSet()
+		for i, diffType := range colDiffTypes {
+			if diffType != diff.None {
+				updatedCols.Add(sch.GetAllCols().GetByIndex(i).Name)
+			}
+		}
+		if updatedCols.Size() == 0 {
+			return "", nil
+		}
+		return SqlRowAsUpdateStmt(row, tableName, sch, updatedCols)
+	case diff.ModifiedOld:
+		// do nothing, we only issue UPDATE for ModifiedNew
+		return "", nil
+	default:
+		return "", fmt.Errorf("unexpected row diff type: %v", rowDiffType)
+	}
+}
+
+// GenerateSqlPatchSchemaStatements examines the table schema changes in the specified TableDelta |td| and returns
+// a slice of SQL path statements that represent the equivalent SQL DDL statements for those schema changes. The
+// specified RootValue, |toRoot|, must be the RootValue that was used as the "To" root when computing the specified
+// TableDelta.
+func GenerateSqlPatchSchemaStatements(ctx *sql.Context, toRoot doltdb.RootValue, td diff.TableDelta) ([]string, error) {
+	toSchemas, err := doltdb.GetAllSchemas(ctx, toRoot)
+	if err != nil {
+		return nil, fmt.Errorf("could not read schemas from toRoot, cause: %s", err.Error())
+	}
+
+	fromSch, toSch, err := td.GetSchemas(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("cannot retrieve schema for table %s, cause: %s", td.ToName, err.Error())
+	}
+
+	var ddlStatements []string
+	if td.IsDrop() {
+		ddlStatements = append(ddlStatements, DropTableStmt(td.FromName))
+	} else if td.IsAdd() {
+		stmt, err := GenerateCreateTableStatement(td.ToName, td.ToSch, td.ToFks, td.ToFksParentSch)
+		if err != nil {
+			return nil, errhand.VerboseErrorFromError(err)
+		}
+		ddlStatements = append(ddlStatements, stmt)
+	} else {
+		stmts, err := generateNonCreateNonDropTableSqlSchemaDiff(td, toSchemas, fromSch, toSch)
+		if err != nil {
+			return nil, err
+		}
+		ddlStatements = append(ddlStatements, stmts...)
+	}
+
+	return ddlStatements, nil
+}
+
+// generateNonCreateNonDropTableSqlSchemaDiff returns any schema diff in SQL statements that is NEITHER 'CREATE TABLE' NOR 'DROP TABLE' statements.
+func generateNonCreateNonDropTableSqlSchemaDiff(td diff.TableDelta, toSchemas map[string]schema.Schema, fromSch, toSch schema.Schema) ([]string, error) {
+	if td.IsAdd() || td.IsDrop() {
+		// use add and drop specific methods
+		return nil, nil
+	}
+
+	var ddlStatements []string
+	if td.FromName != td.ToName {
+		ddlStatements = append(ddlStatements, RenameTableStmt(td.FromName, td.ToName))
+	}
+
+	eq := schema.SchemasAreEqual(fromSch, toSch)
+	if eq && !td.HasFKChanges() {
+		return ddlStatements, nil
+	}
+
+	colDiffs, unionTags := diff.DiffSchColumns(fromSch, toSch)
+	for _, tag := range unionTags {
+		cd := colDiffs[tag]
+		switch cd.DiffType {
+		case diff.SchDiffNone:
+		case diff.SchDiffAdded:
+			ddlStatements = append(ddlStatements, AlterTableAddColStmt(td.ToName, GenerateCreateTableColumnDefinition(*cd.New, sql.CollationID(td.ToSch.GetCollation()))))
+		case diff.SchDiffRemoved:
+			ddlStatements = append(ddlStatements, AlterTableDropColStmt(td.ToName, cd.Old.Name))
+		case diff.SchDiffModified:
+			// Ignore any primary key set changes here
+			if cd.Old.IsPartOfPK != cd.New.IsPartOfPK {
+				continue
+			}
+			if cd.Old.Name != cd.New.Name {
+				ddlStatements = append(ddlStatements, AlterTableRenameColStmt(td.ToName, cd.Old.Name, cd.New.Name))
+			}
+			if !cd.Old.TypeInfo.Equals(cd.New.TypeInfo) {
+				ddlStatements = append(ddlStatements, AlterTableModifyColStmt(td.ToName,
+					GenerateCreateTableColumnDefinition(*cd.New, sql.CollationID(td.ToSch.GetCollation()))))
+			}
+		}
+	}
+
+	// Print changes between a primary key set change. It contains an ALTER TABLE DROP and an ALTER TABLE ADD
+	if !schema.ColCollsAreEqual(fromSch.GetPKCols(), toSch.GetPKCols()) {
+		ddlStatements = append(ddlStatements, AlterTableDropPks(td.ToName))
+		if toSch.GetPKCols().Size() > 0 {
+			ddlStatements = append(ddlStatements, AlterTableAddPrimaryKeys(td.ToName, toSch.GetPKCols().GetColumnNames()))
+		}
+	}
+
+	for _, idxDiff := range diff.DiffSchIndexes(fromSch, toSch) {
+		switch idxDiff.DiffType {
+		case diff.SchDiffNone:
+		case diff.SchDiffAdded:
+			ddlStatements = append(ddlStatements, AlterTableAddIndexStmt(td.ToName, idxDiff.To))
+		case diff.SchDiffRemoved:
+			ddlStatements = append(ddlStatements, AlterTableDropIndexStmt(td.FromName, idxDiff.From))
+		case diff.SchDiffModified:
+			ddlStatements = append(ddlStatements, AlterTableDropIndexStmt(td.FromName, idxDiff.From))
+			ddlStatements = append(ddlStatements, AlterTableAddIndexStmt(td.ToName, idxDiff.To))
+		}
+	}
+
+	for _, fkDiff := range diff.DiffForeignKeys(td.FromFks, td.ToFks) {
+		switch fkDiff.DiffType {
+		case diff.SchDiffNone:
+		case diff.SchDiffAdded:
+			parentSch := toSchemas[fkDiff.To.ReferencedTableName]
+			ddlStatements = append(ddlStatements, AlterTableAddForeignKeyStmt(fkDiff.To, toSch, parentSch))
+		case diff.SchDiffRemoved:
+			from := fkDiff.From
+			ddlStatements = append(ddlStatements, AlterTableDropForeignKeyStmt(from.TableName, from.Name))
+		case diff.SchDiffModified:
+			from := fkDiff.From
+			ddlStatements = append(ddlStatements, AlterTableDropForeignKeyStmt(from.TableName, from.Name))
+
+			parentSch := toSchemas[fkDiff.To.ReferencedTableName]
+			ddlStatements = append(ddlStatements, AlterTableAddForeignKeyStmt(fkDiff.To, toSch, parentSch))
+		}
+	}
+
+	// Handle charset/collation changes
+	toCollation := toSch.GetCollation()
+	fromCollation := fromSch.GetCollation()
+	if toCollation != fromCollation {
+		ddlStatements = append(ddlStatements, AlterTableCollateStmt(td.ToName, fromCollation, toCollation))
+	}
+
+	return ddlStatements, nil
+}
+
 // GenerateCreateTableColumnDefinition returns column definition for CREATE TABLE statement with no indentation
 func GenerateCreateTableColumnDefinition(col schema.Column, tableCollation sql.CollationID) string {
 	colStr := GenerateCreateTableIndentedColumnDefinition(col, tableCollation)

diff --git a/go/libraries/doltcore/table/untyped/sqlexport/sql_diff_writer.go b/go/libraries/doltcore/table/untyped/sqlexport/sql_diff_writer.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/diff"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/sqlfmt"
 	"github.com/dolthub/dolt/go/libraries/doltcore/table/editor"
 	"github.com/dolthub/dolt/go/libraries/utils/iohelp"
 )
@@ -49,7 +50,7 @@ func NewSqlDiffWriter(tableName string, schema schema.Schema, wr io.WriteCloser)
 }
 
 func (w SqlDiffWriter) WriteRow(ctx context.Context, row sql.Row, rowDiffType diff.ChangeType, colDiffTypes []diff.ChangeType) error {
-	stmt, err := diff.GetDataDiffStatement(w.tableName, w.sch, row, rowDiffType, colDiffTypes)
+	stmt, err := sqlfmt.GenerateDataDiffStatement(w.tableName, w.sch, row, rowDiffType, colDiffTypes)
 	if err != nil {
 		return err
 	}