From b5da9dac61b3eec1390610757c3a3dac6b7137eb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Varga=20Bal=C3=A1zs?= <balazs.varga@idata.hu>
Date: Tue, 29 Oct 2024 16:20:07 +0100
Subject: [PATCH] Add MAX_QUERY_MEMORY_BYTES to PGProperty to use it in
 QueryExecutorImpl as limitTupleBytes.

---
 .../main/java/org/postgresql/PGProperty.java  | 23 ++++++++++++++++++-
 .../postgresql/core/v3/QueryExecutorImpl.java |  4 ++--
 .../postgresql/ds/common/BaseDataSource.java  |  8 +++++++
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/pgjdbc/src/main/java/org/postgresql/PGProperty.java b/pgjdbc/src/main/java/org/postgresql/PGProperty.java
index 53cecf280..ef1785dbc 100644
--- a/pgjdbc/src/main/java/org/postgresql/PGProperty.java
+++ b/pgjdbc/src/main/java/org/postgresql/PGProperty.java
@@ -390,7 +390,12 @@ public enum PGProperty {
       "When connections that are not explicitly closed are garbage collected, log the stacktrace from the opening of the connection to trace the leak source"),
 
   /**
-   * Specifies size of buffer during fetching result set. Can be specified as specified size or
+  * The allowed memory size in bytes, what a query can use (iData made attribute)
+  */
+  MAX_QUERY_MEMORY_BYTES("maxQueryMemoryBytes", Long.toString(Long.MAX_VALUE), "The allowed memory size in bytes, what a query can use"),
+
+  /**
+   * Specifies size of buffer during fetching result th. Can be specified as specified size or
    * percent of heap memory.
    */
   MAX_RESULT_BUFFER(
@@ -1010,6 +1015,22 @@ public int getInt(Properties properties) throws PSQLException {
     }
   }
 
+  /**
+  * Return the long value for this connection parameter in the given {@code Properties}
+  *
+  * @param properties properties to take actual value from
+  * @return evaluated value for this connection parameter converted to long
+  * @throws PSQLException if it cannot be converted to long.
+  */
+  public long getLong(Properties properties) throws PSQLException {
+    String value = get(properties);
+    try {
+      return Long.parseLong(value);
+    } catch (NumberFormatException nfe) {
+      throw new PSQLException(GT.tr("{0} parameter value must be a long but was: {1}", getName(), value), PSQLState.INVALID_PARAMETER_VALUE, nfe);
+    }
+  }
+
   /**
    * Set the boolean value for this connection parameter in the given {@link Properties}.
    *
diff --git a/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java b/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java
index 4643923ab..df3b31f2e 100644
--- a/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java
+++ b/pgjdbc/src/main/java/org/postgresql/core/v3/QueryExecutorImpl.java
@@ -171,6 +171,7 @@ public QueryExecutorImpl(PGStream pgStream,
 
     this.allowEncodingChanges = PGProperty.ALLOW_ENCODING_CHANGES.getBoolean(info);
     this.cleanupSavePoints = PGProperty.CLEANUP_SAVEPOINTS.getBoolean(info);
+    this.limitTupleBytes = PGProperty.MAX_QUERY_MEMORY_BYTES.getLong(info);
     // assignment, argument
     this.replicationProtocol = new V3ReplicationProtocol(this, pgStream);
     readStartupMessages();
@@ -3088,8 +3089,7 @@ public boolean getIntegerDateTimes() {
   private long nextUniqueID = 1;
   private final boolean allowEncodingChanges;
   private final boolean cleanupSavePoints;
-
-  public static volatile long limitTupleBytes = Long.MAX_VALUE;
+  private final long limitTupleBytes;
 
   /**
    * <p>The estimated server response size since we last consumed the input stream from the server, in
diff --git a/pgjdbc/src/main/java/org/postgresql/ds/common/BaseDataSource.java b/pgjdbc/src/main/java/org/postgresql/ds/common/BaseDataSource.java
index ce3df39dc..1fbe8142d 100644
--- a/pgjdbc/src/main/java/org/postgresql/ds/common/BaseDataSource.java
+++ b/pgjdbc/src/main/java/org/postgresql/ds/common/BaseDataSource.java
@@ -1736,6 +1736,14 @@ public void setHideUnprivilegedObjects(boolean hideUnprivileged) {
     PGProperty.HIDE_UNPRIVILEGED_OBJECTS.set(properties, hideUnprivileged);
   }
 
+  public @Nullable String getMaxQueryMemoryBytes() {
+    return PGProperty.MAX_QUERY_MEMORY_BYTES.getOrDefault(properties);
+  }
+
+  public void setMaxQueryMemoryBytes(@Nullable String maxQueryMemoryBytes) {
+    PGProperty.MAX_QUERY_MEMORY_BYTES.set(properties, maxQueryMemoryBytes);
+  }
+
   public @Nullable String getMaxResultBuffer() {
     return PGProperty.MAX_RESULT_BUFFER.getOrDefault(properties);
   }