Add MAX_QUERY_MEMORY_BYTES to PGProperty to use it in QueryExecutorIm…

…pl as limitTupleBytes.
idatahu · Nov 14, 2024 · b5da9da · b5da9da
1 parent b60a01c
commit b5da9da
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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);
   }