support postgres comment (#48)

support postgres comment
sunary · Jun 11, 2024 · 18a3a3d · 18a3a3d
1 parent 42d8e46
commit 18a3a3d
Show file tree

Hide file tree

Showing 12 changed files with 163 additions and 88 deletions.
diff --git a/README.md b/README.md
@@ -116,6 +116,7 @@ func main() {
 * Support **generate** comment, using option `sql_builder.WithCommentGenerate()`
 * Support automatic addition of `s` to table names (plural naming convention), using option `sql_builder.WithPluralTableName()`
 * Accept tag convention: `snake_case` or `camelCase`, Eg: `sql:"primary_key"` equalize `sql:"primaryKey"`
+* Custom column name: `sql:"column:column_name"`
 * Primary key for this field: `sql:"primary_key"`
 * Foreign key: `sql:"foreign_key:user_id;references:user_id"`
 * Auto increment: `sql:"auto_increment"`

diff --git a/element/column.go b/element/column.go
@@ -12,7 +12,7 @@ import (
 	"github.com/pingcap/parser/format"
 	"github.com/pingcap/parser/types"
 	sqlite "github.com/rqlite/sql"
-	"github.com/sunary/sqlize/utils"
+	sql_templates "github.com/sunary/sqlize/sql-templates"
 )
 
 const (
@@ -53,7 +53,7 @@ func (c Column) HasDefaultValue() bool {
 }
 
 func (c Column) hashValue() string {
-	strHash := utils.EscapeSqlName(sql.GetDialect(), c.Name)
+	strHash := sql.EscapeSqlName(c.Name)
 	strHash += c.typeDefinition()
 	hash := md5.Sum([]byte(strHash))
 	return hex.EncodeToString(hash[:])
@@ -65,7 +65,7 @@ func (c Column) migrationUp(tbName, after string, ident int) []string {
 		return nil
 
 	case MigrateAddAction:
-		strSql := utils.EscapeSqlName(sql.GetDialect(), c.Name)
+		strSql := sql.EscapeSqlName(c.Name)
 
 		if ident > len(c.Name) {
 			strSql += strings.Repeat(" ", ident-len(c.Name))
@@ -75,31 +75,40 @@ func (c Column) migrationUp(tbName, after string, ident int) []string {
 
 		if ident < 0 {
 			if after != "" {
-				return []string{fmt.Sprintf(sql.AlterTableAddColumnAfterStm(), utils.EscapeSqlName(sql.GetDialect(), tbName), strSql, utils.EscapeSqlName(sql.GetDialect(), after))}
+				return []string{fmt.Sprintf(sql.AlterTableAddColumnAfterStm(), sql.EscapeSqlName(tbName), strSql, sql.EscapeSqlName(after))}
 			}
-			return []string{fmt.Sprintf(sql.AlterTableAddColumnFirstStm(), utils.EscapeSqlName(sql.GetDialect(), tbName), strSql)}
+			return []string{fmt.Sprintf(sql.AlterTableAddColumnFirstStm(), sql.EscapeSqlName(tbName), strSql)}
 		}
 
-		return []string{strSql}
+		return append([]string{strSql}, c.migrationCommentUp(tbName)...)
 
 	case MigrateRemoveAction:
 		if sql.IsSqlite() {
 			return nil
 		}
 
-		return []string{fmt.Sprintf(sql.AlterTableDropColumnStm(), utils.EscapeSqlName(sql.GetDialect(), tbName), utils.EscapeSqlName(sql.GetDialect(), c.Name))}
+		return []string{fmt.Sprintf(sql.AlterTableDropColumnStm(), sql.EscapeSqlName(tbName), sql.EscapeSqlName(c.Name))}
 
 	case MigrateModifyAction:
-		return []string{fmt.Sprintf(sql.AlterTableModifyColumnStm(), utils.EscapeSqlName(sql.GetDialect(), tbName), utils.EscapeSqlName(sql.GetDialect(), c.Name)+c.definition())}
+		return []string{fmt.Sprintf(sql.AlterTableModifyColumnStm(), sql.EscapeSqlName(tbName), sql.EscapeSqlName(c.Name)+c.definition())}
 
 	case MigrateRenameAction:
-		return []string{fmt.Sprintf(sql.AlterTableRenameColumnStm(), utils.EscapeSqlName(sql.GetDialect(), tbName), utils.EscapeSqlName(sql.GetDialect(), c.OldName), utils.EscapeSqlName(sql.GetDialect(), c.Name))}
+		return []string{fmt.Sprintf(sql.AlterTableRenameColumnStm(), sql.EscapeSqlName(tbName), sql.EscapeSqlName(c.OldName), sql.EscapeSqlName(c.Name))}
 
 	default:
 		return nil
 	}
 }
 
+func (c Column) migrationCommentUp(tbName string) []string {
+	if c.Comment == "" || sql.GetDialect() != sql_templates.PostgresDialect {
+		return nil
+	}
+
+	// apply for postgres only
+	return []string{fmt.Sprintf(sql.ColumnComment(), tbName, c.Name, c.Comment)}
+}
+
 func (c Column) migrationDown(tbName, after string) []string {
 	switch c.Action {
 	case MigrateNoAction:

diff --git a/element/foreign_key.go b/element/foreign_key.go
@@ -5,8 +5,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"strings"
-
-	"github.com/sunary/sqlize/utils"
 )
 
 // ForeignKey ...
@@ -32,12 +30,12 @@ func (fk ForeignKey) migrationUp(tbName string) []string {
 
 	case MigrateAddAction:
 		return []string{fmt.Sprintf(sql.CreateForeignKeyStm(),
-			utils.EscapeSqlName(sql.GetDialect(), tbName), utils.EscapeSqlName(sql.GetDialect(), fk.Name), utils.EscapeSqlName(sql.GetDialect(), fk.Column),
-			utils.EscapeSqlName(sql.GetDialect(), fk.RefTable), utils.EscapeSqlName(sql.GetDialect(), fk.RefColumn))}
+			sql.EscapeSqlName(tbName), sql.EscapeSqlName(fk.Name), sql.EscapeSqlName(fk.Column),
+			sql.EscapeSqlName(fk.RefTable), sql.EscapeSqlName(fk.RefColumn))}
 
 	case MigrateRemoveAction:
 		return []string{fmt.Sprintf(sql.DropForeignKeyStm(),
-			utils.EscapeSqlName(sql.GetDialect(), tbName), utils.EscapeSqlName(sql.GetDialect(), fk.Name))}
+			sql.EscapeSqlName(tbName), sql.EscapeSqlName(fk.Name))}
 
 	case MigrateModifyAction:
 		return nil

diff --git a/element/index.go b/element/index.go
@@ -8,7 +8,6 @@ import (
 
 	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/model"
-	"github.com/sunary/sqlize/utils"
 )
 
 // Index ...
@@ -34,20 +33,20 @@ func (i Index) migrationUp(tbName string) []string {
 	case MigrateAddAction:
 		if i.CnsTyp == ast.ConstraintPrimaryKey {
 			return []string{fmt.Sprintf(sql.CreatePrimaryKeyStm(),
-				utils.EscapeSqlName(sql.GetDialect(), tbName),
-				strings.Join(utils.EscapeSqlNames(sql.GetDialect(), i.Columns), ", "))}
+				sql.EscapeSqlName(tbName),
+				strings.Join(sql.EscapeSqlNames(i.Columns), ", "))}
 		}
 
 		switch i.Typ {
 		case ast.IndexKeyTypeNone:
 			return []string{fmt.Sprintf(sql.CreateIndexStm(i.IndexType.String()),
-				utils.EscapeSqlName(sql.GetDialect(), i.Name), utils.EscapeSqlName(sql.GetDialect(), tbName),
-				strings.Join(utils.EscapeSqlNames(sql.GetDialect(), i.Columns), ", "))}
+				sql.EscapeSqlName(i.Name), sql.EscapeSqlName(tbName),
+				strings.Join(sql.EscapeSqlNames(i.Columns), ", "))}
 
 		case ast.IndexKeyTypeUnique:
 			return []string{fmt.Sprintf(sql.CreateUniqueIndexStm(i.IndexType.String()),
-				utils.EscapeSqlName(sql.GetDialect(), i.Name), utils.EscapeSqlName(sql.GetDialect(), tbName),
-				strings.Join(utils.EscapeSqlNames(sql.GetDialect(), i.Columns), ", "))}
+				sql.EscapeSqlName(i.Name), sql.EscapeSqlName(tbName),
+				strings.Join(sql.EscapeSqlNames(i.Columns), ", "))}
 
 		default:
 			return nil
@@ -56,17 +55,17 @@ func (i Index) migrationUp(tbName string) []string {
 	case MigrateRemoveAction:
 		if i.CnsTyp == ast.ConstraintPrimaryKey {
 			return []string{fmt.Sprintf(sql.DropPrimaryKeyStm(),
-				utils.EscapeSqlName(sql.GetDialect(), tbName))}
+				sql.EscapeSqlName(tbName))}
 		}
 
 		if sql.IsSqlite() {
 			return []string{fmt.Sprintf(sql.DropIndexStm(),
-				utils.EscapeSqlName(sql.GetDialect(), i.Name))}
+				sql.EscapeSqlName(i.Name))}
 		}
 
 		return []string{fmt.Sprintf(sql.DropIndexStm(),
-			utils.EscapeSqlName(sql.GetDialect(), i.Name),
-			utils.EscapeSqlName(sql.GetDialect(), tbName))}
+			sql.EscapeSqlName(i.Name),
+			sql.EscapeSqlName(tbName))}
 
 	case MigrateModifyAction:
 		strRems := make([]string, 2)
@@ -78,9 +77,9 @@ func (i Index) migrationUp(tbName string) []string {
 
 	case MigrateRenameAction:
 		return []string{fmt.Sprintf(sql.AlterTableRenameIndexStm(),
-			utils.EscapeSqlName(sql.GetDialect(), tbName),
-			utils.EscapeSqlName(sql.GetDialect(), i.OldName),
-			utils.EscapeSqlName(sql.GetDialect(), i.Name))}
+			sql.EscapeSqlName(tbName),
+			sql.EscapeSqlName(i.OldName),
+			sql.EscapeSqlName(i.Name))}
 
 	default:
 		return nil

diff --git a/element/migration.go b/element/migration.go
@@ -131,6 +131,25 @@ func (m *Migration) RenameColumn(tbName, oldName, newName string) {
 	}
 }
 
+// AddIndex ...
+func (m *Migration) AddComment(tbName, colName, comment string) {
+	if tbName == "" {
+		tbName = m.currentTable
+	}
+
+	id := m.getIndexTable(tbName)
+	if id == -1 {
+		return
+	}
+
+	colIdx := m.Tables[id].getIndexColumn(colName)
+	if colIdx == -1 {
+		return
+	}
+
+	m.Tables[id].Columns[colIdx].Comment = comment
+}
+
 // AddIndex ...
 func (m *Migration) AddIndex(tbName string, idx Index) {
 	if tbName == "" {

diff --git a/element/table.go b/element/table.go
@@ -404,6 +404,7 @@ func (t Table) MigrationColumnUp() ([]string, map[string]struct{}) {
 		}
 
 		strCols := make([]string, 0)
+		commentCols := make([]string, 0)
 		for i := range t.Columns {
 			if t.Columns[i].Action == MigrateAddAction {
 				strCols = append(strCols, " "+t.Columns[i].migrationUp("", "", maxIdent)[0])
@@ -412,12 +413,14 @@ func (t Table) MigrationColumnUp() ([]string, map[string]struct{}) {
 				nCol.Action = MigrateAddAction
 				strCols = append(strCols, " "+nCol.migrationUp("", "", maxIdent)[0])
 			}
+
+			commentCols = append(commentCols, t.Columns[i].migrationCommentUp(t.Name)...)
 		}
 
-		return []string{fmt.Sprintf(sql.CreateTableStm(), utils.EscapeSqlName(sql.GetDialect(), t.Name), strings.Join(strCols, ",\n"), "")}, nil
+		return append([]string{fmt.Sprintf(sql.CreateTableStm(), sql.EscapeSqlName(t.Name), strings.Join(strCols, ",\n"), "")}, commentCols...), nil
 
 	case MigrateRemoveAction:
-		return []string{fmt.Sprintf(sql.DropTableStm(), utils.EscapeSqlName(sql.GetDialect(), t.Name))}, nil
+		return []string{fmt.Sprintf(sql.DropTableStm(), sql.EscapeSqlName(t.Name))}, nil
 
 	case MigrateModifyAction:
 		// TODO