Fix issues & optimize codes

Author: BetaCat0

builder_limit.go

@@ -14,28 +14,25 @@ func (b *Builder) limitWriteTo(w Writer) error {
 if strings.TrimSpace(b.dialect) == "" {
 return errors.New("field `dialect` must be set up when performing LIMIT, try use `Dialect(dbType)` at first")
 }
-if !(b.optype == selectType || b.optype == unionType) {
-return errors.New("LIMIT operation is limited in SELECT and UNION")
-}
 
 if b.limitation != nil {
 limit := b.limitation
 if limit.offset < 0 || limit.limitN <= 0 {
 return errors.New("unexpected offset/limitN")
 }
-
-selects := b.selects
 // erase limit condition
 b.limitation = nil
-// flush writer, both buffer & args
 ow := w.(*BytesWriter)
-ow.writer.Reset()
-ow.args = nil
 
 var final *Builder
 
 switch strings.ToLower(strings.TrimSpace(b.dialect)) {
 case ORACLE:
+// flush writer, both buffer & args
+ow.writer.Reset()
+ow.args = nil
+
+selects := b.selects
 b.selects = append(selects, "ROWNUM RN")
 if limit.offset == 0 {
 final = Dialect(b.dialect).Select(selects...).From("at", b).PK(b.pk...).
@@ -54,14 +51,22 @@ func (b *Builder) limitWriteTo(w Writer) error {
 
 return final.WriteTo(ow)
 case SQLITE, MYSQL, POSTGRES:
-b.WriteTo(ow)
+// if type UNION, we need to write previous content back to current writer
+if b.optype == unionType {
+b.WriteTo(ow)
+}
 
 if limit.offset == 0 {
 fmt.Fprint(ow, " LIMIT ", limit.limitN)
 } else {
 fmt.Fprintf(ow, " LIMIT %v OFFSET %v", limit.limitN, limit.offset)
 }
 case MSSQL:
+// flush writer, both buffer & args
+ow.writer.Reset()
+ow.args = nil
+
+selects := b.selects
 if limit.offset == 0 {
 if len(selects) == 0 {
 selects = append(selects, "*")

builder_select.go

@@ -104,7 +104,6 @@ func (b *Builder) selectWriteTo(w Writer) error {
 if err := b.limitWriteTo(w); err != nil {
 return err
 }
-
 }
 
 return nil

builder_select_test.go

@@ -178,9 +178,9 @@ func TestBuilder_Limit(t *testing.T) {
 // union with limit -- MsSQL style
 sql, args, err = Dialect(MSSQL).Select("a", "b", "c").From("table1").
 PK("id1").Where(Eq{"a": 1}).OrderBy("a ASC").Limit(5, 6).Union("ALL",
-Select("a", "b").From("table1").Where(Eq{"a": 2}).OrderBy("a DESC").Limit(10)).
+Select("a", "b").From("table1").Where(Eq{"b": 2}).OrderBy("a DESC").Limit(10)).
 OrderBy("b DESC").Limit(7).ToSQL()
 assert.NoError(t, err)
-// assert.EqualValues(t, 6, len(args))
+assert.EqualValues(t, 4, len(args))
 fmt.Println(sql, args)
 }

builder_union.go

@@ -28,13 +28,24 @@ func (b *Builder) unionWriteTo(w Writer) error {
 fmt.Fprint(w, "(")
 
 if current.limitation != nil {
-tw := NewWriter()
-if err := current.selectWriteTo(tw); err != nil {
-return err
+switch current.dialect {
+case ORACLE, MSSQL:
+// if dialect is Oracle or Mssql, we need to make a copy of current context
+// coz the result of current will turn into a sub-query when building LIMIT query
+tw := NewWriter()
+if err := current.selectWriteTo(tw); err != nil {
+return err
+}
+
+fmt.Fprintf(w, tw.writer.String())
+w.(*BytesWriter).args = append(w.(*BytesWriter).args, tw.args...)
+case MYSQL, SQLITE, POSTGRES:
+if err := current.selectWriteTo(w); err != nil {
+return err
+}
+default:
+return ErrNotSupportType
 }
-
-fmt.Fprintf(w, tw.writer.String())
-w.(*BytesWriter).args = append(w.(*BytesWriter).args, tw.args...)
 } else {
 if err := current.selectWriteTo(w); err != nil {
 return err