Repository: AbsaOSS/pramen (language: python)

Issue:
Create a schema query for the Sql Query generator
## Background
Sometimes we want to get query the source JDBC database just to fetch the schema. In this case we need to ensure the query doesn't return any records.

The schema query is the same as data query, just with contradicting WHERE condition.

## Feature
Create a schema query for the Sql Query generator.

## Example [Optional]
--
## Proposed Solution [Optional]
Add either `WHERE 0=1` or `AND 0=1` for the data query to get the schema query.


Discussion / hints:

Fix patch:
diff --git a/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGenerator.scala b/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGenerator.scala
index e3800954..10dd1717 100644
--- a/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGenerator.scala
+++ b/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGenerator.scala
@@ -28,6 +28,13 @@ trait SqlGenerator {
     */
   def getDtable(sql: String): String
 
+  /**
+    * Generates a query that returns no records. It should be used to get the schema for a table.
+    *
+    * If the list of columns passed to this method is empty, all columns will be returned.
+    */
+  def getSchemaQuery(tableName: String, columns: Seq[String]): String
+
   /**
     * Generates a query that returns the record count of a table that does not have the information date field.
     */
diff --git a/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGeneratorBase.scala b/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGeneratorBase.scala
index 8d40f711..d60a9db6 100644
--- a/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGeneratorBase.scala
+++ b/pramen/api/src/main/scala/za/co/absa/pramen/api/sql/SqlGeneratorBase.scala
@@ -62,6 +62,12 @@ abstract class SqlGeneratorBase(sqlConfig: SqlConfig) extends SqlGenerator {
 
   def getOffsetWhereCondition(column: String, condition: String, offset: OffsetValue): String
 
+  override def getSchemaQuery(tableName: String, columns: Seq[String]): String = {
+    val dataQuery = getDataQuery(tableName, columns, None)
+
+    s"$dataQuery WHERE 0=1"
+  }
+
   override def getAliasExpression(expression: String, alias: String): String = {
     s"$expression AS ${escape(alias)}"
   }
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
index 429d6b21..b7134ead 100644
--- 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
@@ -20,12 +20,11 @@ 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.api.offset.{OffsetInfo, OffsetValue}
+import za.co.absa.pramen.api.offset.OffsetValue
 import za.co.absa.pramen.core.config.Keys
-import za.co.absa.pramen.core.reader.model.{OffsetInfoParser, TableReaderJdbcConfig}
+import za.co.absa.pramen.core.reader.model.TableReaderJdbcConfig
 import za.co.absa.pramen.core.utils.{ConfigUtils, JdbcNativeUtils, JdbcSparkUtils, TimeUtils}
 
-import java.time.format.DateTimeFormatter
 import java.time.{Instant, LocalDate}
 import scala.annotation.tailrec
 import scala.util.{Failure, Success, Try}
@@ -214,8 +213,11 @@ class TableReaderJdbc(jdbcReaderConfig: TableReaderJdbcConfig,
     }
 
     if (isDataQuery && jdbcReaderConfig.enableSchemaMetadata) {
-      log.info(s"Reading JDBC metadata from the query: $sql")
-      JdbcSparkUtils.withJdbcMetadata(jdbcReaderConfig.jdbcConfig, sql) { (connection, jdbcMetadata) =>
+      val schemaQuery = tableOpt match {
+        case Some(table) => sqlGen.getSchemaQuery(table, Seq.empty)
+        case _ => JdbcSparkUtils.getSchemaQuery(sql)
+      }
+      JdbcSparkUtils.withJdbcMetadata(jdbcReaderConfig.jdbcConfig, schemaQuery) { (connection, jdbcMetadata) =>
         val schemaWithMetadata = JdbcSparkUtils.addMetadataFromJdbc(df.schema, jdbcMetadata)
         val schemaWithColumnDescriptions = tableOpt match {
           case Some(table) =>
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 5f1112f9..307b0d23 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
@@ -99,7 +99,11 @@ class TableReaderJdbcNative(jdbcReaderConfig: TableReaderJdbcConfig,
     }
 
     if (jdbcReaderConfig.enableSchemaMetadata) {
-      JdbcSparkUtils.withJdbcMetadata(jdbcConfig, sql) { (connection, _) =>
+      val schemaQuery = tableOpt match {
+        case Some(table) => sqlGen.getSchemaQuery(table, Seq.empty)
+        case _ => JdbcSparkUtils.getSchemaQuery(sql)
+      }
+      JdbcSparkUtils.withJdbcMetadata(jdbcConfig, schemaQuery) { (connection, _) =>
         val schemaWithColumnDescriptions = tableOpt match {
           case Some(table) =>
             log.info(s"Reading JDBC metadata descriptions the table: $table")
@@ -129,7 +133,9 @@ class TableReaderJdbcNative(jdbcReaderConfig: TableReaderJdbcConfig,
     }
 
     if (jdbcReaderConfig.enableSchemaMetadata) {
-      JdbcSparkUtils.withJdbcMetadata(jdbcReaderConfig.jdbcConfig, sql) { (connection, _) =>
+      val schemaQuery = sqlGen.getSchemaQuery(tableName, columns)
+
+      JdbcSparkUtils.withJdbcMetadata(jdbcReaderConfig.jdbcConfig, schemaQuery) { (connection, _) =>
         log.info(s"Reading JDBC metadata descriptions the table: $tableName")
         df = spark.createDataFrame(df.rdd,
           JdbcSparkUtils.addColumnDescriptionsFromJdbc(df.schema, sqlGen.unquote(tableName), connection))
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/sql/SqlGeneratorMicrosoft.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/sql/SqlGeneratorMicrosoft.scala
index 92603b70..f952effa 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/sql/SqlGeneratorMicrosoft.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/sql/SqlGeneratorMicrosoft.scala
@@ -44,11 +44,17 @@ class SqlGeneratorMicrosoft(sqlConfig: SqlConfig) extends SqlGenerator {
     }
   }
 
-  def getCountQuery(tableName: String): String = {
+  override def getSchemaQuery(tableName: String, columns: Seq[String]): String = {
+    val dataQuery = getDataQuery(tableName, columns, None)
+
+    s"$dataQuery WHERE 0=1"
+  }
+
+  override def getCountQuery(tableName: String): String = {
     s"SELECT ${getAliasExpression("COUNT(*)", "CNT")} FROM ${escape(tableName)} WITH (NOLOCK)"
   }
 
-  def getCountQuery(tableName: String, infoDateBegin: LocalDate, infoDateEnd: LocalDate): String = {
+  override def getCountQuery(tableName: String, infoDateBegin: LocalDate, infoDateEnd: LocalDate): String = {
     val where = getWhere(infoDateBegin, infoDateEnd)
     s"SELECT ${getAliasExpression("COUNT(*)", "CNT")} FROM ${escape(tableName)} WITH (NOLOCK) WHERE $where"
   }
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcSparkUtils.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcSparkUtils.scala
index 6460ca6c..7d98b397 100644
--- a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcSparkUtils.scala
+++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/JdbcSparkUtils.scala
@@ -179,29 +179,20 @@ object JdbcSparkUtils {
     * metadata of the query.
     *
     * @param jdbcConfig  a JDBC configuration.
-    * @param nativeQuery a SQL query in the dialect native to the database.
+    * @param schemaQuery a SQL query in the dialect native to the database which does not return records.
     * @param action      the action to execute on a connection + resultset metadata.
     */
   def withJdbcMetadata(jdbcConfig: JdbcConfig,
-                       nativeQuery: String)
+                       schemaQuery: String)
                       (action: (Connection, ResultSetMetaData) => Unit): Unit = {
-    val (url, connection) = JdbcNativeUtils.getConnection(jdbcConfig)
+    val (_, connection) = JdbcNativeUtils.getConnection(jdbcConfig)
 
     connection.setAutoCommit(false)
 
-    /** If not filtered out, some JDBC drivers will try to receive all data before closing the result set.
-      * ToDo Fix this properly using SQL generators by adding a generator for schema query. */
-    val q = if (nativeQuery.toLowerCase.contains(" where ")) {
-      nativeQuery + " AND 0=1"
-    } else {
-      nativeQuery + " WHERE 0=1"
-    }
-
-    log.info(s"Successfully connected to JDBC URL: $url")
-    log.info(s"Getting metadata for: $q")
+    log.info(s"Getting metadata for: $schemaQuery")
 
     try {
-      withMetadataResultSet(connection, q) { rs =>
+      withMetadataResultSet(connection, schemaQuery) { rs =>
         action(connection, rs.getMetaData)
       }
     } finally {
@@ -209,6 +200,14 @@ object JdbcSparkUtils {
     }
   }
 
+  private[core] def getSchemaQuery(sql: String): String = {
+    if (sql.toLowerCase.contains(" where ")) {
+      sql + " AND 0=1"
+    } else {
+      sql + " WHERE 0=1"
+    }
+  }
+
   /**
     * Executes a query against a JDBC connection and allows running an action on the result set.
     * Handles the closure of created objects.


Test patch:
diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/mocks/SqlGeneratorDummy.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/mocks/SqlGeneratorDummy.scala
index cf94fc24..e26dce9d 100644
--- a/pramen/core/src/test/scala/za/co/absa/pramen/core/mocks/SqlGeneratorDummy.scala
+++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/mocks/SqlGeneratorDummy.scala
@@ -26,6 +26,8 @@ class SqlGeneratorDummy(sqlConfig: SqlConfig) extends SqlGenerator {
 
   override def getDtable(sql: String): String = null
 
+  override def getSchemaQuery(tableName: String, columns: Seq[String]): String = null
+
   override def getCountQuery(tableName: String): String = null
 
   override def getCountQuery(tableName: String, infoDateBegin: LocalDate, infoDateEnd: LocalDate): String = null
diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorGenericSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorGenericSuite.scala
index f5bd44c0..46ef8b1f 100644
--- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorGenericSuite.scala
+++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorGenericSuite.scala
@@ -45,6 +45,14 @@ class SqlGeneratorGenericSuite extends AnyWordSpec {
   val genNum: SqlGenerator = getSqlGenerator("generic", sqlConfigNumber)
   val genEscaped: SqlGenerator = getSqlGenerator("generic", sqlConfigEscape)
 
+  "generate schema query without list of columns specified" in {
+    assert(genDate.getSchemaQuery("A", Seq.empty) == "SELECT * FROM A WHERE 0=1")
+  }
+
+  "generate schema queries when list of columns is specified" in {
+    assert(genEscaped.getSchemaQuery("A", columns) == "SELECT \"A\", \"D\", \"Column with spaces\" FROM \"A\" WHERE 0=1")
+  }
+
   "generate count queries without date ranges" in {
     assert(genDate.getCountQuery("A") == "SELECT COUNT(*) AS CNT FROM A")
   }
diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorMicrosoftSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorMicrosoftSuite.scala
index 77b50f43..817803ca 100644
--- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorMicrosoftSuite.scala
+++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/sql/SqlGeneratorMicrosoftSuite.scala
@@ -48,6 +48,14 @@ class SqlGeneratorMicrosoftSuite extends AnyWordSpec {
   val genEscaped: SqlGenerator = getSqlGenerator(driver, sqlConfigEscape)
   val genEscaped2: SqlGenerator = getSqlGenerator(driver,  DummySqlConfigFactory.getDummyConfig(infoDateColumn = "[Info date]", identifierQuotingPolicy = QuotingPolicy.Auto))
 
+  "generate schema query without list of columns specified" in {
+    assert(genDate.getSchemaQuery("A", Seq.empty) == "SELECT * FROM A WITH (NOLOCK) WHERE 0=1")
+  }
+
+  "generate schema queries when list of columns is specified" in {
+    assert(genEscaped.getSchemaQuery("A", columns) == "SELECT [A], [D], [Column with spaces] FROM [A] WITH (NOLOCK) WHERE 0=1")
+  }
+
   "generate count queries without date ranges" in {
     assert(genDate.getCountQuery("A") == "SELECT COUNT(*) AS CNT FROM A WITH (NOLOCK)")
   }
diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/utils/JdbcSparkUtilsSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/utils/JdbcSparkUtilsSuite.scala
index b5723126..f388b922 100644
--- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/utils/JdbcSparkUtilsSuite.scala
+++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/utils/JdbcSparkUtilsSuite.scala
@@ -143,6 +143,22 @@ class JdbcSparkUtilsSuite extends AnyWordSpec with BeforeAndAfterAll with SparkT
     }
   }
 
+  "getSchemaQuery" should {
+    "work for queries without where clause" in {
+      val expected = s"SELECT * FROM A WHERE 0=1"
+      val actual = JdbcSparkUtils.getSchemaQuery("SELECT * FROM A")
+
+      assert(actual == expected)
+    }
+
+    "work for queries with where clause" in {
+      val expected = s"SELECT * FROM A WHERE X>0 AND 0=1"
+      val actual = JdbcSparkUtils.getSchemaQuery("SELECT * FROM A WHERE X>0")
+
+      assert(actual == expected)
+    }
+  }
+
   "withMetadataResultSet" should {
     "provide the resultset object for the query" in {
       JdbcSparkUtils.withMetadataResultSet(getConnection, s"SELECT * FROM ${RdbExampleTable.Company.tableName}") { rs =>