#722 Autoclose JDBC connections on job timeout

pramen/core/src/main/resources/reference.conf

@@ -61,7 +61,7 @@ pramen {
   # It is not always possible. When a table is initially created, MSCK REPAIR is always used to pick up all partitions.
   # Also ADD PARTTITION is only for Parquet format.
   # This option can be overridden per metatable.
-  hive.prefer.add.partition = false
+  hive.prefer.add.partition = true
 
   # If enabled, the job will wait for the output table to become available before running a job
   # If the number of seconds <=0 the waiting will be infinite

pramen/core/src/main/scala/za/co/absa/pramen/core/runner/task/ThreadClosableRegistry.scala

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package za.co.absa.pramen.core.runner.task
+
+import org.slf4j.LoggerFactory
+
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
+
+object ThreadClosableRegistry {
+  private val log = LoggerFactory.getLogger(this.getClass)
+  private val closeables = new java.util.LinkedList[(Long, AutoCloseable)]
+
+  /**
+    * Registers a closeable resource for the current thread.
+    * The resource will be automatically closed when [[cleanupThread]] is called for this thread.
+    *
+    * @param closeable The AutoCloseable resource to register
+    */
+  def registerCloseable(closeable: AutoCloseable): Unit = synchronized {
+    val threadId = Thread.currentThread().getId
+
+    val iterator = closeables.iterator()
+    var alreadyRegistered = false
+    while (iterator.hasNext && !alreadyRegistered) {
+      alreadyRegistered = iterator.next()._2 == closeable
+    }
+    if (!alreadyRegistered) {
+      closeables.add((threadId, closeable))
+    }
+  }
+
+  /**
+    * Unregisters a closeable resource from the registry.
+    * This removes the resource regardless of which thread it was registered from.
+    *
+    * @param closeable The AutoCloseable resource to unregister
+    */
+  def unregisterCloseable(closeable: AutoCloseable): Unit = synchronized {
+    val iterator = closeables.iterator()
+    while (iterator.hasNext) {
+      val (_, c) = iterator.next()
+      if (c == closeable) {
+        iterator.remove()
+        return
+      }
+    }
+  }
+
+  /**
+    * Closes all registered resources for the specified thread in LIFO (Last-In-First-Out) order.
+    * This method is typically called when a thread times out or completes execution.
+    * Any exceptions during closing are logged but do not prevent other resources from being closed.
+    *
+    * @param threadId The ID of the thread whose resources should be cleaned up
+    */
+  def cleanupThread(threadId: Long): Unit = synchronized {
+    val threadCloseables = closeables.asScala.filter(_._1 == threadId).map(_._2).toList
+    threadCloseables.reverse.foreach { closeable =>
+      try {
+        closeable.close()
+      } catch {
+        case NonFatal(ex) =>
+          log.warn(s"Error closing resource for thread $threadId.", ex)
+      } finally {
+        unregisterCloseable(closeable)
+      }
+    }
+  }
+}

pramen/core/src/main/scala/za/co/absa/pramen/core/utils/ThreadUtils.scala

@@ -16,22 +16,26 @@
 
 package za.co.absa.pramen.core.utils
 
+import org.slf4j.LoggerFactory
 import za.co.absa.pramen.core.exceptions.TimeoutException
+import za.co.absa.pramen.core.runner.task.ThreadClosableRegistry
 import za.co.absa.pramen.core.utils.impl.ThreadWithException
 
 import java.lang.Thread.UncaughtExceptionHandler
 import scala.concurrent.duration.Duration
 
 object ThreadUtils {
+  private val log = LoggerFactory.getLogger(this.getClass)
+
   /**
     * Executes an action with a timeout. If the timeout is breached the task is killed (using Thread.interrupt())
     *
     * If the task times out, an exception is thrown.
     *
     * Any exception is passed to the caller.
     *
-    * @param timeout  The task timeout.
-    * @param action   An action to execute.
+    * @param timeout The task timeout.
+    * @param action  An action to execute.
     */
   @throws[TimeoutException]
   def runWithTimeout(timeout: Duration)(action: => Unit): Unit = {
@@ -54,6 +58,15 @@ object ThreadUtils {
 
     if (thread.isAlive) {
       val stackTrace = thread.getStackTrace
+
+      try {
+        // Execute cleanup BEFORE interrupt - e.g. close the JDBC connection/statement
+        ThreadClosableRegistry.cleanupThread(thread.getId)
+      } catch {
+        case ex: Throwable =>
+          log.warn(s"Exception during timeout cleanup: ${ex.getMessage}")
+      }
+
       thread.interrupt()
 
       val prettyTimeout = TimeUtils.prettyPrintElapsedTimeShort(timeout.toMillis)

pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/QueryExecutorJdbc.scala

@@ -19,6 +19,7 @@ package za.co.absa.pramen.core.utils.hive
 import org.slf4j.LoggerFactory
 import za.co.absa.pramen.core.reader.JdbcUrlSelector
 import za.co.absa.pramen.core.reader.model.JdbcConfig
+import za.co.absa.pramen.core.runner.task.ThreadClosableRegistry
 
 import java.sql._
 import scala.util.control.NonFatal
@@ -66,15 +67,43 @@ class QueryExecutorJdbc(jdbcUrlSelector: JdbcUrlSelector, optimizedExistQuery: B
     executeActionOnConnection { conn =>
       val statement = conn.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)
 
+      val autoCloseStatement: AutoCloseable = new AutoCloseable {
+        val statementClosed = new java.util.concurrent.atomic.AtomicBoolean(false)
+
+        override def close(): Unit = {
+          if (statementClosed.compareAndSet(false, true)) {
+            try {
+              log.info(s"Cancelling SQL statement: $query")
+              statement.cancel()
+            } finally {
+              log.info(s"Closing the SQL statement...")
+              statement.close()
+            }
+          }
+        }
+      }
+
+      ThreadClosableRegistry.registerCloseable(autoCloseStatement)
+
       try {
         statement.execute(query)
       } finally {
-        statement.close()
+        ThreadClosableRegistry.unregisterCloseable(autoCloseStatement)
+        autoCloseStatement.close()
       }
     }
   }
 
-  override def close(): Unit = if (connection != null) connection.close()
+  override def close(): Unit = {
+    if (connection != null) {
+      ThreadClosableRegistry.unregisterCloseable(connection)
+      try {
+        connection.close()
+      } catch {
+        case NonFatal(ex) => log.warn("Failed to close JDBC connection", ex)
+      }
+    }
+  }
 
   private[core] def executeActionOnConnection(action: Connection => Boolean): Boolean = {
     val currentConnection = getConnection(forceReconnect = false)
@@ -97,7 +126,9 @@ class QueryExecutorJdbc(jdbcUrlSelector: JdbcUrlSelector, optimizedExistQuery: B
     if (connection == null || forceReconnect) {
       val (newConnection, url) = jdbcUrlSelector.getWorkingConnection(retries)
       log.info(s"Selected query executor connection: $url")
+      close()
       connection = newConnection
+      ThreadClosableRegistry.registerCloseable(connection)
     }
     connection
   }

pramen/core/src/test/scala/za/co/absa/pramen/core/metastore/MetastoreSuite.scala

@@ -306,7 +306,7 @@ class MetastoreSuite extends AnyWordSpec with SparkTestBase with TextComparisonF
         m.repairOrCreateHiveTable("table_hive_parquet", infoDate, Option(schema), hh, recreate = false)
 
         assert(qe.queries.length == 1)
-        assert(qe.queries.exists(_.contains("REPAIR")))
+        assert(qe.queries.exists(_.contains("ALTER TABLE")))
       }
 
       "do nothing for a delta since it does not need repairing" in {