refactor select statement

2023-07-11 11:33:53 +08:00 · 2023-07-11 11:33:53 +08:00 · 13cff1f270
parent 04e3462ed0
commit 13cff1f270
5 changed files with 189 additions and 161 deletions
--- a/dialects/dialect.go
+++ b/dialects/dialect.go
@ -85,8 +85,6 @@ type Dialect interface {
 	AddColumnSQL(tableName string, col *schemas.Column) string
 	ModifyColumnSQL(tableName string, col *schemas.Column) string

-	ForUpdateSQL(query string) string
-
 	Filters() []Filter
 	SetParams(params map[string]string)
 }
@ -245,11 +243,6 @@ func (db *Base) ModifyColumnSQL(tableName string, col *schemas.Column) string {
 	return fmt.Sprintf("ALTER TABLE %s MODIFY COLUMN %s", db.quoter.Quote(tableName), s)
 }

-// ForUpdateSQL returns for updateSQL
-func (db *Base) ForUpdateSQL(query string) string {
-	return query + " FOR UPDATE"
-}
-
 // SetParams set params
 func (db *Base) SetParams(params map[string]string) {
 }
--- a/dialects/mssql.go
+++ b/dialects/mssql.go
@ -669,10 +669,6 @@ func (db *mssql) CreateTableSQL(ctx context.Context, queryer core.Queryer, table
 	return b.String(), true, nil
 }

-func (db *mssql) ForUpdateSQL(query string) string {
-	return query
-}
-
 func (db *mssql) Filters() []Filter {
 	return []Filter{}
 }
--- a/dialects/sqlite3.go
+++ b/dialects/sqlite3.go
@ -193,11 +193,11 @@ func (db *sqlite3) Features() *DialectFeatures {
 func (db *sqlite3) SetQuotePolicy(quotePolicy QuotePolicy) {
 	switch quotePolicy {
 	case QuotePolicyNone:
-		var q = sqlite3Quoter
+		q := sqlite3Quoter
 		q.IsReserved = schemas.AlwaysNoReserve
 		db.quoter = q
 	case QuotePolicyReserved:
-		var q = sqlite3Quoter
+		q := sqlite3Quoter
 		q.IsReserved = db.IsReserved
 		db.quoter = q
 	case QuotePolicyAlways:
@ -291,10 +291,6 @@ func (db *sqlite3) DropIndexSQL(tableName string, index *schemas.Index) string {
 	return fmt.Sprintf("DROP INDEX %v", db.Quoter().Quote(idxName))
 }

-func (db *sqlite3) ForUpdateSQL(query string) string {
-	return query
-}
-
 func (db *sqlite3) IsColumnExist(queryer core.Queryer, ctx context.Context, tableName, colName string) (bool, error) {
 	query := "SELECT * FROM " + tableName + " LIMIT 0"
 	rows, err := queryer.QueryContext(ctx, query)
@ -320,7 +316,7 @@ func (db *sqlite3) IsColumnExist(queryer core.Queryer, ctx context.Context, tabl
 // splitColStr splits a sqlite col strings as fields
 func splitColStr(colStr string) []string {
 	colStr = strings.TrimSpace(colStr)
-	var results = make([]string, 0, 10)
+	results := make([]string, 0, 10)
 	var lastIdx int
 	var hasC, hasQuote bool
 	for i, c := range colStr {
--- a/internal/statements/query.go
+++ b/internal/statements/query.go
@ -7,6 +7,7 @@ package statements
 import (
 	"errors"
 	"fmt"
+	"io"
 	"reflect"
 	"strings"

@ -59,7 +60,7 @@ func (statement *Statement) GenQuerySQL(sqlOrArgs ...interface{}) (string, []int
 		return "", nil, err
 	}

-	return statement.genSelectSQL(columnStr, true, true)
+	return statement.genSelectSQL(columnStr, true)
 }

 // GenSumSQL generates sum SQL
@ -87,7 +88,7 @@ func (statement *Statement) GenSumSQL(bean interface{}, columns ...string) (stri
 		return "", nil, err
 	}

-	return statement.genSelectSQL(sumSelect, true, true)
+	return statement.genSelectSQL(sumSelect, true)
 }

 // GenGetSQL generates Get SQL
@ -139,7 +140,7 @@ func (statement *Statement) GenGetSQL(bean interface{}) (string, []interface{},
 		}
 	}

-	return statement.genSelectSQL(columnStr, true, true)
+	return statement.genSelectSQL(columnStr, true)
 }

 // GenCountSQL generates the SQL for counting
@ -176,7 +177,7 @@ func (statement *Statement) GenCountSQL(beans ...interface{}) (string, []interfa
 		subQuerySelect = selectSQL
 	}

-	sqlStr, condArgs, err := statement.genSelectSQL(subQuerySelect, false, false)
+	sqlStr, condArgs, err := statement.genSelectSQL(subQuerySelect, false)
 	if err != nil {
 		return "", nil, err
 	}
@ -218,155 +219,194 @@ func (statement *Statement) writeLimitOffset(w builder.Writer) error {
 	return nil
 }

-func (statement *Statement) genSelectSQL(columnStr string, needLimit, needOrderBy bool) (string, []interface{}, error) {
-	var (
-		distinct      string
-		dialect       = statement.dialect
-		top, whereStr string
-		mssqlCondi    = builder.NewWriter()
-	)
-
-	if statement.IsDistinct && !strings.HasPrefix(columnStr, "count") {
-		distinct = "DISTINCT "
+func (statement *Statement) writeTop(w builder.Writer) error {
+	if statement.dialect.URI().DBType != schemas.MSSQL {
+		return nil
 	}

-	condWriter := builder.NewWriter()
-	if err := statement.cond.WriteTo(statement.QuoteReplacer(condWriter)); err != nil {
-		return "", nil, err
+	if statement.LimitN == nil {
+		return nil
+	}
+	_, err := fmt.Fprintf(w, " TOP %d", *statement.LimitN)
+	return err
+}
+
+func (statement *Statement) writeDistinct(w builder.Writer) error {
+	if statement.IsDistinct && !strings.HasPrefix(statement.SelectStr, "count(") {
+		_, err := fmt.Fprint(w, " DISTINCT")
+		return err
+	}
+	return nil
+}
+
+func (statement *Statement) writeSelect(w *builder.BytesWriter, columnStr string) error {
+	if _, err := fmt.Fprintf(w, "SELECT "); err != nil {
+		return err
+	}
+	if err := statement.writeDistinct(w); err != nil {
+		return err
+	}
+	if err := statement.writeTop(w); err != nil {
+		return err
 	}

-	if condWriter.Len() > 0 {
-		whereStr = " WHERE "
+	_, err := fmt.Fprint(w, columnStr)
+	return err
+}
+
+func (statement *Statement) writeWhere(w *builder.BytesWriter) error {
+	if !statement.cond.IsValid() {
+		return statement.writeMssqlPaginationCond(w)
+	}
+	if _, err := fmt.Fprint(w, " WHERE "); err != nil {
+		return err
+	}
+	if err := statement.cond.WriteTo(statement.QuoteReplacer(w)); err != nil {
+		return err
+	}
+	return statement.writeMssqlPaginationCond(w)
+}
+
+func (statement *Statement) writeMssqlPaginationCond(w *builder.BytesWriter) error {
+	if statement.dialect.URI().DBType != schemas.MSSQL || statement.LimitN == nil {
+		return nil
 	}

-	pLimitN := statement.LimitN
-	if dialect.URI().DBType == schemas.MSSQL {
-		if pLimitN != nil {
-			LimitNValue := *pLimitN
-			top = fmt.Sprintf("TOP %d ", LimitNValue)
+	mssqlCondi := builder.NewWriter()
+
+	if err := statement.writeTop(mssqlCondi); err != nil {
+		return err
+	}
+
+	if statement.Start > 0 {
+		if statement.RefTable == nil {
+			return errors.New("unsupported query limit without reference table")
 		}
-		if statement.Start > 0 {
-			if statement.RefTable == nil {
-				return "", nil, errors.New("Unsupported query limit without reference table")
+		var column string
+		if len(statement.RefTable.PKColumns()) == 0 {
+			for _, index := range statement.RefTable.Indexes {
+				if len(index.Cols) == 1 {
+					column = index.Cols[0]
+					break
+				}
 			}
-			var column string
-			if len(statement.RefTable.PKColumns()) == 0 {
-				for _, index := range statement.RefTable.Indexes {
-					if len(index.Cols) == 1 {
-						column = index.Cols[0]
-						break
-					}
-				}
-				if len(column) == 0 {
-					column = statement.RefTable.ColumnsSeq()[0]
-				}
+			if len(column) == 0 {
+				column = statement.RefTable.ColumnsSeq()[0]
+			}
+		} else {
+			column = statement.RefTable.PKColumns()[0].Name
+		}
+		if statement.needTableName() {
+			if len(statement.TableAlias) > 0 {
+				column = fmt.Sprintf("%s.%s", statement.TableAlias, column)
 			} else {
-				column = statement.RefTable.PKColumns()[0].Name
-			}
-			if statement.needTableName() {
-				if len(statement.TableAlias) > 0 {
-					column = fmt.Sprintf("%s.%s", statement.TableAlias, column)
-				} else {
-					column = fmt.Sprintf("%s.%s", statement.TableName(), column)
-				}
+				column = fmt.Sprintf("%s.%s", statement.TableName(), column)
 			}
+		}

-			if _, err := fmt.Fprintf(mssqlCondi, "(%s NOT IN (SELECT TOP %d %s",
-				column, statement.Start, column); err != nil {
-				return "", nil, err
-			}
-			if err := statement.writeFrom(mssqlCondi); err != nil {
-				return "", nil, err
-			}
-			if whereStr != "" {
-				if _, err := fmt.Fprint(mssqlCondi, whereStr); err != nil {
-					return "", nil, err
-				}
-				if err := utils.WriteBuilder(mssqlCondi, statement.QuoteReplacer(condWriter)); err != nil {
-					return "", nil, err
-				}
-			}
-			if needOrderBy {
-				if err := statement.WriteOrderBy(mssqlCondi); err != nil {
-					return "", nil, err
-				}
-			}
-			if err := statement.WriteGroupBy(mssqlCondi); err != nil {
-				return "", nil, err
-			}
-			if _, err := fmt.Fprint(mssqlCondi, "))"); err != nil {
-				return "", nil, err
-			}
+		if _, err := fmt.Fprintf(mssqlCondi, "(%s NOT IN (SELECT TOP %d %s",
+			column, statement.Start, column); err != nil {
+			return err
+		}
+		if err := statement.writeFrom(mssqlCondi); err != nil {
+			return err
+		}
+		if err := statement.writeWhere(mssqlCondi); err != nil {
+			return err
+		}
+		if err := statement.WriteOrderBy(mssqlCondi); err != nil {
+			return err
+		}
+		if err := statement.WriteGroupBy(mssqlCondi); err != nil {
+			return err
+		}
+		if _, err := fmt.Fprint(mssqlCondi, "))"); err != nil {
+			return err
 		}
 	}

+	if statement.cond.IsValid() {
+		if _, err := fmt.Fprint(w, " AND "); err != nil {
+			return err
+		}
+	} else {
+		if _, err := fmt.Fprint(w, " WHERE "); err != nil {
+			return err
+		}
+	}
+
+	return utils.WriteBuilder(w, mssqlCondi)
+}
+
+func (statement *Statement) writeOracleLimit(w *builder.BytesWriter, columnStr string) error {
+	if statement.LimitN == nil {
+		return nil
+	}
+
+	oldString := w.String()
+	w.Reset()
+	rawColStr := columnStr
+	if rawColStr == "*" {
+		rawColStr = "at.*"
+	}
+	_, err := fmt.Fprintf(w, "SELECT %v FROM (SELECT %v,ROWNUM RN FROM (%v) at WHERE ROWNUM <= %d) aat WHERE RN > %d",
+		columnStr, rawColStr, oldString, statement.Start+*statement.LimitN, statement.Start)
+	return err
+}
+
+func (statement *Statement) genSelectSQL(columnStr string, needLimit bool) (string, []interface{}, error) {
 	buf := builder.NewWriter()
-	if _, err := fmt.Fprintf(buf, "SELECT %v%v%v", distinct, top, columnStr); err != nil {
+	if err := statement.writeSelect(buf, columnStr); err != nil {
 		return "", nil, err
 	}
 	if err := statement.writeFrom(buf); err != nil {
 		return "", nil, err
 	}
-	if whereStr != "" {
-		if _, err := fmt.Fprint(buf, whereStr); err != nil {
-			return "", nil, err
-		}
-		if err := utils.WriteBuilder(buf, statement.QuoteReplacer(condWriter)); err != nil {
-			return "", nil, err
-		}
+	if err := statement.writeWhere(buf); err != nil {
+		return "", nil, err
 	}
-	if mssqlCondi.Len() > 0 {
-		if len(whereStr) > 0 {
-			if _, err := fmt.Fprint(buf, " AND "); err != nil {
-				return "", nil, err
-			}
-		} else {
-			if _, err := fmt.Fprint(buf, " WHERE "); err != nil {
-				return "", nil, err
-			}
-		}
-
-		if err := utils.WriteBuilder(buf, mssqlCondi); err != nil {
-			return "", nil, err
-		}
-	}
-
 	if err := statement.WriteGroupBy(buf); err != nil {
 		return "", nil, err
 	}
 	if err := statement.writeHaving(buf); err != nil {
 		return "", nil, err
 	}
-	if needOrderBy {
-		if err := statement.WriteOrderBy(buf); err != nil {
-			return "", nil, err
-		}
+	if err := statement.WriteOrderBy(buf); err != nil {
+		return "", nil, err
 	}
+
+	dialect := statement.dialect
 	if needLimit {
-		if dialect.URI().DBType != schemas.MSSQL && dialect.URI().DBType != schemas.ORACLE {
+		if dialect.URI().DBType == schemas.ORACLE {
+			if err := statement.writeOracleLimit(buf, columnStr); err != nil {
+				return "", nil, err
+			}
+		} else if dialect.URI().DBType != schemas.MSSQL {
 			if err := statement.writeLimitOffset(buf); err != nil {
 				return "", nil, err
 			}
-		} else if dialect.URI().DBType == schemas.ORACLE {
-			if pLimitN != nil {
-				oldString := buf.String()
-				buf.Reset()
-				rawColStr := columnStr
-				if rawColStr == "*" {
-					rawColStr = "at.*"
-				}
-				fmt.Fprintf(buf, "SELECT %v FROM (SELECT %v,ROWNUM RN FROM (%v) at WHERE ROWNUM <= %d) aat WHERE RN > %d",
-					columnStr, rawColStr, oldString, statement.Start+*pLimitN, statement.Start)
-			}
 		}
 	}
-	if statement.IsForUpdate {
-		return dialect.ForUpdateSQL(buf.String()), buf.Args(), nil
+
+	if err := statement.writeForUpdate(buf); err != nil {
+		return "", nil, err
 	}

 	return buf.String(), buf.Args(), nil
 }

+func (statement *Statement) writeForUpdate(w io.Writer) error {
+	if !statement.IsForUpdate {
+		return nil
+	}
+
+	if statement.dialect.URI().DBType != schemas.MYSQL {
+		return errors.New("only support mysql for update")
+	}
+	_, err := fmt.Fprint(w, " FOR UPDATE")
+	return err
+}
+
 // GenExistSQL generates Exist SQL
 func (statement *Statement) GenExistSQL(bean ...interface{}) (string, []interface{}, error) {
 	if statement.RawSQL != "" {
@ -457,6 +497,35 @@ func (statement *Statement) GenExistSQL(bean ...interface{}) (string, []interfac
 	return buf.String(), buf.Args(), nil
 }

+func (statement *Statement) genSelect() string {
+	if len(statement.SelectStr) > 0 {
+		return statement.SelectStr
+	}
+
+	columnStr := statement.ColumnStr()
+	if statement.JoinStr == "" {
+		if columnStr == "" {
+			if statement.GroupByStr != "" {
+				columnStr = statement.quoteColumnStr(statement.GroupByStr)
+			} else {
+				columnStr = statement.genColumnStr()
+			}
+		}
+	} else {
+		if columnStr == "" {
+			if statement.GroupByStr != "" {
+				columnStr = statement.quoteColumnStr(statement.GroupByStr)
+			} else {
+				columnStr = "*"
+			}
+		}
+	}
+	if columnStr == "" {
+		columnStr = "*"
+	}
+	return columnStr
+}
+
 // GenFindSQL generates Find SQL
 func (statement *Statement) GenFindSQL(autoCond builder.Cond) (string, []interface{}, error) {
 	if statement.RawSQL != "" {
@ -467,33 +536,7 @@ func (statement *Statement) GenFindSQL(autoCond builder.Cond) (string, []interfa
 		return "", nil, ErrTableNotFound
 	}

-	columnStr := statement.ColumnStr()
-	if len(statement.SelectStr) > 0 {
-		columnStr = statement.SelectStr
-	} else {
-		if statement.JoinStr == "" {
-			if columnStr == "" {
-				if statement.GroupByStr != "" {
-					columnStr = statement.quoteColumnStr(statement.GroupByStr)
-				} else {
-					columnStr = statement.genColumnStr()
-				}
-			}
-		} else {
-			if columnStr == "" {
-				if statement.GroupByStr != "" {
-					columnStr = statement.quoteColumnStr(statement.GroupByStr)
-				} else {
-					columnStr = "*"
-				}
-			}
-		}
-		if columnStr == "" {
-			columnStr = "*"
-		}
-	}
-
 	statement.cond = statement.cond.And(autoCond)

-	return statement.genSelectSQL(columnStr, true, true)
+	return statement.genSelectSQL(statement.genSelect(), true)
 }
--- a/internal/statements/statement.go
+++ b/internal/statements/statement.go
@ -300,7 +300,7 @@ func (statement *Statement) WriteGroupBy(w builder.Writer) error {

 // Having generate "Having conditions" statement
 func (statement *Statement) Having(conditions string) *Statement {
-	statement.HavingStr = fmt.Sprintf("HAVING %v", statement.ReplaceQuote(conditions))
+	statement.HavingStr = conditions
 	return statement
 }

@ -308,7 +308,7 @@ func (statement *Statement) writeHaving(w builder.Writer) error {
 	if statement.HavingStr == "" {
 		return nil
 	}
-	_, err := fmt.Fprint(w, " ", statement.HavingStr)
+	_, err := fmt.Fprint(w, " HAVING ", statement.ReplaceQuote(statement.HavingStr))
 	return err
 }