snapshot of https://github.com/ent/ent.git (master, 2026-09-15; blobs > 131072B omitted)

dialect/sql/driver_test.go

@@ -0,0 +1,93 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"context"
+	"testing"
+
+	"entgo.io/ent/dialect"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithVars(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	db.SetMaxOpenConns(1)
+	drv := OpenDB(dialect.Postgres, db)
+	mock.ExpectExec("SET foo = 'bar'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	rows := &Rows{}
+	err = drv.Query(
+		WithVar(context.Background(), "foo", "bar"),
+		"SELECT 1",
+		[]any{},
+		rows,
+	)
+	require.NoError(t, err)
+	require.NoError(t, rows.Close(), "rows should be closed to release the connection")
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	mock.ExpectExec("SET foo = 'bar'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("SET foo = 'baz'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	err = drv.Query(
+		WithVar(WithVar(context.Background(), "foo", "bar"), "foo", "baz"),
+		"SELECT 1",
+		[]any{},
+		rows,
+	)
+	require.NoError(t, err)
+	require.NoError(t, rows.Close(), "rows should be closed to release the connection")
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	mock.ExpectBegin()
+	mock.ExpectExec("SET foo = 'bar'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
+	mock.ExpectCommit()
+	tx, err := drv.Tx(context.Background())
+	require.NoError(t, err)
+	err = tx.Query(
+		WithVar(context.Background(), "foo", "bar"),
+		"SELECT 1",
+		[]any{},
+		rows,
+	)
+	require.NoError(t, err)
+	require.NoError(t, tx.Commit())
+	require.NoError(t, mock.ExpectationsWereMet())
+	// Rows should not be closed to release the session,
+	// as a transaction is always scoped to a single connection.
+
+	mock.ExpectExec("SET foo = 'qux'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("INSERT INTO users DEFAULT VALUES").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	err = drv.Exec(
+		WithVar(context.Background(), "foo", "qux"),
+		"INSERT INTO users DEFAULT VALUES",
+		[]any{},
+		nil,
+	)
+	require.NoError(t, err)
+	require.NoError(t, mock.ExpectationsWereMet())
+	// No rows are returned, so no need to close them.
+
+	mock.ExpectExec("SET foo = 'foo'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("INSERT INTO users DEFAULT VALUES").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	err = drv.Exec(
+		WithVar(context.Background(), "foo", "foo"),
+		"INSERT INTO users DEFAULT VALUES",
+		[]any{},
+		nil,
+	)
+	require.NoError(t, err)
+	require.NoError(t, mock.ExpectationsWereMet())
+	// No rows are returned, so no need to close them.
+}