From c23f2b013a4fccfb2687db0be52cdc8848ad2d09 Mon Sep 17 00:00:00 2001
From: Fan Yang <fanyang@apecloud.com>
Date: Thu, 19 Dec 2024 17:11:50 +0800
Subject: [PATCH] fix: close (intead of release) the duckdb connection (#301)

---
 adapter/adapter.go             |  6 +++---
 backend/connpool.go            | 12 +++++++++++-
 backend/session.go             |  4 ++--
 pgserver/connection_handler.go |  5 +++--
 test/bats/postgres/cli.bats    | 17 ++++++++++++++++-
 test/bats/postgres/helper.bash |  4 ++--
 6 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/adapter/adapter.go b/adapter/adapter.go
index 51a7fdf0..1bbacd0d 100644
--- a/adapter/adapter.go
+++ b/adapter/adapter.go
@@ -15,7 +15,7 @@ type ConnectionHolder interface {
 	GetCatalogTxn(ctx context.Context, options *stdsql.TxOptions) (*stdsql.Tx, error)
 	TryGetTxn() *stdsql.Tx
 	CloseTxn()
-	CloseBackendConn()
+	CloseConn()
 }
 
 func GetConn(ctx *sql.Context) (*stdsql.Conn, error) {
@@ -26,8 +26,8 @@ func GetCatalogConn(ctx *sql.Context) (*stdsql.Conn, error) {
 	return ctx.Session.(ConnectionHolder).GetCatalogConn(ctx)
 }
 
-func CloseBackendConn(ctx *sql.Context) {
-	ctx.Session.(ConnectionHolder).CloseBackendConn()
+func CloseConn(ctx *sql.Context) {
+	ctx.Session.(ConnectionHolder).CloseConn()
 }
 
 func GetTxn(ctx *sql.Context, options *stdsql.TxOptions) (*stdsql.Tx, error) {
diff --git a/backend/connpool.go b/backend/connpool.go
index db634b87..45232e92 100644
--- a/backend/connpool.go
+++ b/backend/connpool.go
@@ -16,6 +16,7 @@ package backend
 import (
 	"context"
 	stdsql "database/sql"
+	"database/sql/driver"
 	"errors"
 	"fmt"
 	"strings"
@@ -110,7 +111,16 @@ func (p *ConnectionPool) CloseConn(id uint32) error {
 	entry, ok := p.conns.Load(id)
 	if ok {
 		conn := entry.(*stdsql.Conn)
-		if err := conn.Close(); err != nil {
+		if err := conn.Raw(func(driverConn any) error {
+			// When driver.ErrBadConn is returned here,
+			// the connection will not be put back into
+			// the pool and will be closed instead.
+			return driver.ErrBadConn
+		}); err != nil && !errors.Is(err, driver.ErrBadConn) {
+			logrus.WithError(err).Warn("Failed to close connection during Raw function call")
+			return err
+		}
+		if err := conn.Close(); err != nil && !errors.Is(err, stdsql.ErrConnDone) {
 			logrus.WithError(err).Warn("Failed to close connection")
 			return err
 		}
diff --git a/backend/session.go b/backend/session.go
index 47215137..afb9d940 100644
--- a/backend/session.go
+++ b/backend/session.go
@@ -231,8 +231,8 @@ func (sess *Session) CloseTxn() {
 	sess.pool.CloseTxn(sess.ID())
 }
 
-// CloseBackendConn implements adapter.ConnectionHolder.
-func (sess *Session) CloseBackendConn() {
+// CloseConn implements adapter.ConnectionHolder.
+func (sess *Session) CloseConn() {
 	sess.pool.CloseConn(sess.ID())
 }
 
diff --git a/pgserver/connection_handler.go b/pgserver/connection_handler.go
index d26403b6..ec558bc3 100644
--- a/pgserver/connection_handler.go
+++ b/pgserver/connection_handler.go
@@ -117,9 +117,10 @@ func NewConnectionHandler(conn net.Conn, handler mysql.Handler, engine *gms.Engi
 func (h *ConnectionHandler) closeBackendConn() {
 	ctx, err := h.duckHandler.NewContext(context.Background(), h.mysqlConn, "")
 	if err != nil {
-		fmt.Println(err.Error())
+		h.logger.WithError(err).Error("Failed to create context for closing backend connection")
+		return
 	}
-	adapter.CloseBackendConn(ctx)
+	adapter.CloseConn(ctx)
 }
 
 // HandleConnection handles a connection's session, reading messages, executing queries, and sending responses.
diff --git a/test/bats/postgres/cli.bats b/test/bats/postgres/cli.bats
index 861d6da1..7fc0287a 100644
--- a/test/bats/postgres/cli.bats
+++ b/test/bats/postgres/cli.bats
@@ -4,4 +4,19 @@ load helper
 
 @test "cli_show_all_tables" {
     psql_exec "SHOW ALL TABLES;" | grep -q '__sys__'
-}
\ No newline at end of file
+}
+
+# `DISCARD ALL` should clear all temp tables
+@test "discard_all_clears_temp_tables" {
+    # Run all SQL statements in a single session
+    run psql_exec_stdin <<-EOF
+        CREATE TEMP TABLE tt (id int);
+        INSERT INTO tt VALUES (1), (2);
+        SELECT COUNT(*) FROM tt;
+        DISCARD ALL;
+        SELECT COUNT(*) FROM tt;
+EOF
+
+    [ "$status" -ne 0 ]
+    [[ "${output}" == *"Table with name tt does not exist"* ]]
+}
diff --git a/test/bats/postgres/helper.bash b/test/bats/postgres/helper.bash
index ba2bc9f2..7d344b0b 100644
--- a/test/bats/postgres/helper.bash
+++ b/test/bats/postgres/helper.bash
@@ -8,9 +8,9 @@ PG_USER=${PG_USER:-"postgres"}
 psql_exec() {
     local query="$1"
     shift
-    psql -h "$PG_HOST" -U "$PG_USER" -F ',' --no-align --field-separator ',' --pset footer=off "$@" -c "$query"
+    psql -h "$PG_HOST" -U "$PG_USER" -F ',' --no-align --field-separator ',' -t --pset footer=off -v "ON_ERROR_STOP=1" "$@" -c "$query"
 }
 
 psql_exec_stdin() {
-    psql -h "$PG_HOST" -U "$PG_USER" -F ',' --no-align --field-separator ',' --pset footer=off "$@"
+    psql -h "$PG_HOST" -U "$PG_USER" -F ',' --no-align --field-separator ',' -t --pset footer=off -v "ON_ERROR_STOP=1" "$@"
 }
\ No newline at end of file