GitHub Issue 1187: Auto-disable PostgreSQL JDBC result set caching

api/src/org/labkey/api/data/ContainerManager.java

@@ -1090,17 +1090,18 @@ private static Map<String, Container> getChildrenMap(Container parent)
         {
             try (DbScope.Transaction t = ensureTransaction())
             {
-                List<Container> children = new SqlSelector(CORE.getSchema(),
+                List<GUID> ids = new ArrayList<>();
+                ContainerFactory factory = new ContainerFactory();
+                // ContainerFactory only builds from a ResultSet (fromMap is unsupported), so iterate rows directly
+                new SqlSelector(CORE.getSchema(),
                         "SELECT * FROM " + CORE.getTableInfoContainers() + " WHERE Parent = ? ORDER BY SortOrder, LOWER(Name)",
-                        parent.getId()).getArrayList(Container.class);
-
-                childIds = new ArrayList<>(children.size());
-                for (Container c : children)
-                {
-                    childIds.add(c.getEntityId());
+                        parent.getId()).forEach(rs -> {
+                    Container c = factory.handle(rs);
+                    ids.add(c.getEntityId());
                     _addToCache(c);
-                }
-                childIds = Collections.unmodifiableList(childIds);
+                });
+
+                childIds = Collections.unmodifiableList(ids);
                 CACHE_CHILDREN.put(parent.getEntityId(), childIds);
                 // No database changes to commit, but need to decrement the transaction counter
                 t.commit();

api/src/org/labkey/api/data/DbScope.java

@@ -1202,6 +1202,11 @@ public synchronized boolean release(Connection conn)
 
             return 0 == _refCount;
         }
+
+        public synchronized int getRefCount()
+        {
+            return _refCount;
+        }
     }
 
     /**
@@ -1259,6 +1264,20 @@ private Connection getCurrentConnection(@Nullable Logger log) throws SQLExceptio
         return getConnectionHolder().get(log);
     }
 
+    /**
+     * @return true if this thread already holds the shared, ref-counted thread connection — i.e., some code up the stack
+     * called {@link #getConnection()} and hasn't released it. Reports the existing {@link ConnectionHolder} ref count
+     * that governs {@link ConnectionType#Thread} sharing.
+     * <p>
+     * A caller that borrows the thread connection and temporarily changes its state (e.g. disabling JDBC caching for a
+     * streaming read) must do so only when this returns false, so it is the outermost borrower and can restore the
+     * original state via runOnClose, which {@link ConnectionType#Thread} fires when the last holder releases it.
+     */
+    public boolean isThreadConnectionActive()
+    {
+        return getConnectionHolder().getRefCount() > 0;
+    }
+
     /**
      * Get a fresh connection directly from the pool... not part of the current transaction, not shared with the thread, etc.
      */

api/src/org/labkey/api/data/SqlExecutingSelector.java

@@ -19,6 +19,8 @@
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.labkey.api.cache.CacheManager;
+import org.labkey.api.cache.Throttle;
 import org.labkey.api.data.dialect.SqlDialect;
 import org.labkey.api.data.dialect.SqlDialect.ExecutionPlanType;
 import org.labkey.api.data.dialect.StatementWrapper;
@@ -33,9 +35,12 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static org.labkey.api.util.ExceptionUtil.CALCULATED_COLUMN_SQL_TAG;
 
@@ -46,10 +51,20 @@ public abstract class SqlExecutingSelector<FACTORY extends SqlFactory, SELECTOR
 {
     private static final Logger LOGGER = LogHelper.getLogger(SqlExecutingSelector.class, "Log warnings about SQL exceptions");
 
+    // Warn when this many (or more) rows are pulled into a Java collection; suggests switching to a streaming method
+    private static final int LARGE_RESULT_THRESHOLD = 10_000;
+
+    // Throttles the large-result warning to at most once per day per unique call stack, so a legitimately large (but
+    // expected) load doesn't flood the log. Keyed by a signature of the call stack; see getArrayList().
+    private static final Throttle<LargeResultWarning> LARGE_RESULT_WARNING_THROTTLE = new Throttle<>("SqlSelector large result warnings", 1000, CacheManager.DAY,
+            w -> LOGGER.warn("{} rows loaded into a collection via {}. Consider switching to streaming variants to reduce memory usage. SQL: {}",
+                    w.rowCount, w.selectorClass, w.sql, w.stackTrace));
+
     int _maxRows = Table.ALL_ROWS;
     protected long _offset = Table.NO_OFFSET;
     @Nullable Map<String, Object> _namedParameters = null;
-    private ConnectionFactory _connectionFactory = super::getConnection;
+    private @Nullable ConnectionFactory _connectionFactory = null; // null means "no explicit choice"; see getEffectiveConnectionFactory()
+    private boolean _jdbcCachingExplicitlySet = false;
     private Integer _fetchSize = null; // By default, use the standard fetch size
 
     private @Nullable AsyncQueryRequest<?> _asyncRequest = null;
@@ -79,21 +94,60 @@ public interface ConnectionFactory
     @Override
     public Connection getConnection() throws SQLException
     {
-        return _connectionFactory.get();
+        // Public callers may hold the returned Connection beyond this call, so treat it as "escaping": don't borrow the
+        // thread's shared connection. Internal callers that fully consume and close the ResultSet within a single
+        // selector call use getConnection(true); see getEffectiveConnectionFactory().
+        return getConnection(false);
+    }
+
+    Connection getConnection(boolean selfContained) throws SQLException
+    {
+        return getEffectiveConnectionFactory(selfContained).get();
+    }
+
+    /**
+     * Determines which {@link ConnectionFactory} to use for this query, resolved lazily so the transaction check
+     * reflects execution-time state. An explicit {@link #setJdbcCaching(boolean)} call or a Connection supplied at
+     * construction is honored; otherwise JDBC caching is disabled by default (via the dialect) so the driver won't
+     * buffer the whole ResultSet in memory. The dialect returns null — use the shared Connection with default caching —
+     * when that default is unnecessary or unsafe: in a transaction, non-PostgreSQL, or not a SELECT.
+     * <p>
+     * {@code selfContained} is passed through to {@link SqlDialect#getConnectionFactory}, which documents how it governs
+     * whether the thread's shared connection may be borrowed.
+     */
+    private ConnectionFactory getEffectiveConnectionFactory(boolean selfContained)
+    {
+        // Honor an explicit setJdbcCaching() call (which populated _connectionFactory)...
+        if (_jdbcCachingExplicitlySet)
+            return _connectionFactory;
+
+        // ...or a Connection supplied at construction time (super::getConnection returns the stashed _conn)
+        if (null != _conn)
+            return super::getConnection;
+
+        ConnectionFactory factory = getScope().getSqlDialect().getConnectionFactory(false, selfContained, getScope(),
+            new SQLFragment("SELECT FakeColumn FROM FakeTable") /* SqlExecutingSelector always generates SELECT statements */);
+
+        return null != factory ? factory : super::getConnection;
     }
 
     /**
      * <p>Calling this method with cache=false ensures that the JDBC driver will not cache the produced ResultSet in
      * memory, which is useful when potentially working with very large (e.g., > 100MB) ResultSets. Calling it with
-     * cache=true (the default setting) ensures the JDBC driver's default caching behavior.</p>
+     * cache=true ensures the JDBC driver's default caching behavior.</p>
      *
      * <p>By default, the PostgreSQL JDBC driver caches every ResultSet in its entirety. This can lead to
      * OutOfMemoryErrors when working with very large ResultSets. When the underlying database is PostgreSQL, calling
      * this method with false instructs this SqlExecutingSelector to use an unshared Connection and configure it with
      * special settings that disable the driver caching. The trade-off is that the underlying database query will not
      * use the shared Connection that other code on the thread (up or down the call stack) may be using, making
-     * Connection exhaustion more likely; that's why JDBC caching is on by default. Calling this method is not
-     * compatible with passing in an explicit Connection to the constructor.</p>
+     * Connection exhaustion more likely. Calling this method is not compatible with passing in an explicit Connection to
+     * the constructor.</p>
+     *
+     * <p>When neither this method nor an explicit Connection is supplied, JDBC caching is disabled by default whenever
+     * it's safe (PostgreSQL, no active transaction, SELECT) — see {@link #getEffectiveConnectionFactory(boolean)}. Callers that
+     * require the driver's default caching (e.g. to share the thread's Connection) must opt in by calling this with
+     * cache=true.</p>
      *
      * <p>When the underlying database is not PostgreSQL, calling this method has no effect, other than validating that
      * the stashed Connection is null.</p>
@@ -106,13 +160,63 @@ public SELECTOR setJdbcCaching(boolean cache)
         if (null != _conn)
             throw new IllegalStateException("Calling setJdbcCaching() is not valid when a Connection has already been provided");
 
-        ConnectionFactory factory = getScope().getSqlDialect().getConnectionFactory(cache, getScope(),
+        // Explicitly disabling caching is documented to use an unshared Connection, so this path is never self-contained
+        ConnectionFactory factory = getScope().getSqlDialect().getConnectionFactory(cache, false, getScope(),
             new SQLFragment("SELECT FakeColumn FROM FakeTable") /* SqlExecutingSelector always generates SELECT statements */);
         _connectionFactory = null != factory ? factory : super::getConnection;
+        _jdbcCachingExplicitlySet = true;
 
         return getThis();
     }
 
+    /**
+     * Overridden to warn when a large number of rows is pulled into a Java collection. Loading many rows into memory
+     * (here plus, potentially, in the JDBC driver's buffer) is a common source of OutOfMemoryErrors; callers should
+     * generally prefer a streaming method — {@link #forEach(Class, Selector.ForEachBlock)}, {@link #forEachBatch}, or {@link #uncachedStream} — that
+     * processes rows without materializing them all at once. {@code getArray}, {@code getCollection},
+     * {@code getMapArray}, {@code stream}, and {@code getMapCollection} all delegate here, so they're covered as well.
+     */
+    @Override
+    public @NotNull <E> ArrayList<E> getArrayList(Class<E> clazz)
+    {
+        ArrayList<E> result = super.getArrayList(clazz);
+
+        if (result.size() >= LARGE_RESULT_THRESHOLD)
+        {
+            Throwable stackTrace = new Throwable("Stack trace for large collection load");
+            // Log the parameterized SQL only so bound parameter values stay out of the log
+            SQLFragment sql = getSqlFactory(false).getSql();
+            LARGE_RESULT_WARNING_THROTTLE.execute(new LargeResultWarning(getStackKey(stackTrace), result.size(),
+                    getClass().getSimpleName(), sql == null ? null : sql.getSQL(), stackTrace));
+        }
+
+        return result;
+    }
+
+    // Builds a stable key from a Throwable's stack trace so identical call stacks map to the same throttle entry
+    private static String getStackKey(Throwable t)
+    {
+        return Arrays.stream(t.getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("\n"));
+    }
+
+    // Carries the fields needed to build the large-result warning, but hashes/compares only on the call-stack signature
+    // so the throttle dedupes per unique call stack rather than per (stack + row count + SQL) combination.
+    private record LargeResultWarning(String stackKey, int rowCount, String selectorClass, String sql,
+                                      Throwable stackTrace)
+    {
+        @Override
+        public boolean equals(Object o)
+        {
+            return o instanceof LargeResultWarning w && stackKey.equals(w.stackKey);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return stackKey.hashCode();
+        }
+    }
+
     /**
      * Set a ResultSet fetch size that differs from the default value (1,000 rows on PostgreSQL). This is normally a
      * fine fetch size, but not when dealing with rows containing large TEXT or BYTEA columns.
@@ -396,7 +500,9 @@ public <T> T handleResultSet(ResultSetHandler<T> handler)
                 if (null != _sql)
                 {
                     DbScope scope = getScope();
-                    conn = getConnection();
+                    // _closeResultSet means this factory fully consumes and closes the ResultSet within this call, so
+                    // the connection can be borrowed from the thread and restored rather than dedicated to the caller
+                    conn = getConnection(_closeResultSet);
 
                     try
                     {