From f9dd164cc0cdd7b3e8ad3fd9b7ace874924a6433 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Wa=C5=9B?= <jan.was@starburstdata.com>
Date: Sun, 6 Oct 2024 13:24:43 +0200
Subject: [PATCH] Support EXECUTE IMMEDIATE

Use EXECUTE IMMEDIATE sent in the HTTP request body, instead of putting
the query text in HTTP headers. This should allow sending large query
text. It can be enabled by setting the `explicitPrepare` option to
true in the connection string.
---
 trino/etc/config.properties |  2 ++
 trino/etc/jvm.config        |  2 +-
 trino/integration_test.go   | 22 ++++++++++++++++++++++
 trino/trino.go              | 19 +++++++++++++++++--
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/trino/etc/config.properties b/trino/etc/config.properties
index d4d09aa..9b696f1 100644
--- a/trino/etc/config.properties
+++ b/trino/etc/config.properties
@@ -12,3 +12,5 @@ http-server.https.port=8443
 http-server.authentication.allow-insecure-over-http=true
 http-server.https.keystore.path=/etc/trino/secrets/certificate_with_key.pem
 internal-communication.shared-secret=gotrino
+
+query.max-length=2000000
diff --git a/trino/etc/jvm.config b/trino/etc/jvm.config
index bb4dca9..cdf7bde 100644
--- a/trino/etc/jvm.config
+++ b/trino/etc/jvm.config
@@ -1,4 +1,4 @@
--Xmx1G
+-Xmx4G
 -XX:+UseG1GC
 -XX:G1HeapRegionSize=32M
 -XX:+UseGCOverheadLimit
diff --git a/trino/integration_test.go b/trino/integration_test.go
index d78c1ae..10dc417 100644
--- a/trino/integration_test.go
+++ b/trino/integration_test.go
@@ -1112,3 +1112,25 @@ func TestIntegrationDayToHourIntervalMilliPrecision(t *testing.T) {
 		})
 	}
 }
+
+func TestIntegrationLargeQuery(t *testing.T) {
+	dsn := *integrationServerFlag
+	dsn += "?explicitPrepare=false"
+	db := integrationOpen(t, dsn)
+	defer db.Close()
+	rows, err := db.Query("SELECT ?"+strings.Repeat(", 'b', 'c', 'd', 'e', 'f'", 40000), 42)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rows.Close()
+	count := 0
+	for rows.Next() {
+		count++
+	}
+	if rows.Err() != nil {
+		t.Fatal(err)
+	}
+	if count != 1 {
+		t.Fatal("not enough rows returned:", count)
+	}
+}
diff --git a/trino/trino.go b/trino/trino.go
index 5204717..528381b 100644
--- a/trino/trino.go
+++ b/trino/trino.go
@@ -141,6 +141,7 @@ const (
 	sslCertPathConfig               = "SSLCertPath"
 	sslCertConfig                   = "SSLCert"
 	accessTokenConfig               = "accessToken"
+	explicitPrepareConfig           = "explicitPrepare"
 )
 
 var (
@@ -282,6 +283,7 @@ type Conn struct {
 	kerberosRemoteServiceName string
 	progressUpdater           ProgressUpdater
 	progressUpdaterPeriod     queryProgressCallbackPeriod
+	useExplicitPrepare        bool
 }
 
 var (
@@ -298,6 +300,10 @@ func newConn(dsn string) (*Conn, error) {
 	query := serverURL.Query()
 
 	kerberosEnabled, _ := strconv.ParseBool(query.Get(kerberosEnabledConfig))
+	useExplicitPrepare := true
+	if query.Get(explicitPrepareConfig) != "" {
+		useExplicitPrepare, _ = strconv.ParseBool(query.Get(explicitPrepareConfig))
+	}
 
 	var kerberosClient *client.Client
 
@@ -356,6 +362,7 @@ func newConn(dsn string) (*Conn, error) {
 		kerberosClient:            kerberosClient,
 		kerberosEnabled:           kerberosEnabled,
 		kerberosRemoteServiceName: query.Get(kerberosRemoteServiceNameConfig),
+		useExplicitPrepare:        useExplicitPrepare,
 	}
 
 	var user string
@@ -867,7 +874,7 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 
 				hs.Add(arg.Name, headerValue)
 			} else {
-				if hs.Get(preparedStatementHeader) == "" {
+				if st.conn.useExplicitPrepare && hs.Get(preparedStatementHeader) == "" {
 					for _, v := range st.conn.httpHeaders.Values(preparedStatementHeader) {
 						hs.Add(preparedStatementHeader, v)
 					}
@@ -880,7 +887,11 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 			return nil, ErrInvalidProgressCallbackHeader
 		}
 		if len(ss) > 0 {
-			query = "EXECUTE " + preparedStatementName + " USING " + strings.Join(ss, ", ")
+			if st.conn.useExplicitPrepare {
+				query = "EXECUTE " + preparedStatementName + " USING " + strings.Join(ss, ", ")
+			} else {
+				query = "EXECUTE IMMEDIATE " + formatStringLiteral(st.query) + " USING " + strings.Join(ss, ", ")
+			}
 		}
 	}
 
@@ -1028,6 +1039,10 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 	return &sr, handleResponseError(resp.StatusCode, sr.Error)
 }
 
+func formatStringLiteral(query string) string {
+	return "'" + strings.ReplaceAll(query, "'", "''") + "'"
+}
+
 type driverRows struct {
 	ctx     context.Context
 	stmt    *driverStmt