From 3f4916ff8b4d957b682f1fb1d2e37c2355d85280 Mon Sep 17 00:00:00 2001
From: Pedro Henrique <system.pedrohenrique@gmail.com>
Date: Fri, 9 Dec 2022 16:08:55 -0300
Subject: [PATCH] dialect/sql/schema: fix issue with WriteDriver when using
 Postgres with Query (#3162)

---
 dialect/sql/schema/writer.go      | 1 +
 dialect/sql/schema/writer_test.go | 6 ++++++
 2 files changed, 7 insertions(+)

diff --git a/dialect/sql/schema/writer.go b/dialect/sql/schema/writer.go
index cecd3b776..ce2917d62 100644
--- a/dialect/sql/schema/writer.go
+++ b/dialect/sql/schema/writer.go
@@ -110,6 +110,7 @@ func (w *WriteDriver) Query(ctx context.Context, query string, args, res any) er
 		if rr, ok := res.(*sql.Rows); ok {
 			*rr = sql.Rows{ColumnScanner: noRows{}}
 		}
+		return nil
 	}
 	switch w.Driver.(type) {
 	case nil, nopDriver:
diff --git a/dialect/sql/schema/writer_test.go b/dialect/sql/schema/writer_test.go
index fd95d4dc4..deaf28179 100644
--- a/dialect/sql/schema/writer_test.go
+++ b/dialect/sql/schema/writer_test.go
@@ -64,6 +64,12 @@ func TestWriteDriver(t *testing.T) {
 	err = w.Exec(ctx, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id`, nil, nil)
 	require.NoError(t, err)
 	require.Equal(t, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id;`+"\n", b.String())
+
+	// batchCreator uses tx.Query when doing an insert
+	b.Reset()
+	err = w.Query(ctx, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id`, nil, nil)
+	require.NoError(t, err)
+	require.Equal(t, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id;`+"\n", b.String())
 }
 
 func TestDirWriter(t *testing.T) {