Fix postgres

This commit is contained in:
Lunny Xiao 2020-03-25 17:04:43 +08:00
parent 32b224abe5
commit 6a71ed3696
No known key found for this signature in database
GPG Key ID: C3B7C91B632F738A
2 changed files with 23 additions and 19 deletions

View File

@ -32,8 +32,9 @@ type URI struct {
// SetSchema set schema // SetSchema set schema
func (uri *URI) SetSchema(schema string) { func (uri *URI) SetSchema(schema string) {
// hack me
if uri.DBType == schemas.POSTGRES { if uri.DBType == schemas.POSTGRES {
uri.Schema = schema uri.Schema = strings.TrimSpace(schema)
} }
} }

View File

@ -781,11 +781,14 @@ type postgres struct {
func (db *postgres) Init(uri *URI) error { func (db *postgres) Init(uri *URI) error {
db.quoter = postgresQuoter db.quoter = postgresQuoter
if uri.Schema == "" { return db.Base.Init(db, uri)
uri.Schema = DefaultPostgresSchema
} }
return db.Base.Init(db, uri) func (db *postgres) getSchema() string {
if db.uri.Schema != "" {
return db.uri.Schema
}
return DefaultPostgresSchema
} }
func (db *postgres) needQuote(name string) bool { func (db *postgres) needQuote(name string) bool {
@ -928,32 +931,32 @@ func (db *postgres) CreateTableSQL(table *schemas.Table, tableName string) ([]st
} }
func (db *postgres) IndexCheckSQL(tableName, idxName string) (string, []interface{}) { func (db *postgres) IndexCheckSQL(tableName, idxName string) (string, []interface{}) {
if len(db.uri.Schema) == 0 { if len(db.getSchema()) == 0 {
args := []interface{}{tableName, idxName} args := []interface{}{tableName, idxName}
return `SELECT indexname FROM pg_indexes WHERE tablename = ? AND indexname = ?`, args return `SELECT indexname FROM pg_indexes WHERE tablename = ? AND indexname = ?`, args
} }
args := []interface{}{db.uri.Schema, tableName, idxName} args := []interface{}{db.getSchema(), tableName, idxName}
return `SELECT indexname FROM pg_indexes ` + return `SELECT indexname FROM pg_indexes ` +
`WHERE schemaname = ? AND tablename = ? AND indexname = ?`, args `WHERE schemaname = ? AND tablename = ? AND indexname = ?`, args
} }
func (db *postgres) IsTableExist(queryer core.Queryer, ctx context.Context, tableName string) (bool, error) { func (db *postgres) IsTableExist(queryer core.Queryer, ctx context.Context, tableName string) (bool, error) {
if len(db.uri.Schema) == 0 { if len(db.getSchema()) == 0 {
return db.HasRecords(queryer, ctx, `SELECT tablename FROM pg_tables WHERE tablename = $1`, tableName) return db.HasRecords(queryer, ctx, `SELECT tablename FROM pg_tables WHERE tablename = $1`, tableName)
} }
return db.HasRecords(queryer, ctx, `SELECT tablename FROM pg_tables WHERE schemaname = $1 AND tablename = $2`, return db.HasRecords(queryer, ctx, `SELECT tablename FROM pg_tables WHERE schemaname = $1 AND tablename = $2`,
db.uri.Schema, tableName) db.getSchema(), tableName)
} }
func (db *postgres) ModifyColumnSQL(tableName string, col *schemas.Column) string { func (db *postgres) ModifyColumnSQL(tableName string, col *schemas.Column) string {
if len(db.uri.Schema) == 0 || strings.Contains(tableName, ".") { if len(db.getSchema()) == 0 || strings.Contains(tableName, ".") {
return fmt.Sprintf("alter table %s ALTER COLUMN %s TYPE %s", return fmt.Sprintf("alter table %s ALTER COLUMN %s TYPE %s",
tableName, col.Name, db.SQLType(col)) tableName, col.Name, db.SQLType(col))
} }
return fmt.Sprintf("alter table %s.%s ALTER COLUMN %s TYPE %s", return fmt.Sprintf("alter table %s.%s ALTER COLUMN %s TYPE %s",
db.uri.Schema, tableName, col.Name, db.SQLType(col)) db.getSchema(), tableName, col.Name, db.SQLType(col))
} }
func (db *postgres) DropIndexSQL(tableName string, index *schemas.Index) string { func (db *postgres) DropIndexSQL(tableName string, index *schemas.Index) string {
@ -970,17 +973,17 @@ func (db *postgres) DropIndexSQL(tableName string, index *schemas.Index) string
idxName = fmt.Sprintf("IDX_%v_%v", tableName, index.Name) idxName = fmt.Sprintf("IDX_%v_%v", tableName, index.Name)
} }
} }
if db.uri.Schema != "" { if db.getSchema() != "" {
idxName = db.uri.Schema + "." + idxName idxName = db.getSchema() + "." + idxName
} }
return fmt.Sprintf("DROP INDEX %v", db.Quoter().Quote(idxName)) return fmt.Sprintf("DROP INDEX %v", db.Quoter().Quote(idxName))
} }
func (db *postgres) IsColumnExist(queryer core.Queryer, ctx context.Context, tableName, colName string) (bool, error) { func (db *postgres) IsColumnExist(queryer core.Queryer, ctx context.Context, tableName, colName string) (bool, error) {
args := []interface{}{db.uri.Schema, tableName, colName} args := []interface{}{db.getSchema(), tableName, colName}
query := "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE table_schema = $1 AND table_name = $2" + query := "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE table_schema = $1 AND table_name = $2" +
" AND column_name = $3" " AND column_name = $3"
if len(db.uri.Schema) == 0 { if len(db.getSchema()) == 0 {
args = []interface{}{tableName, colName} args = []interface{}{tableName, colName}
query = "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = $1" + query = "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = $1" +
" AND column_name = $2" " AND column_name = $2"
@ -1009,7 +1012,7 @@ FROM pg_attribute f
LEFT JOIN INFORMATION_SCHEMA.COLUMNS s ON s.column_name=f.attname AND c.relname=s.table_name LEFT JOIN INFORMATION_SCHEMA.COLUMNS s ON s.column_name=f.attname AND c.relname=s.table_name
WHERE n.nspname= s.table_schema AND c.relkind = 'r'::char AND c.relname = $1%s AND f.attnum > 0 ORDER BY f.attnum;` WHERE n.nspname= s.table_schema AND c.relkind = 'r'::char AND c.relname = $1%s AND f.attnum > 0 ORDER BY f.attnum;`
schema := db.uri.Schema schema := db.getSchema()
if schema != "" { if schema != "" {
s = fmt.Sprintf(s, "AND s.table_schema = $2") s = fmt.Sprintf(s, "AND s.table_schema = $2")
args = append(args, schema) args = append(args, schema)
@ -1139,7 +1142,7 @@ WHERE n.nspname= s.table_schema AND c.relkind = 'r'::char AND c.relname = $1%s A
func (db *postgres) GetTables(queryer core.Queryer, ctx context.Context) ([]*schemas.Table, error) { func (db *postgres) GetTables(queryer core.Queryer, ctx context.Context) ([]*schemas.Table, error) {
args := []interface{}{} args := []interface{}{}
s := "SELECT tablename FROM pg_tables" s := "SELECT tablename FROM pg_tables"
schema := db.uri.Schema schema := db.getSchema()
if schema != "" { if schema != "" {
args = append(args, schema) args = append(args, schema)
s = s + " WHERE schemaname = $1" s = s + " WHERE schemaname = $1"
@ -1179,8 +1182,8 @@ func getIndexColName(indexdef string) []string {
func (db *postgres) GetIndexes(queryer core.Queryer, ctx context.Context, tableName string) (map[string]*schemas.Index, error) { func (db *postgres) GetIndexes(queryer core.Queryer, ctx context.Context, tableName string) (map[string]*schemas.Index, error) {
args := []interface{}{tableName} args := []interface{}{tableName}
s := fmt.Sprintf("SELECT indexname, indexdef FROM pg_indexes WHERE tablename=$1") s := fmt.Sprintf("SELECT indexname, indexdef FROM pg_indexes WHERE tablename=$1")
if len(db.uri.Schema) != 0 { if len(db.getSchema()) != 0 {
args = append(args, db.uri.Schema) args = append(args, db.getSchema())
s = s + " AND schemaname=$2" s = s + " AND schemaname=$2"
} }