Partitioned Postgis - Expose pgjdbc prepareThreshold as a connection param (#10787)

* Adds `prepare_threshold` datastore param and uses it as a pgjdbc driver connection property

Closes #10786

Author: aheyne

docs/user/postgis/usage.rst

@@ -18,6 +18,10 @@ Parameter                               Type     Description
                                                  for more information. Setting this timeout may help prevent abandoned queries from slowing down database operations.
 ``read_access_roles``                   String   A comma-separated list of roles that should be granted read-only access to any new schemas. These roles must already exist in the
                                                  database.
+``prepare_threshold``                   Integer  The pgjdbc ``prepareThreshold`` connection property. By default (``5``), the JDBC driver only requests the binary result wire
+                                                 format after a statement has been server-prepared, so one-shot queries are always decoded as the slower text format. Set to
+                                                 ``-1`` to force the binary format on the first execution. Leave unset to use the pgjdbc default. See the
+                                                 `pgjdbc documentation <https://jdbc.postgresql.org/documentation/use/>`__ for more information.
 ``geomesa.metrics.registry``            String   Specify the type of registry used to publish metrics. Must be one of ``none``,
                                                  ``prometheus``, or ``cloudwatch``. See :ref:`geomesa_metrics` for registry details.
 ``geomesa.metrics.registry.config``     String   Override the default registry config. See :ref:`geomesa_metrics` for configuration details.

geomesa-gt/geomesa-gt-partitioning/src/main/scala/org/locationtech/geomesa/gt/partition/postgis/PartitionedPostgisDataStoreFactory.scala

@@ -31,7 +31,7 @@ import scala.util.control.NonFatal
 class PartitionedPostgisDataStoreFactory extends PostgisNGDataStoreFactory with LazyLogging {
 
   import JDBCDataStoreFactory.{DATABASE, USER}
-  import PartitionedPostgisDataStoreParams.{DbType, IdleInTransactionTimeout, PreparedStatements, ReadAccessRoles}
+  import PartitionedPostgisDataStoreParams.{DbType, IdleInTransactionTimeout, PrepareThreshold, PreparedStatements, ReadAccessRoles}
   import org.locationtech.geomesa.index.geotools.GeoMesaDataStoreFactory.{MetricsRegistryConfigParam, MetricsRegistryParam}
 
   import scala.collection.JavaConverters._
@@ -44,7 +44,7 @@ class PartitionedPostgisDataStoreFactory extends PostgisNGDataStoreFactory with
 
   override protected def setupParameters(parameters: java.util.Map[String, AnyRef]): Unit = {
     super.setupParameters(parameters)
-    Seq(DbType, IdleInTransactionTimeout, PreparedStatements, ReadAccessRoles, MetricsRegistryParam, MetricsRegistryConfigParam)
+    Seq(DbType, IdleInTransactionTimeout, PreparedStatements, ReadAccessRoles, PrepareThreshold, MetricsRegistryParam, MetricsRegistryConfigParam)
         .foreach(p => parameters.put(p.key, p))
   }
 
@@ -180,18 +180,25 @@ class PartitionedPostgisDataStoreFactory extends PostgisNGDataStoreFactory with
   override protected def createSQLDialect(dataStore: JDBCDataStore, params: java.util.Map[String, _]): SQLDialect =
     new PostGISDialect(dataStore)
 
-  private def createConnectionOptions(params: java.util.Map[String, _]): Map[String, String] = {
+  private[postgis] def createConnectionOptions(params: java.util.Map[String, _]): Map[String, String] = {
     val options =
       Seq(IdleInTransactionTimeout)
         .flatMap(p => p.opt(params).map(t => s"-c ${p.key}=${t.millis}"))
 
     logger.debug(s"Connection options: ${options.mkString(" ")}")
 
-    if (options.isEmpty) {
-      Map.empty
-    } else {
-      Map("options" -> options.mkString(" "))
-    }
+    val serverOptions =
+      if (options.isEmpty) { Map.empty[String, String] } else { Map("options" -> options.mkString(" ")) }
+
+    // prepareThreshold is a pgjdbc driver connection property (not a server GUC), so it is added
+    // directly rather than through the '-c <guc>' options string. -1 forces binary result wire
+    // format on the first execution, which one-shot queries never reach with the default of 5.
+    val driverOptions =
+      Option(PrepareThreshold.lookUp(params).asInstanceOf[Integer])
+        .map(v => "prepareThreshold" -> v.toString)
+        .toMap
+
+    serverOptions ++ driverOptions
   }
 }
 

geomesa-gt/geomesa-gt-partitioning/src/main/scala/org/locationtech/geomesa/gt/partition/postgis/PartitionedPostgisDataStoreParams.scala

@@ -52,6 +52,16 @@ object PartitionedPostgisDataStoreParams {
       false
     )
 
+  val PrepareThreshold =
+    new Param(
+      "prepare_threshold",
+      classOf[Integer],
+      "pgjdbc prepareThreshold connection property. -1 forces binary result wire format on the first " +
+          "execution, which one-shot queries never reach with the pgjdbc default of 5. Leave unset to use " +
+          "the pgjdbc default. See https://jdbc.postgresql.org/documentation/use/",
+      false
+    )
+
   // note: need a default string constructor so geotools can create it from the param
   class Timeout(repr: String) {
     private val duration = Duration(repr)

geomesa-gt/geomesa-gt-partitioning/src/test/scala/org/locationtech/geomesa/gt/partition/postgis/PartitionedPostgisDataStoreFactoryTest.scala

@@ -0,0 +1,58 @@
+/***********************************************************************
+ * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, Inc.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Apache License, Version 2.0
+ * which accompanies this distribution and is available at
+ * https://www.apache.org/licenses/LICENSE-2.0
+ ***********************************************************************/
+
+package org.locationtech.geomesa.gt.partition.postgis
+
+import org.junit.runner.RunWith
+import org.specs2.mutable.Specification
+import org.specs2.runner.JUnitRunner
+
+@RunWith(classOf[JUnitRunner])
+class PartitionedPostgisDataStoreFactoryTest extends Specification {
+
+  import scala.collection.JavaConverters._
+
+  "PartitionedPostgisDataStoreFactory" should {
+
+    "advertise the prepare_threshold parameter" in {
+      // guards the setupParameters registration: if the param is dropped from that Seq,
+      // GeoTools silently stops exposing it and this fails
+      val factory = new PartitionedPostgisDataStoreFactory()
+      val keys = factory.getParametersInfo.map(_.key).toSeq
+      keys must contain(PartitionedPostgisDataStoreParams.PrepareThreshold.key)
+      PartitionedPostgisDataStoreParams.PrepareThreshold.key mustEqual "prepare_threshold"
+    }
+
+    "emit prepare_threshold as the pgjdbc driver property prepareThreshold" in {
+      val factory = new PartitionedPostgisDataStoreFactory()
+      val params = Map[String, AnyRef](PartitionedPostgisDataStoreParams.PrepareThreshold.key -> Int.box(-1))
+      val options = factory.createConnectionOptions(params.asJava)
+      // the datastore key is snake_case but the emitted pgjdbc connection property keeps its own spelling
+      options.get("prepareThreshold") must beSome("-1")
+      // it's a driver property, not a server GUC, so it must not leak into the '-c' options string
+      options.get("options") must beNone
+    }
+
+    "not emit prepareThreshold when the parameter is unset" in {
+      val factory = new PartitionedPostgisDataStoreFactory()
+      val options = factory.createConnectionOptions(Map.empty[String, AnyRef].asJava)
+      options.get("prepareThreshold") must beNone
+    }
+
+    "route idle_in_transaction_session_timeout to a server GUC, not a driver property" in {
+      // sibling param: proves the driver-property vs. '-c <guc>' split is behaving as intended
+      val factory = new PartitionedPostgisDataStoreFactory()
+      val params =
+        Map[String, AnyRef](PartitionedPostgisDataStoreParams.IdleInTransactionTimeout.key -> "2 minutes")
+      val options = factory.createConnectionOptions(params.asJava)
+      options.get("options") must beSome
+      options("options") must contain("-c idle_in_transaction_session_timeout=120000")
+      options.get("prepareThreshold") must beNone
+    }
+  }
+}