Don't use spark to get count query.

AbsaOSS · Aug 13, 2024 · e68ebfb · e68ebfb
1 parent 068a35f
commit e68ebfb
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/TableReaderJdbc.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/reader/TableReaderJdbc.scala
@@ -22,7 +22,7 @@ import org.slf4j.LoggerFactory
 import za.co.absa.pramen.api.Query
 import za.co.absa.pramen.core.config.Keys
 import za.co.absa.pramen.core.reader.model.TableReaderJdbcConfig
-import za.co.absa.pramen.core.utils.{ConfigUtils, JdbcSparkUtils, TimeUtils}
+import za.co.absa.pramen.core.utils.{ConfigUtils, JdbcNativeUtils, JdbcSparkUtils, TimeUtils}
 
 import java.time.format.DateTimeFormatter
 import java.time.{Instant, LocalDate}
@@ -109,8 +109,18 @@ class TableReaderJdbc(jdbcReaderConfig: TableReaderJdbcConfig,
 
   private[core] def getCountForSql(sql: String, infoDateBegin: LocalDate, infoDateEnd: LocalDate): Long = {
     val filteredSql = TableReaderJdbcNative.getFilteredSql(sql, infoDateBegin, infoDateEnd)
-    val countSql = s"SELECT COUNT(*) AS CNT FROM ($filteredSql)"
-    getWithRetry[Long](countSql, isDataQuery = false, jdbcRetries, None)(df => BigDecimal(df.collect()(0)(0).toString).toLong)
+    val countSql = s"SELECT COUNT(*) FROM ($filteredSql)"
+    var count = 0L
+
+    JdbcNativeUtils.withResultSet(jdbcUrlSelector, countSql, jdbcRetries) { rs =>
+      if (!rs.next())
+        throw new IllegalStateException(s"No rows returned by the count query: $countSql")
+      else {
+        count = rs.getLong(1)
+      }
+    }
+
+    count
   }
 
   private[core] def getDataForTable(tableName: String, infoDateBegin: LocalDate, infoDateEnd: LocalDate, columns: Seq[String]): DataFrame = {