From f34e0bfadf5e4e24dcaf570f5dd43a5620a7f208 Mon Sep 17 00:00:00 2001
From: XIE Long <long.xie.lnu@gmail.com>
Date: Sat, 2 Oct 2021 09:20:30 +0800
Subject: [PATCH] Add mysql "upsert" with column aliases

---
 exp/insert_clauses.go               | 20 +++++++++++++++++++
 insert_dataset.go                   |  7 ++++++-
 insert_dataset_test.go              | 31 +++++++++++++++++++++++++++++
 sqlgen/insert_sql_generator.go      |  4 ++++
 sqlgen/insert_sql_generator_test.go | 11 ++++++++++
 5 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/exp/insert_clauses.go b/exp/insert_clauses.go
index 4412d6bc..ac7a330d 100644
--- a/exp/insert_clauses.go
+++ b/exp/insert_clauses.go
@@ -28,6 +28,10 @@ type (
 		HasRows() bool
 		SetRows(rows []interface{}) InsertClauses
 
+		HasAlias() bool
+		Alias() IdentifierExpression
+		SetAlias(ie IdentifierExpression) InsertClauses
+
 		Vals() [][]interface{}
 		HasVals() bool
 		SetVals(vals [][]interface{}) InsertClauses
@@ -41,6 +45,7 @@ type (
 		cols         ColumnListExpression
 		into         Expression
 		returning    ColumnListExpression
+		alias        IdentifierExpression
 		rows         []interface{}
 		values       [][]interface{}
 		from         AppendableExpression
@@ -62,6 +67,7 @@ func (ic *insertClauses) clone() *insertClauses {
 		cols:         ic.cols,
 		into:         ic.into,
 		returning:    ic.returning,
+		alias:        ic.alias,
 		rows:         ic.rows,
 		values:       ic.values,
 		from:         ic.from,
@@ -117,6 +123,20 @@ func (ic *insertClauses) HasReturning() bool {
 	return ic.returning != nil && !ic.returning.IsEmpty()
 }
 
+func (ic *insertClauses) HasAlias() bool {
+	return ic.alias != nil
+}
+
+func (ic *insertClauses) Alias() IdentifierExpression {
+	return ic.alias
+}
+
+func (ic *insertClauses) SetAlias(ie IdentifierExpression) InsertClauses {
+	ret := ic.clone()
+	ret.alias = ie
+	return ret
+}
+
 func (ic *insertClauses) SetReturning(cl ColumnListExpression) InsertClauses {
 	ret := ic.clone()
 	ret.returning = cl
diff --git a/insert_dataset.go b/insert_dataset.go
index 3ed270db..5404ee82 100644
--- a/insert_dataset.go
+++ b/insert_dataset.go
@@ -242,7 +242,12 @@ func (id *InsertDataset) AppendSQL(b sb.SQLBuilder) {
 }
 
 func (id *InsertDataset) GetAs() exp.IdentifierExpression {
-	return nil
+	return id.clauses.Alias()
+}
+
+// Sets the alias for this dataset. This is typically used when using a Dataset as MySQL upsert
+func (id *InsertDataset) As(alias string) *InsertDataset {
+	return id.copy(id.clauses.SetAlias(T(alias)))
 }
 
 func (id *InsertDataset) ReturnsColumns() bool {
diff --git a/insert_dataset_test.go b/insert_dataset_test.go
index b7e671b7..2f43c464 100644
--- a/insert_dataset_test.go
+++ b/insert_dataset_test.go
@@ -369,6 +369,37 @@ func (ids *insertDatasetSuite) TestOnConflict() {
 	)
 }
 
+func (ids *insertDatasetSuite) TestAs() {
+	du := goqu.DoUpdate("other_items", goqu.Record{"new.a": 1})
+
+	bd := goqu.Insert("items").As("new")
+	ids.assertCases(
+		insertTestCase{
+			ds: bd.OnConflict(nil),
+			clauses: exp.NewInsertClauses().SetInto(goqu.C("items")).
+				SetAlias(exp.NewIdentifierExpression("", "new", "")),
+		},
+		insertTestCase{
+			ds: bd.OnConflict(goqu.DoNothing()),
+			clauses: exp.NewInsertClauses().SetInto(goqu.C("items")).
+				SetAlias(exp.NewIdentifierExpression("", "new", "")).
+				SetOnConflict(goqu.DoNothing()),
+		},
+		insertTestCase{
+			ds: bd.OnConflict(du),
+			clauses: exp.NewInsertClauses().
+				SetAlias(exp.NewIdentifierExpression("", "new", "")).
+				SetInto(goqu.C("items")).SetOnConflict(du),
+		},
+		insertTestCase{
+			ds: bd,
+			clauses: exp.NewInsertClauses().
+				SetAlias(exp.NewIdentifierExpression("", "new", "")).
+				SetInto(goqu.C("items")),
+		},
+	)
+}
+
 func (ids *insertDatasetSuite) TestClearOnConflict() {
 	du := goqu.DoUpdate("other_items", goqu.Record{"a": 1})
 
diff --git a/sqlgen/insert_sql_generator.go b/sqlgen/insert_sql_generator.go
index 86dc5370..1c6105b9 100644
--- a/sqlgen/insert_sql_generator.go
+++ b/sqlgen/insert_sql_generator.go
@@ -100,6 +100,10 @@ func (isg *insertSQLGenerator) InsertSQL(b sb.SQLBuilder, ic exp.InsertClauses)
 	default:
 		isg.defaultValuesSQL(b)
 	}
+	if ic.HasAlias() {
+		b.Write(isg.DialectOptions().AsFragment)
+		isg.ExpressionSQLGenerator().Generate(b, ic.Alias())
+	}
 	isg.onConflictSQL(b, ic.OnConflict())
 }
 
diff --git a/sqlgen/insert_sql_generator_test.go b/sqlgen/insert_sql_generator_test.go
index 4c92e1ce..95479c56 100644
--- a/sqlgen/insert_sql_generator_test.go
+++ b/sqlgen/insert_sql_generator_test.go
@@ -257,6 +257,9 @@ func (igs *insertSQLGeneratorSuite) TestGenerate_onConflict() {
 		})
 	icDn := ic.SetOnConflict(exp.NewDoNothingConflictExpression())
 	icDu := ic.SetOnConflict(exp.NewDoUpdateConflictExpression("test", exp.Record{"a": "b"}))
+	icAsDu := ic.SetAlias(exp.NewIdentifierExpression("", "new", "")).SetOnConflict(
+		exp.NewDoUpdateConflictExpression("test", exp.Record{"a": exp.NewIdentifierExpression("", "new", "a")}),
+	)
 	icDoc := ic.SetOnConflict(exp.NewDoUpdateConflictExpression("on constraint test", exp.Record{"a": "b"}))
 	icDuw := ic.SetOnConflict(
 		exp.NewDoUpdateConflictExpression("test", exp.Record{"a": "b"}).Where(exp.Ex{"foo": true}),
@@ -283,6 +286,14 @@ func (igs *insertSQLGeneratorSuite) TestGenerate_onConflict() {
 			args:       []interface{}{"a1", "b"},
 		},
 
+		insertTestCase{clause: icAsDu, sql: `INSERT INTO "test" ("a") VALUES ('a1') AS "new" on conflict (test) do update set "a"="new"."a"`},
+		insertTestCase{
+			clause:     icAsDu,
+			sql:        `INSERT INTO "test" ("a") VALUES (?) AS "new" on conflict (test) do update set "a"="new"."a"`,
+			isPrepared: true,
+			args:       []interface{}{"a1"},
+		},
+
 		insertTestCase{clause: icDoc, sql: `INSERT INTO "test" ("a") VALUES ('a1') on conflict on constraint test do update set "a"='b'`},
 		insertTestCase{
 			clause:     icDoc,