Merge branch 'master' of https://github.com/projectleo/gorm

projectleo · Jul 27, 2022 · 603f1c3 · 603f1c3
2 parents 2c91e85 + cbcbf63
commit 603f1c3
Show file tree

Hide file tree

Showing 21 changed files with 189 additions and 129 deletions.
diff --git a/.idea/.gitignore b/.idea/.gitignore
diff --git a/.idea/gorm.iml b/.idea/gorm.iml
diff --git a/.idea/markdown.xml b/.idea/markdown.xml
diff --git a/.idea/modules.xml b/.idea/modules.xml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/callback_create.go b/callback_create.go
@@ -130,7 +130,7 @@ func createCallback(scope *Scope) {
 
 		// execute create sql: no primaryField
 		if primaryField == nil {
-			if result, err := scope.SQLDB().Exec(scope.SQL, scope.SQLVars...); scope.Err(err) == nil {
+			if result, err := scope.SQLDB().ExecContext(scope.Context(), scope.SQL, scope.SQLVars...); scope.Err(err) == nil {
 				// set rows affected count
 				scope.db.RowsAffected, _ = result.RowsAffected()
 
@@ -146,7 +146,7 @@ func createCallback(scope *Scope) {
 
 		// execute create sql: lastInsertID implemention for majority of dialects
 		if lastInsertIDReturningSuffix == "" && lastInsertIDOutputInterstitial == "" {
-			if result, err := scope.SQLDB().Exec(scope.SQL, scope.SQLVars...); scope.Err(err) == nil {
+			if result, err := scope.SQLDB().ExecContext(scope.Context(), scope.SQL, scope.SQLVars...); scope.Err(err) == nil {
 				// set rows affected count
 				scope.db.RowsAffected, _ = result.RowsAffected()
 
@@ -162,7 +162,7 @@ func createCallback(scope *Scope) {
 
 		// execute create sql: dialects with additional lastInsertID requirements (currently postgres & mssql)
 		if primaryField.Field.CanAddr() {
-			if err := scope.SQLDB().QueryRow(scope.SQL, scope.SQLVars...).Scan(primaryField.Field.Addr().Interface()); scope.Err(err) == nil {
+			if err := scope.SQLDB().QueryRowContext(scope.Context(), scope.SQL, scope.SQLVars...).Scan(primaryField.Field.Addr().Interface()); scope.Err(err) == nil {
 				primaryField.IsBlank = false
 				scope.db.RowsAffected = 1
 			}

diff --git a/callback_query.go b/callback_query.go
@@ -65,7 +65,7 @@ func queryCallback(scope *Scope) {
 			scope.SQL = fmt.Sprint(str) + scope.SQL
 		}
 
-		if rows, err := scope.SQLDB().Query(scope.SQL, scope.SQLVars...); scope.Err(err) == nil {
+		if rows, err := scope.SQLDB().QueryContext(scope.Context(), scope.SQL, scope.SQLVars...); scope.Err(err) == nil {
 			defer rows.Close()
 
 			columns, _ := rows.Columns()

diff --git a/callback_row_query.go b/callback_row_query.go
@@ -29,9 +29,9 @@ func rowQueryCallback(scope *Scope) {
 		}
 
 		if rowResult, ok := result.(*RowQueryResult); ok {
-			rowResult.Row = scope.SQLDB().QueryRow(scope.SQL, scope.SQLVars...)
+			rowResult.Row = scope.SQLDB().QueryRowContext(scope.Context(), scope.SQL, scope.SQLVars...)
 		} else if rowsResult, ok := result.(*RowsQueryResult); ok {
-			rowsResult.Rows, rowsResult.Error = scope.SQLDB().Query(scope.SQL, scope.SQLVars...)
+			rowsResult.Rows, rowsResult.Error = scope.SQLDB().QueryContext(scope.Context(), scope.SQL, scope.SQLVars...)
 		}
 	}
 }
diff --git a/customize_column_test.go b/customize_column_test.go
@@ -26,7 +26,7 @@ func TestCustomizeColumn(t *testing.T) {
 	DB.AutoMigrate(&CustomizeColumn{})
 
 	scope := DB.NewScope(&CustomizeColumn{})
-	if !scope.Dialect().HasColumn(scope.TableName(), col) {
+	if !scope.Dialect().HasColumn(scope.Context(), scope.TableName(), col) {
 		t.Errorf("CustomizeColumn should have column %s", col)
 	}
 

diff --git a/dialect.go b/dialect.go
@@ -1,6 +1,7 @@
 package gorm
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"reflect"
@@ -24,17 +25,17 @@ type Dialect interface {
 	DataTypeOf(field *StructField) string
 
 	// HasIndex check has index or not
-	HasIndex(tableName string, indexName string) bool
+	HasIndex(ctx context.Context, tableName string, indexName string) bool
 	// HasForeignKey check has foreign key or not
-	HasForeignKey(tableName string, foreignKeyName string) bool
+	HasForeignKey(ctx context.Context, tableName string, foreignKeyName string) bool
 	// RemoveIndex remove index
-	RemoveIndex(tableName string, indexName string) error
+	RemoveIndex(ctx context.Context, tableName string, indexName string) error
 	// HasTable check has table or not
-	HasTable(tableName string) bool
+	HasTable(ctx context.Context, tableName string) bool
 	// HasColumn check has column or not
-	HasColumn(tableName string, columnName string) bool
+	HasColumn(ctx context.Context, tableName string, columnName string) bool
 	// ModifyColumn modify column's type
-	ModifyColumn(tableName string, columnName string, typ string) error
+	ModifyColumn(ctx context.Context, tableName string, columnName string, typ string) error
 
 	// LimitAndOffsetSQL return generated SQL with Limit and Offset, as mssql has special case
 	LimitAndOffsetSQL(limit, offset interface{}) (string, error)
@@ -54,7 +55,7 @@ type Dialect interface {
 	NormalizeIndexAndColumn(indexName, columnName string) (string, string)
 
 	// CurrentDatabase return current database name
-	CurrentDatabase() string
+	CurrentDatabase(ctx context.Context) string
 }
 
 var dialectsMap = map[string]Dialect{}
@@ -67,9 +68,9 @@ func newDialect(name string, db SQLCommon) Dialect {
 	}
 
 	fmt.Printf("`%v` is not officially supported, running under compatibility mode.\n", name)
-	commontDialect := &commonDialect{}
-	commontDialect.SetDB(db)
-	return commontDialect
+	cd := &commonDialect{}
+	cd.SetDB(db)
+	return cd
 }
 
 // RegisterDialect register new dialect
@@ -138,10 +139,10 @@ var ParseFieldStructForDialect = func(field *StructField, dialect Dialect) (fiel
 	return fieldValue, dataType, size, strings.TrimSpace(additionalType)
 }
 
-func currentDatabaseAndTable(dialect Dialect, tableName string) (string, string) {
+func currentDatabaseAndTable(ctx context.Context, dialect Dialect, tableName string) (string, string) {
 	if strings.Contains(tableName, ".") {
 		splitStrings := strings.SplitN(tableName, ".", 2)
 		return splitStrings[0], splitStrings[1]
 	}
-	return dialect.CurrentDatabase(), tableName
+	return dialect.CurrentDatabase(ctx), tableName
 }
diff --git a/dialect_common.go b/dialect_common.go
@@ -1,6 +1,7 @@
 package gorm
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"regexp"
@@ -99,43 +100,43 @@ func (s *commonDialect) DataTypeOf(field *StructField) string {
 	return fmt.Sprintf("%v %v", sqlType, additionalType)
 }
 
-func (s commonDialect) HasIndex(tableName string, indexName string) bool {
+func (s commonDialect) HasIndex(ctx context.Context, tableName string, indexName string) bool {
 	var count int
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
-	s.db.QueryRow("SELECT count(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE table_schema = ? AND table_name = ? AND index_name = ?", currentDatabase, tableName, indexName).Scan(&count)
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE table_schema = ? AND table_name = ? AND index_name = ?", currentDatabase, tableName, indexName).Scan(&count)
 	return count > 0
 }
 
-func (s commonDialect) RemoveIndex(tableName string, indexName string) error {
-	_, err := s.db.Exec(fmt.Sprintf("DROP INDEX %v", indexName))
+func (s commonDialect) RemoveIndex(ctx context.Context, tableName string, indexName string) error {
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf("DROP INDEX %v", indexName))
 	return err
 }
 
-func (s commonDialect) HasForeignKey(tableName string, foreignKeyName string) bool {
+func (s commonDialect) HasForeignKey(ctx context.Context, tableName string, foreignKeyName string) bool {
 	return false
 }
 
-func (s commonDialect) HasTable(tableName string) bool {
+func (s commonDialect) HasTable(ctx context.Context, tableName string) bool {
 	var count int
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
-	s.db.QueryRow("SELECT count(*) FROM INFORMATION_SCHEMA.TABLES WHERE table_schema = ? AND table_name = ?", currentDatabase, tableName).Scan(&count)
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM INFORMATION_SCHEMA.TABLES WHERE table_schema = ? AND table_name = ?", currentDatabase, tableName).Scan(&count)
 	return count > 0
 }
 
-func (s commonDialect) HasColumn(tableName string, columnName string) bool {
+func (s commonDialect) HasColumn(ctx context.Context, tableName string, columnName string) bool {
 	var count int
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
-	s.db.QueryRow("SELECT count(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE table_schema = ? AND table_name = ? AND column_name = ?", currentDatabase, tableName, columnName).Scan(&count)
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE table_schema = ? AND table_name = ? AND column_name = ?", currentDatabase, tableName, columnName).Scan(&count)
 	return count > 0
 }
 
-func (s commonDialect) ModifyColumn(tableName string, columnName string, typ string) error {
-	_, err := s.db.Exec(fmt.Sprintf("ALTER TABLE %v ALTER COLUMN %v TYPE %v", tableName, columnName, typ))
+func (s commonDialect) ModifyColumn(ctx context.Context, tableName string, columnName string, typ string) error {
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %v ALTER COLUMN %v TYPE %v", tableName, columnName, typ))
 	return err
 }
 
-func (s commonDialect) CurrentDatabase() (name string) {
-	s.db.QueryRow("SELECT DATABASE()").Scan(&name)
+func (s commonDialect) CurrentDatabase(ctx context.Context) (name string) {
+	s.db.QueryRowContext(ctx, "SELECT DATABASE()").Scan(&name)
 	return
 }
 

diff --git a/dialect_mysql.go b/dialect_mysql.go
@@ -1,6 +1,7 @@
 package gorm
 
 import (
+	"context"
 	"crypto/sha1"
 	"database/sql"
 	"fmt"
@@ -129,13 +130,13 @@ func (s *mysql) DataTypeOf(field *StructField) string {
 	return fmt.Sprintf("%v %v", sqlType, additionalType)
 }
 
-func (s mysql) RemoveIndex(tableName string, indexName string) error {
-	_, err := s.db.Exec(fmt.Sprintf("DROP INDEX %v ON %v", indexName, s.Quote(tableName)))
+func (s mysql) RemoveIndex(ctx context.Context, tableName string, indexName string) error {
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf("DROP INDEX %v ON %v", indexName, s.Quote(tableName)))
 	return err
 }
 
-func (s mysql) ModifyColumn(tableName string, columnName string, typ string) error {
-	_, err := s.db.Exec(fmt.Sprintf("ALTER TABLE %v MODIFY COLUMN %v %v", tableName, columnName, typ))
+func (s mysql) ModifyColumn(ctx context.Context, tableName string, columnName string, typ string) error {
+	_, err := s.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %v MODIFY COLUMN %v %v", tableName, columnName, typ))
 	return err
 }
 
@@ -162,18 +163,18 @@ func (s mysql) LimitAndOffsetSQL(limit, offset interface{}) (sql string, err err
 	return
 }
 
-func (s mysql) HasForeignKey(tableName string, foreignKeyName string) bool {
+func (s mysql) HasForeignKey(ctx context.Context, tableName string, foreignKeyName string) bool {
 	var count int
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
-	s.db.QueryRow("SELECT count(*) FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE CONSTRAINT_SCHEMA=? AND TABLE_NAME=? AND CONSTRAINT_NAME=? AND CONSTRAINT_TYPE='FOREIGN KEY'", currentDatabase, tableName, foreignKeyName).Scan(&count)
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE CONSTRAINT_SCHEMA=? AND TABLE_NAME=? AND CONSTRAINT_NAME=? AND CONSTRAINT_TYPE='FOREIGN KEY'", currentDatabase, tableName, foreignKeyName).Scan(&count)
 	return count > 0
 }
 
-func (s mysql) HasTable(tableName string) bool {
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
+func (s mysql) HasTable(ctx context.Context, tableName string) bool {
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
 	var name string
 	// allow mysql database name with '-' character
-	if err := s.db.QueryRow(fmt.Sprintf("SHOW TABLES FROM `%s` WHERE `Tables_in_%s` = ?", currentDatabase, currentDatabase), tableName).Scan(&name); err != nil {
+	if err := s.db.QueryRowContext(ctx, fmt.Sprintf("SHOW TABLES FROM `%s` WHERE `Tables_in_%s` = ?", currentDatabase, currentDatabase), tableName).Scan(&name); err != nil {
 		if err == sql.ErrNoRows {
 			return false
 		}
@@ -183,28 +184,28 @@ func (s mysql) HasTable(tableName string) bool {
 	}
 }
 
-func (s mysql) HasIndex(tableName string, indexName string) bool {
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
-	if rows, err := s.db.Query(fmt.Sprintf("SHOW INDEXES FROM `%s` FROM `%s` WHERE Key_name = ?", tableName, currentDatabase), indexName); err != nil {
+func (s mysql) HasIndex(ctx context.Context, tableName string, indexName string) bool {
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
+	if rows, err := s.db.QueryContext(ctx, fmt.Sprintf("SHOW INDEXES FROM `%s` FROM `%s` WHERE Key_name = ?", tableName, currentDatabase), indexName); err != nil {
 		panic(err)
 	} else {
 		defer rows.Close()
 		return rows.Next()
 	}
 }
 
-func (s mysql) HasColumn(tableName string, columnName string) bool {
-	currentDatabase, tableName := currentDatabaseAndTable(&s, tableName)
-	if rows, err := s.db.Query(fmt.Sprintf("SHOW COLUMNS FROM `%s` FROM `%s` WHERE Field = ?", tableName, currentDatabase), columnName); err != nil {
+func (s mysql) HasColumn(ctx context.Context, tableName string, columnName string) bool {
+	currentDatabase, tableName := currentDatabaseAndTable(ctx, &s, tableName)
+	if rows, err := s.db.QueryContext(ctx, fmt.Sprintf("SHOW COLUMNS FROM `%s` FROM `%s` WHERE Field = ?", tableName, currentDatabase), columnName); err != nil {
 		panic(err)
 	} else {
 		defer rows.Close()
 		return rows.Next()
 	}
 }
 
-func (s mysql) CurrentDatabase() (name string) {
-	s.db.QueryRow("SELECT DATABASE()").Scan(&name)
+func (s mysql) CurrentDatabase(ctx context.Context) (name string) {
+	s.db.QueryRowContext(ctx, "SELECT DATABASE()").Scan(&name)
 	return
 }
 

diff --git a/dialect_postgres.go b/dialect_postgres.go
@@ -1,6 +1,7 @@
 package gorm
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"reflect"
@@ -91,32 +92,32 @@ func (s *postgres) DataTypeOf(field *StructField) string {
 	return fmt.Sprintf("%v %v", sqlType, additionalType)
 }
 
-func (s postgres) HasIndex(tableName string, indexName string) bool {
+func (s postgres) HasIndex(ctx context.Context, tableName string, indexName string) bool {
 	var count int
-	s.db.QueryRow("SELECT count(*) FROM pg_indexes WHERE tablename = $1 AND indexname = $2 AND schemaname = CURRENT_SCHEMA()", tableName, indexName).Scan(&count)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM pg_indexes WHERE tablename = $1 AND indexname = $2 AND schemaname = CURRENT_SCHEMA()", tableName, indexName).Scan(&count)
 	return count > 0
 }
 
-func (s postgres) HasForeignKey(tableName string, foreignKeyName string) bool {
+func (s postgres) HasForeignKey(ctx context.Context, tableName string, foreignKeyName string) bool {
 	var count int
-	s.db.QueryRow("SELECT count(con.conname) FROM pg_constraint con WHERE $1::regclass::oid = con.conrelid AND con.conname = $2 AND con.contype='f'", tableName, foreignKeyName).Scan(&count)
+	s.db.QueryRowContext(ctx, "SELECT count(con.conname) FROM pg_constraint con WHERE $1::regclass::oid = con.conrelid AND con.conname = $2 AND con.contype='f'", tableName, foreignKeyName).Scan(&count)
 	return count > 0
 }
 
-func (s postgres) HasTable(tableName string) bool {
+func (s postgres) HasTable(ctx context.Context, tableName string) bool {
 	var count int
-	s.db.QueryRow("SELECT count(*) FROM INFORMATION_SCHEMA.tables WHERE table_name = $1 AND table_type = 'BASE TABLE' AND table_schema = CURRENT_SCHEMA()", tableName).Scan(&count)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM INFORMATION_SCHEMA.tables WHERE table_name = $1 AND table_type = 'BASE TABLE' AND table_schema = CURRENT_SCHEMA()", tableName).Scan(&count)
 	return count > 0
 }
 
-func (s postgres) HasColumn(tableName string, columnName string) bool {
+func (s postgres) HasColumn(ctx context.Context, tableName string, columnName string) bool {
 	var count int
-	s.db.QueryRow("SELECT count(*) FROM INFORMATION_SCHEMA.columns WHERE table_name = $1 AND column_name = $2 AND table_schema = CURRENT_SCHEMA()", tableName, columnName).Scan(&count)
+	s.db.QueryRowContext(ctx, "SELECT count(*) FROM INFORMATION_SCHEMA.columns WHERE table_name = $1 AND column_name = $2 AND table_schema = CURRENT_SCHEMA()", tableName, columnName).Scan(&count)
 	return count > 0
 }
 
-func (s postgres) CurrentDatabase() (name string) {
-	s.db.QueryRow("SELECT CURRENT_DATABASE()").Scan(&name)
+func (s postgres) CurrentDatabase(ctx context.Context) (name string) {
+	s.db.QueryRowContext(ctx, "SELECT CURRENT_DATABASE()").Scan(&name)
 	return
 }