Add SetSchema for engine (go-xorm#876)

* add SetSchema for engine

* fix user

* fix postgres with schema

* fix test

* fix test

* fix test

* fix tablename

* refactor tableName

* fix schema support

* improve the interface of EngineInterface

Author: lunny

circle.yml

@@ -17,6 +17,7 @@ database:
     - createdb -p 5432 -e -U postgres xorm_test1
     - createdb -p 5432 -e -U postgres xorm_test2
     - createdb -p 5432 -e -U postgres xorm_test3
+    - psql xorm_test postgres -c "create schema xorm"
 
 test:
   override:
@@ -30,7 +31,9 @@ test:
     - go test -v -race -db="mymysql" -conn_str="xorm_test/root/" -cache=true -coverprofile=coverage3-2.txt -covermode=atomic
     - go test -v -race -db="postgres" -conn_str="dbname=xorm_test sslmode=disable" -coverprofile=coverage4-1.txt -covermode=atomic
     - go test -v -race -db="postgres" -conn_str="dbname=xorm_test sslmode=disable" -cache=true -coverprofile=coverage4-2.txt -covermode=atomic
-    - gocovmerge coverage1-1.txt coverage1-2.txt coverage2-1.txt coverage2-2.txt coverage3-1.txt coverage3-2.txt coverage4-1.txt coverage4-2.txt > coverage.txt
+    - go test -v -race -db="postgres" -conn_str="dbname=xorm_test sslmode=disable" -schema=xorm -coverprofile=coverage5-1.txt -covermode=atomic
+    - go test -v -race -db="postgres" -conn_str="dbname=xorm_test sslmode=disable" -schema=xorm -cache=true -coverprofile=coverage5-2.txt -covermode=atomic
+    - gocovmerge coverage1-1.txt coverage1-2.txt coverage2-1.txt coverage2-2.txt coverage3-1.txt coverage3-2.txt coverage4-1.txt coverage4-2.txt coverage5-1.txt coverage5-2.txt > coverage.txt
     - cd /home/ubuntu/.go_workspace/src/github.com/go-xorm/tests && ./sqlite3.sh
     - cd /home/ubuntu/.go_workspace/src/github.com/go-xorm/tests && ./mysql.sh
     - cd /home/ubuntu/.go_workspace/src/github.com/go-xorm/tests && ./postgres.sh

dialect_postgres.go

@@ -895,6 +895,7 @@ func (db *postgres) TableCheckSql(tableName string) (string, []interface{}) {
 		args := []interface{}{tableName}
 		return `SELECT tablename FROM pg_tables WHERE tablename = ?`, args
 	}
+
 	args := []interface{}{db.Schema, tableName}
 	return `SELECT tablename FROM pg_tables WHERE schemaname = ? AND tablename = ?`, args
 }
@@ -912,6 +913,9 @@ func (db *postgres) DropIndexSql(tableName string, index *core.Index) string {
 	quote := db.Quote
 	idxName := index.Name
 
+	tableName = strings.Replace(tableName, `"`, "", -1)
+	tableName = strings.Replace(tableName, `.`, "_", -1)
+
 	if !strings.HasPrefix(idxName, "UQE_") &&
 		!strings.HasPrefix(idxName, "IDX_") {
 		if index.Type == core.UniqueType {
@@ -920,6 +924,9 @@ func (db *postgres) DropIndexSql(tableName string, index *core.Index) string {
 			idxName = fmt.Sprintf("IDX_%v_%v", tableName, index.Name)
 		}
 	}
+	if db.Uri.Schema != "" {
+		idxName = db.Uri.Schema + "." + idxName
+	}
 	return fmt.Sprintf("DROP INDEX %v", quote(idxName))
 }
 
@@ -960,7 +967,7 @@ WHERE c.relkind = 'r'::char AND c.relname = $1%s AND f.attnum > 0 ORDER BY f.att
 	var f string
 	if len(db.Schema) != 0 {
 		args = append(args, db.Schema)
-		f = "AND s.table_schema = $2"
+		f = " AND s.table_schema = $2"
 	}
 	s = fmt.Sprintf(s, f)
 
@@ -1085,11 +1092,11 @@ func (db *postgres) GetTables() ([]*core.Table, error) {
 func (db *postgres) GetIndexes(tableName string) (map[string]*core.Index, error) {
 	args := []interface{}{tableName}
 	s := fmt.Sprintf("SELECT indexname, indexdef FROM pg_indexes WHERE tablename=$1")
-	db.LogSQL(s, args)
 	if len(db.Schema) != 0 {
 		args = append(args, db.Schema)
 		s = s + " AND schemaname=$2"
 	}
+	db.LogSQL(s, args)
 
 	rows, err := db.DB().Query(s, args...)
 	if err != nil {

engine.go

@@ -536,46 +536,6 @@ func (engine *Engine) dumpTables(tables []*core.Table, w io.Writer, tp ...core.D
 	return nil
 }
 
-func (engine *Engine) tableName(beanOrTableName interface{}) (string, error) {
-	v := rValue(beanOrTableName)
-	if v.Type().Kind() == reflect.String {
-		return beanOrTableName.(string), nil
-	} else if v.Type().Kind() == reflect.Struct {
-		return engine.tbName(v), nil
-	}
-	return "", errors.New("bean should be a struct or struct's point")
-}
-
-func (engine *Engine) tbSchemaName(v string) string {
-	// Add schema name as prefix of table name.
-	// Only for postgres database.
-	if engine.dialect.DBType() == core.POSTGRES &&
-		engine.dialect.URI().Schema != "" &&
-		engine.dialect.URI().Schema != postgresPublicSchema &&
-		strings.Index(v, ".") == -1 {
-		return engine.dialect.URI().Schema + "." + v
-	}
-	return v
-}
-
-func (engine *Engine) tbName(v reflect.Value) string {
-	if tb, ok := v.Interface().(TableName); ok {
-		return engine.tbSchemaName(tb.TableName())
-
-	}
-
-	if v.Type().Kind() == reflect.Ptr {
-		if tb, ok := reflect.Indirect(v).Interface().(TableName); ok {
-			return engine.tbSchemaName(tb.TableName())
-		}
-	} else if v.CanAddr() {
-		if tb, ok := v.Addr().Interface().(TableName); ok {
-			return engine.tbSchemaName(tb.TableName())
-		}
-	}
-	return engine.tbSchemaName(engine.TableMapper.Obj2Table(reflect.Indirect(v).Type().Name()))
-}
-
 // Cascade use cascade or not
 func (engine *Engine) Cascade(trueOrFalse ...bool) *Session {
 	session := engine.NewSession()
@@ -859,7 +819,7 @@ func (engine *Engine) TableInfo(bean interface{}) *Table {
 	if err != nil {
 		engine.logger.Error(err)
 	}
-	return &Table{tb, engine.tbName(v)}
+	return &Table{tb, engine.TableName(bean)}
 }
 
 func addIndex(indexName string, table *core.Table, col *core.Column, indexType int) {
@@ -895,20 +855,8 @@ var (
 func (engine *Engine) mapType(v reflect.Value) (*core.Table, error) {
 	t := v.Type()
 	table := engine.newTable()
-	if tb, ok := v.Interface().(TableName); ok {
-		table.Name = tb.TableName()
-	} else {
-		if v.CanAddr() {
-			if tb, ok = v.Addr().Interface().(TableName); ok {
-				table.Name = tb.TableName()
-			}
-		}
-		if table.Name == "" {
-			table.Name = engine.TableMapper.Obj2Table(t.Name())
-		}
-	}
-
 	table.Type = t
+	table.Name = engine.tbNameForMap(v)
 
 	var idFieldColName string
 	var hasCacheTag, hasNoCacheTag bool
@@ -1186,7 +1134,7 @@ func (engine *Engine) ClearCacheBean(bean interface{}, id string) error {
 	if t.Kind() != reflect.Struct {
 		return errors.New("error params")
 	}
-	tableName := engine.tbName(v)
+	tableName := engine.TableName(bean)
 	table, err := engine.autoMapType(v)
 	if err != nil {
 		return err
@@ -1210,7 +1158,7 @@ func (engine *Engine) ClearCache(beans ...interface{}) error {
 		if t.Kind() != reflect.Struct {
 			return errors.New("error params")
 		}
-		tableName := engine.tbName(v)
+		tableName := engine.TableName(bean)
 		table, err := engine.autoMapType(v)
 		if err != nil {
 			return err
@@ -1237,13 +1185,13 @@ func (engine *Engine) Sync(beans ...interface{}) error {
 
 	for _, bean := range beans {
 		v := rValue(bean)
-		tableName := engine.tbName(v)
+		tableNameNoSchema := engine.tbNameNoSchema(v.Interface())
 		table, err := engine.autoMapType(v)
 		if err != nil {
 			return err
 		}
 
-		isExist, err := session.Table(bean).isTableExist(tableName)
+		isExist, err := session.Table(bean).isTableExist(tableNameNoSchema)
 		if err != nil {
 			return err
 		}
@@ -1269,12 +1217,12 @@ func (engine *Engine) Sync(beans ...interface{}) error {
 			}
 		} else {
 			for _, col := range table.Columns() {
-				isExist, err := engine.dialect.IsColumnExist(tableName, col.Name)
+				isExist, err := engine.dialect.IsColumnExist(tableNameNoSchema, col.Name)
 				if err != nil {
 					return err
 				}
 				if !isExist {
-					if err := session.statement.setRefValue(v); err != nil {
+					if err := session.statement.setRefBean(bean); err != nil {
 						return err
 					}
 					err = session.addColumn(col.Name)
@@ -1285,35 +1233,35 @@ func (engine *Engine) Sync(beans ...interface{}) error {
 			}
 
 			for name, index := range table.Indexes {
-				if err := session.statement.setRefValue(v); err != nil {
+				if err := session.statement.setRefBean(bean); err != nil {
 					return err
 				}
 				if index.Type == core.UniqueType {
-					isExist, err := session.isIndexExist2(tableName, index.Cols, true)
+					isExist, err := session.isIndexExist2(tableNameNoSchema, index.Cols, true)
 					if err != nil {
 						return err
 					}
 					if !isExist {
-						if err := session.statement.setRefValue(v); err != nil {
+						if err := session.statement.setRefBean(bean); err != nil {
 							return err
 						}
 
-						err = session.addUnique(tableName, name)
+						err = session.addUnique(tableNameNoSchema, name)
 						if err != nil {
 							return err
 						}
 					}
 				} else if index.Type == core.IndexType {
-					isExist, err := session.isIndexExist2(tableName, index.Cols, false)
+					isExist, err := session.isIndexExist2(tableNameNoSchema, index.Cols, false)
 					if err != nil {
 						return err
 					}
 					if !isExist {
-						if err := session.statement.setRefValue(v); err != nil {
+						if err := session.statement.setRefBean(bean); err != nil {
 							return err
 						}
 
-						err = session.addIndex(tableName, name)
+						err = session.addIndex(tableNameNoSchema, name)
 						if err != nil {
 							return err
 						}
@@ -1649,6 +1597,11 @@ func (engine *Engine) SetTZDatabase(tz *time.Location) {
 	engine.DatabaseTZ = tz
 }
 
+// SetSchema sets the schema of database
+func (engine *Engine) SetSchema(schema string) {
+	engine.dialect.URI().Schema = schema
+}
+
 // Unscoped always disable struct tag "deleted"
 func (engine *Engine) Unscoped() *Session {
 	session := engine.NewSession()

engine_table.go

@@ -0,0 +1,109 @@
+// Copyright 2018 The Xorm Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package xorm
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+
+	"github.com/go-xorm/core"
+)
+
+// TableNameWithSchema will automatically add schema prefix on table name
+func (engine *Engine) tbNameWithSchema(v string) string {
+	// Add schema name as prefix of table name.
+	// Only for postgres database.
+	if engine.dialect.DBType() == core.POSTGRES &&
+		engine.dialect.URI().Schema != "" &&
+		engine.dialect.URI().Schema != postgresPublicSchema &&
+		strings.Index(v, ".") == -1 {
+		return engine.dialect.URI().Schema + "." + v
+	}
+	return v
+}
+
+// TableName returns table name with schema prefix if has
+func (engine *Engine) TableName(bean interface{}, includeSchema ...bool) string {
+	tbName := engine.tbNameNoSchema(bean)
+	if len(includeSchema) > 0 && includeSchema[0] {
+		tbName = engine.tbNameWithSchema(tbName)
+	}
+
+	return tbName
+}
+
+// tbName get some table's table name
+func (session *Session) tbNameNoSchema(table *core.Table) string {
+	if len(session.statement.AltTableName) > 0 {
+		return session.statement.AltTableName
+	}
+
+	return table.Name
+}
+
+func (engine *Engine) tbNameForMap(v reflect.Value) string {
+	t := v.Type()
+	if tb, ok := v.Interface().(TableName); ok {
+		return tb.TableName()
+	}
+	if v.CanAddr() {
+		if tb, ok := v.Addr().Interface().(TableName); ok {
+			return tb.TableName()
+		}
+	}
+	return engine.TableMapper.Obj2Table(t.Name())
+}
+
+func (engine *Engine) tbNameNoSchema(tablename interface{}) string {
+	switch tablename.(type) {
+	case []string:
+		t := tablename.([]string)
+		if len(t) > 1 {
+			return fmt.Sprintf("%v AS %v", engine.Quote(t[0]), engine.Quote(t[1]))
+		} else if len(t) == 1 {
+			return engine.Quote(t[0])
+		}
+	case []interface{}:
+		t := tablename.([]interface{})
+		l := len(t)
+		var table string
+		if l > 0 {
+			f := t[0]
+			switch f.(type) {
+			case string:
+				table = f.(string)
+			case TableName:
+				table = f.(TableName).TableName()
+			default:
+				v := rValue(f)
+				t := v.Type()
+				if t.Kind() == reflect.Struct {
+					table = engine.tbNameForMap(v)
+				} else {
+					table = engine.Quote(fmt.Sprintf("%v", f))
+				}
+			}
+		}
+		if l > 1 {
+			return fmt.Sprintf("%v AS %v", engine.Quote(table),
+				engine.Quote(fmt.Sprintf("%v", t[1])))
+		} else if l == 1 {
+			return engine.Quote(table)
+		}
+	case TableName:
+		return tablename.(TableName).TableName()
+	case string:
+		return tablename.(string)
+	default:
+		v := rValue(tablename)
+		t := v.Type()
+		if t.Kind() == reflect.Struct {
+			return engine.tbNameForMap(v)
+		}
+		return engine.Quote(fmt.Sprintf("%v", tablename))
+	}
+	return ""
+}

interface.go

@@ -87,13 +87,15 @@ type EngineInterface interface {
 	SetDefaultCacher(core.Cacher)
 	SetLogLevel(core.LogLevel)
 	SetMapper(core.IMapper)
+	SetSchema(string)
 	SetTZDatabase(tz *time.Location)
 	SetTZLocation(tz *time.Location)
 	ShowSQL(show ...bool)
 	Sync(...interface{}) error
 	Sync2(...interface{}) error
 	StoreEngine(storeEngine string) *Session
 	TableInfo(bean interface{}) *Table
+	TableName(interface{}, ...bool) string
 	UnMapType(reflect.Type)
 }