Repository: AbsaOSS/pramen (language: python)

Issue:
Large PostgreSQL queries result in Out of Memory
## Describe the bug
This happens since PostgreSql driver loads all query results in memory.

## Code and/or configuration snippet that caused the issue

```hocon
    jdbc = {
      driver = "org.postgresql.Driver"
      url = "jdbc:postgresql://..."
    }
```

options like
```hocon
option {
    adaptiveFetch = false
    defaultRowFetchSize = 100
    adaptiveFetchMaximum = 200
  }
}
```

don't help.

Fetch sizes are ignored.

## Expected behavior
PostgreSql queries should not load full results into mmory, but load data in chunks.

## Context
According to this:
https://stackoverflow.com/questions/1504341/jdbc-large-postgresql-query-give-out-of-memory

The forward-only cursor shoud help:
```scala
connection.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)
```



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index e947f26b1..16b10385e 100644
--- a/README.md
+++ b/README.md
@@ -501,6 +501,9 @@ is determined by the pipeline configuration.
       
       # Any option passed as 'option.' will be passed to the JDBC driver. Example:
       #option.database = "test_db"
+      
+      # (Optional) Autocommit, false by default. Used only when 'use.jdbc.native = true'
+      #autocommit = false
     }
 
     # Any option passed as '.option' here will be passed to the Spark reader as options. For example,
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/JdbcUrlSelectorImpl.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/JdbcUrlSelectorImpl.scala
index 37a64cc4c..16fe0559e 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/JdbcUrlSelectorImpl.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/JdbcUrlSelectorImpl.scala
@@ -54,19 +54,24 @@ class JdbcUrlSelectorImpl(val jdbcConfig: JdbcConfig) extends JdbcUrlSelector{
 
   def logConnectionSettings(): Unit = {
     log.info(s"JDBC Configuration:")
-    log.info(s"Driver:            ${jdbcConfig.driver}")
-    log.info(s"URL:               $currentUrl")
-    jdbcConfig.database.foreach(d => log.info(s"Database:          $d"))
-    jdbcConfig.user.foreach(user => log.info(s"User:              $user"))
-    jdbcConfig.password.foreach(_ =>log.info(s"Password:          [redacted]"))
+    log.info(s"Driver:             ${jdbcConfig.driver}")
+    log.info(s"URL:                $currentUrl")
 
-    jdbcConfig.retries.foreach(n => log.info(s"Retry attempts:    $n"))
-    jdbcConfig.connectionTimeoutSeconds.foreach(n => log.info(s"Timeout:           $n seconds"))
+    jdbcConfig.database.foreach(d  => log.info(s"Database:           $d"))
+    jdbcConfig.user.foreach(user   => log.info(s"User:               $user"))
+    jdbcConfig.password.foreach(_  => log.info(s"Password:           [redacted]"))
+    jdbcConfig.fetchSize.foreach(n => log.info(s"Fetch size:         $n"))
+
+    log.info(s"Auto commit:        ${jdbcConfig.autoCommit}")
+    log.info(s"Sanitize date/time: ${jdbcConfig.sanitizeDateTime}")
+
+    jdbcConfig.retries.foreach(n =>   log.info(s"Retry attempts:     $n"))
+    jdbcConfig.connectionTimeoutSeconds.foreach(n => log.info(s"Timeout:            $n seconds"))
 
     if (numberOfUrls > 1) {
       log.info("URL redundancy configuration:")
-      jdbcConfig.primaryUrl.foreach(s => log.info(s"Primary URL:       $s"))
-      jdbcConfig.fallbackUrls.zipWithIndex.foreach { case (s, i) => log.info(s"Pool URL ${i + 1}:        $s") }
+      jdbcConfig.primaryUrl.foreach(s => log.info(s"Primary URL:        $s"))
+      jdbcConfig.fallbackUrls.zipWithIndex.foreach { case (s, i) => log.info(s"Pool URL ${i + 1}:         $s") }
     }
     if (jdbcConfig.extraOptions.nonEmpty) {
       log.info("JDBC extra properties:")
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/TableReaderJdbcNative.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/TableReaderJdbcNative.scala
index 5ec0e6a7c..8923b77ca 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/TableReaderJdbcNative.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/TableReaderJdbcNative.scala
@@ -20,7 +20,7 @@ import com.typesafe.config.Config
 import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.slf4j.LoggerFactory
 import za.co.absa.pramen.api.Query
-import za.co.absa.pramen.core.reader.model.TableReaderJdbcConfig
+import za.co.absa.pramen.core.reader.model.{JdbcConfig, TableReaderJdbcConfig}
 import za.co.absa.pramen.core.utils.{JdbcNativeUtils, StringUtils, TimeUtils}
 
 import java.time.format.DateTimeFormatter
@@ -35,13 +35,14 @@ class TableReaderJdbcNative(jdbcReaderConfig: TableReaderJdbcConfig,
   private val log = LoggerFactory.getLogger(this.getClass)
 
   private val jdbcConfig = jdbcReaderConfig.jdbcConfig
-  private val infoDateFormatPattern = jdbcReaderConfig.infoDateFormat
-  private val infoDateFormatter = DateTimeFormatter.ofPattern(infoDateFormatPattern)
+
   private val url = jdbcUrlSelector.getWorkingUrl(jdbcConfig.retries.getOrElse(jdbcUrlSelector.getNumberOfUrls))
 
   logConfiguration()
 
-  private[core] def getJdbcReaderConfig: TableReaderJdbcConfig = jdbcReaderConfig
+  private[core] def getJdbcReaderConfig: TableReaderJdbcConfig = {
+    jdbcReaderConfig.copy(jdbcConfig = jdbcConfig)
+  }
 
   override def getRecordCount(query: Query, infoDateBegin: LocalDate, infoDateEnd: LocalDate): Long = {
     val start = Instant.now()
@@ -92,15 +93,28 @@ class TableReaderJdbcNative(jdbcReaderConfig: TableReaderJdbcConfig,
 }
 
 object TableReaderJdbcNative {
+  val FETCH_SIZE_KEY = "option.fetchsize"
+
   def apply(conf: Config,
             parent: String = "")
            (implicit spark: SparkSession): TableReaderJdbcNative = {
-    val tableReaderJdbc = TableReaderJdbcConfig.load(conf, parent)
+    val tableReaderJdbcOrig = TableReaderJdbcConfig.load(conf, parent)
+    val jdbcConfig = getJdbcConfig(tableReaderJdbcOrig, conf)
+    val tableReaderJdbc = tableReaderJdbcOrig.copy(jdbcConfig = jdbcConfig)
     val urlSelector = JdbcUrlSelector(tableReaderJdbc.jdbcConfig)
 
     new TableReaderJdbcNative(tableReaderJdbc, urlSelector, conf)
   }
 
+  private[core] def getJdbcConfig(tableReaderJdbcConfig: TableReaderJdbcConfig, conf: Config): JdbcConfig = {
+    val originConfig = tableReaderJdbcConfig.jdbcConfig
+    if (conf.hasPath(FETCH_SIZE_KEY)) {
+      originConfig.copy(fetchSize = Option(conf.getInt(FETCH_SIZE_KEY)))
+    } else {
+      originConfig
+    }
+  }
+
   def getFilteredSql(sqlExpression: String, infoDateBegin: LocalDate, infoDateEnd: LocalDate): String = {
     val f1 = StringUtils.replaceFormattedDateExpression(sqlExpression, "dateFrom", infoDateBegin)
     val f2 = StringUtils.replaceFormattedDateExpression(f1, "dateTo", infoDateEnd)
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/model/JdbcConfig.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/model/JdbcConfig.scala
index 024b76cf8..584f58553 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/model/JdbcConfig.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/model/JdbcConfig.scala
@@ -29,6 +29,8 @@ case class JdbcConfig(
                        user: Option[String] = None,
                        password: Option[String] = None,
                        retries: Option[Int] = None,
+                       fetchSize: Option[Int] = None,
+                       autoCommit: Boolean = false,
                        connectionTimeoutSeconds: Option[Int] = None,
                        sanitizeDateTime: Boolean = true,
                        incorrectDecimalsAsString: Boolean = false,
@@ -46,6 +48,8 @@ object JdbcConfig {
   val JDBC_USER = "jdbc.user"
   val JDBC_PASSWORD = "jdbc.password"
   val JDBC_RETRIES = "jdbc.retries"
+  val JDBC_FETCH_SIZE = "jdbc.fetchsize"
+  val JDBC_AUTOCOMMIT = "jdbc.autocommit"
   val JDBC_CONNECTION_TIMEOUT = "jdbc.connection.timeout"
   val JDBC_SANITIZE_DATETIME = "jdbc.sanitize.datetime"
   val JDBC_INCORRECT_PRECISION_AS_STRING = "jdbc.incorrect.precision.as.string"
@@ -77,6 +81,8 @@ object JdbcConfig {
       user = ConfigUtils.getOptionString(conf, JDBC_USER),
       password = ConfigUtils.getOptionString(conf, JDBC_PASSWORD),
       retries = ConfigUtils.getOptionInt(conf, JDBC_RETRIES),
+      fetchSize = ConfigUtils.getOptionInt(conf, JDBC_FETCH_SIZE),
+      autoCommit = ConfigUtils.getOptionBoolean(conf, JDBC_AUTOCOMMIT).getOrElse(false),
       connectionTimeoutSeconds = ConfigUtils.getOptionInt(conf, JDBC_CONNECTION_TIMEOUT),
       sanitizeDateTime = ConfigUtils.getOptionBoolean(conf, JDBC_SANITIZE_DATETIME).getOrElse(true),
       incorrectDecimalsAsString = ConfigUtils.getOptionBoolean(conf, JDBC_INCORRECT_PRECISION_AS_STRING).getOrElse(false),
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/source/JdbcSource.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/source/JdbcSource.scala
index 7a6cdf89f..545219308 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/source/JdbcSource.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/source/JdbcSource.scala
@@ -49,12 +49,14 @@ class JdbcSource(sourceConfig: Config,
   }
 
   private[core] def getReader(query: Query, isCountQuery: Boolean): TableReader = {
-    val urlSelector = JdbcUrlSelector(jdbcReaderConfig.jdbcConfig)
+    val jdbcConfig = TableReaderJdbcNative.getJdbcConfig(jdbcReaderConfig, sourceConfig)
+    val jdbcReaderConfigNative = jdbcReaderConfig.copy(jdbcConfig = jdbcConfig)
+    val urlSelector = JdbcUrlSelector(jdbcConfig)
 
     query match {
       case Query.Table(dbTable) if jdbcReaderConfig.useJdbcNative && !isCountQuery =>
         log.info(s"Using TableReaderJdbcNative to read the table: $dbTable")
-        new TableReaderJdbcNative(jdbcReaderConfig, urlSelector, sourceConfig)
+        new TableReaderJdbcNative(jdbcReaderConfigNative, urlSelector, sourceConfig)
       case Query.Table(dbTable) =>
         log.info(s"Using TableReaderJdbc to read the table: $dbTable")
         new TableReaderJdbc(jdbcReaderConfig, urlSelector, sourceConfig)
@@ -63,7 +65,7 @@ class JdbcSource(sourceConfig: Config,
         new TableReaderJdbc(jdbcReaderConfig, urlSelector, sourceConfig)
       case Query.Sql(sql)  =>
         log.info(s"Using TableReaderJdbcNative to read the query: $sql")
-        new TableReaderJdbcNative(jdbcReaderConfig, urlSelector, sourceConfig)
+        new TableReaderJdbcNative(jdbcReaderConfigNative, urlSelector, sourceConfig)
       case q =>
         throw new IllegalArgumentException(s"Unexpected '${q.name}' spec for the JDBC reader. Only 'table' or 'sql' are supported. Config path: $sourceConfigParentPath")
     }
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcNativeUtils.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcNativeUtils.scala
index bd0adf06d..2221aa1d5 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcNativeUtils.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcNativeUtils.scala
@@ -127,7 +127,16 @@ object JdbcNativeUtils {
     val connection = getJdbcConnection(jdbcConfig, url)
 
     val statement = try {
-      connection.createStatement(ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY)
+      connection.setAutoCommit(jdbcConfig.autoCommit)
+
+      if (jdbcConfig.driver == "org.postgresql.Driver")
+        // Special handling of PostgreSQL driver that loads.
+        // PostgreSQL loads full query results info memory ignoring fetch sizes and other driver options by default.
+        // This is the only cursor that makes it respect fetch sizes.
+        // If we encounter more exceptions, we can improve this 'patch' into a properly handling method.
+        connection.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)
+      else
+        connection.createStatement(ResultSet.TYPE_SCROLL_INSENSITIVE, ResultSet.CONCUR_READ_ONLY)
     } catch {
       case _: java.sql.SQLException =>
         // Fallback with more permissible result type.
@@ -137,6 +146,10 @@ object JdbcNativeUtils {
         throw ex
     }
 
+    jdbcConfig.fetchSize.foreach(n =>
+      statement.setFetchSize(n)
+    )
+
     statement.executeQuery(query)
   }
 
diff --git a/pramen/project/Versions.scala b/pramen/project/Versions.scala
index 8de3124a7..626c3d9ba 100644
--- a/pramen/project/Versions.scala
+++ b/pramen/project/Versions.scala
@@ -22,7 +22,7 @@ object Versions {
   val defaultSparkVersionForScala213 = "3.4.2"
 
   val typesafeConfigVersion = "1.4.0"
-  val postgreSqlDriverVersion = "42.3.8"
+  val postgreSqlDriverVersion = "42.7.3"
   val msSqlDriverVersion = "1.3.1"
   val mongoDbScalaDriverVersion = "2.7.0"
   val hsqlDbVersion = "2.7.1"


Test patch:
diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/reader/TableReaderJdbcNativeSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/reader/TableReaderJdbcNativeSuite.scala
index 275d9ffa2..165681494 100644
--- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/reader/TableReaderJdbcNativeSuite.scala
+++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/reader/TableReaderJdbcNativeSuite.scala
@@ -37,6 +37,7 @@ class TableReaderJdbcNativeSuite extends AnyWordSpec with RelationalDbFixture wi
        |    connection.string = "$url"
        |    user = "$user"
        |    password = "$password"
+       |    autocommit = true
        |  }
        |
        |  has.information.date.column = true
@@ -44,6 +45,7 @@ class TableReaderJdbcNativeSuite extends AnyWordSpec with RelationalDbFixture wi
        |  information.date.column = "FOUNDED"
        |  information.date.type = "date"
        |  information.date.format = "YYYY-MM-dd"
+       |  option.fetchsize = 1000
        |
        |}
        |reader_legacy {
@@ -102,6 +104,8 @@ class TableReaderJdbcNativeSuite extends AnyWordSpec with RelationalDbFixture wi
       val reader = getReader
       assert(reader != null)
       assert(reader.getJdbcReaderConfig.infoDateFormat == "YYYY-MM-dd")
+      assert(reader.getJdbcReaderConfig.jdbcConfig.autoCommit)
+      assert(reader.getJdbcReaderConfig.jdbcConfig.fetchSize.get == 1000)
     }
 
     "work with legacy config" in {
@@ -114,6 +118,8 @@ class TableReaderJdbcNativeSuite extends AnyWordSpec with RelationalDbFixture wi
       val reader = TableReaderJdbcNative(conf.getConfig("reader_minimal"), "reader_minimal")
       assert(reader.getJdbcReaderConfig.infoDateFormat == "yyyy-MM-dd")
       assert(reader.getJdbcReaderConfig.jdbcConfig.sanitizeDateTime)
+      assert(!reader.getJdbcReaderConfig.jdbcConfig.autoCommit)
+      assert(reader.getJdbcReaderConfig.jdbcConfig.fetchSize.isEmpty)
     }
 
     "throw an exception if config is missing" in {