Fixes the schema generation for JDBC integration (#470)

zaleslaw · web-flow · commit 835de4b46c54 · 2023-10-16T12:31:21.000+02:00
* Adjusted SQL table/schema reading to accommodate table name

The update introduces 'tableName' in the 'JdbcOptions' annotation class and modifies conditions in 'GenerateDataSchemaTask.kt' and 'DataSchemaGenerator.kt' to handle empty 'tableName'. It also includes driver initialisation for Mariadb. This alteration allows greater flexibility when specifying a SQL source, accommodating both SQL queries and direct table names.

* Correct the DbUrl extraction and add JDBC class to DbType
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/ImportDataSchema.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/ImportDataSchema.kt
@@ -52,6 +52,7 @@ public annotation class CsvOptions(
 public annotation class JdbcOptions(
     public val user: String = "", // TODO: I'm not sure about the default parameters
     public val password: String = "", // TODO: I'm not sure about the default parameters)
+    public val tableName: String = "",
     public val sqlQuery: String = ""
 )
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/ImportDataSchema.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/ImportDataSchema.kt
@@ -52,6 +52,7 @@ public annotation class CsvOptions(
 public annotation class JdbcOptions(
     public val user: String = "", // TODO: I'm not sure about the default parameters
     public val password: String = "", // TODO: I'm not sure about the default parameters)
+    public val tableName: String = "",
     public val sqlQuery: String = ""
 )
 
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/DbType.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/DbType.kt
@@ -11,6 +11,17 @@ import org.jetbrains.kotlinx.dataframe.io.TableMetadata
  * @property [dbTypeInJdbcUrl] The name of the database as specified in the JDBC URL.
  */
 public abstract class DbType(public val dbTypeInJdbcUrl: String) {
+
+
+    /**
+     * Represents the JDBC driver class name for a given database type.
+     *
+     * NOTE: It's important for usage in dataframe-gradle-plugin for force class loading.
+     *
+     * @return The JDBC driver class name as a [String].
+     */
+    public abstract val driverClassName: String
+
     /**
      * Converts the data from the given [ResultSet] into the specified [TableColumnMetadata] type.
      *
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/H2.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/H2.kt
@@ -18,6 +18,9 @@ import kotlin.reflect.typeOf
  * NOTE: All date and timestamp related types are converted to String to avoid java.sql.* types.
  */
 public object H2 : DbType("h2") {
+    override val driverClassName: String
+        get() = "org.h2.Driver"
+
     override fun convertDataFromResultSet(rs: ResultSet, tableColumnMetadata: TableColumnMetadata): Any? {
         val name = tableColumnMetadata.name
         return when (tableColumnMetadata.sqlTypeName) {
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/MariaDb.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/MariaDb.kt
@@ -13,6 +13,9 @@ import kotlin.reflect.typeOf
  * and to generate the corresponding column schema.
  */
 public object MariaDb : DbType("mariadb") {
+    override val driverClassName: String
+        get() = "org.mariadb.jdbc.Driver"
+
     override fun convertDataFromResultSet(rs: ResultSet, tableColumnMetadata: TableColumnMetadata): Any? {
         val name = tableColumnMetadata.name
         return when (tableColumnMetadata.sqlTypeName) {
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/MySql.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/MySql.kt
@@ -16,6 +16,9 @@ import kotlin.reflect.typeOf
  * and to generate the corresponding column schema.
  */
 public object MySql : DbType("mysql") {
+    override val driverClassName: String
+        get() = "com.mysql.jdbc.Driver"
+
     override fun convertDataFromResultSet(rs: ResultSet, tableColumnMetadata: TableColumnMetadata): Any? {
         val name = tableColumnMetadata.name
         return when (tableColumnMetadata.sqlTypeName) {
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/PostgreSql.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/PostgreSql.kt
@@ -16,6 +16,9 @@ import kotlin.reflect.typeOf
  * and to generate the corresponding column schema.
  */
 public object PostgreSql : DbType("postgresql") {
+    override val driverClassName: String
+        get() = "org.postgresql.Driver"
+
     override fun convertDataFromResultSet(rs: ResultSet, tableColumnMetadata: TableColumnMetadata): Any? {
         val name = tableColumnMetadata.name
         return when (tableColumnMetadata.sqlTypeName) {
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/Sqlite.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/Sqlite.kt
@@ -13,6 +13,9 @@ import kotlin.reflect.typeOf
  * and to generate the corresponding column schema.
  */
 public object Sqlite : DbType("sqlite") {
+    override val driverClassName: String
+        get() = "org.sqlite.JDBC"
+
     override fun convertDataFromResultSet(rs: ResultSet, tableColumnMetadata: TableColumnMetadata): Any? {
         val name = tableColumnMetadata.name
         return when (tableColumnMetadata.sqlTypeName) {
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/util.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/db/util.kt
@@ -9,7 +9,7 @@ import java.sql.SQLException
  * @return the corresponding [DbType].
  * @throws RuntimeException if the url is null.
  */
-public fun extractDBTypeFromURL(url: String?): DbType {
+public fun extractDBTypeFromUrl(url: String?): DbType {
     if (url != null) {
         return when {
             H2.dbTypeInJdbcUrl in url -> H2
@@ -24,3 +24,14 @@ public fun extractDBTypeFromURL(url: String?): DbType {
         throw SQLException("Database URL could not be null. The existing value is $url")
     }
 }
+
+/**
+ * Retrieves the driver class name from the given JDBC URL.
+ *
+ * @param [url] The JDBC URL to extract the driver class name from.
+ * @return The driver class name as a [String].
+ */
+public fun driverClassNameFromUrl(url: String): String {
+    val dbType = extractDBTypeFromUrl(url)
+    return dbType.driverClassName
+}
diff --git a/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readJdbc.kt b/dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readJdbc.kt
@@ -11,9 +11,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.impl.schema.DataFrameSchemaImpl
 import org.jetbrains.kotlinx.dataframe.io.db.DbType
-import org.jetbrains.kotlinx.dataframe.io.db.H2
-import org.jetbrains.kotlinx.dataframe.io.db.Sqlite
-import org.jetbrains.kotlinx.dataframe.io.db.extractDBTypeFromURL
+import org.jetbrains.kotlinx.dataframe.io.db.extractDBTypeFromUrl
 import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
 
 private val logger = KotlinLogging.logger {}
@@ -115,7 +113,7 @@ public fun DataFrame.Companion.readSqlTable(connection: Connection, tableName: S
     if (limit > 0) preparedQuery += " LIMIT $limit"
 
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     connection.createStatement().use { st ->
         logger.debug { "Connection with url:${url} is established successfully." }
@@ -182,7 +180,7 @@ public fun DataFrame.Companion.readSqlQuery(connection: Connection, sqlQuery: St
  */
 public fun DataFrame.Companion.readSqlQuery(connection: Connection, sqlQuery: String, limit: Int): AnyFrame {
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     var internalSqlQuery = sqlQuery
     if (limit > 0) internalSqlQuery += " LIMIT $limit"
@@ -247,7 +245,7 @@ public fun DataFrame.Companion.readResultSet(resultSet: ResultSet, connection: C
  */
 public fun DataFrame.Companion.readResultSet(resultSet: ResultSet, connection: Connection, limit: Int): AnyFrame {
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     return readResultSet(resultSet, dbType, limit)
 }
@@ -300,7 +298,7 @@ public fun DataFrame.Companion.readAllSqlTables(connection: Connection): List<An
 public fun DataFrame.Companion.readAllSqlTables(connection: Connection, limit: Int): List<AnyFrame> {
     val metaData = connection.metaData
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     // exclude a system and other tables without data, but it looks like it supported badly for many databases
     val tables = metaData.getTables(null, null, null, arrayOf("TABLE"))
@@ -348,7 +346,7 @@ public fun DataFrame.Companion.getSchemaForSqlTable(
     tableName: String
 ): DataFrameSchema {
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     connection.createStatement().use {
         logger.debug { "Connection with url:${connection.metaData.url} is established successfully." }
@@ -383,7 +381,7 @@ public fun DataFrame.Companion.getSchemaForSqlQuery(dbConfig: DatabaseConfigurat
  */
 public fun DataFrame.Companion.getSchemaForSqlQuery(connection: Connection, sqlQuery: String): DataFrameSchema {
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     connection.createStatement().use { st ->
         st.executeQuery(sqlQuery).use { rs ->
@@ -419,7 +417,7 @@ public fun DataFrame.Companion.getSchemaForResultSet(resultSet: ResultSet, dbTyp
  */
 public fun DataFrame.Companion.getSchemaForResultSet(resultSet: ResultSet, connection: Connection): DataFrameSchema {
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     val tableColumns = getTableColumnsMetadata(resultSet)
     return buildSchemaByTableColumns(tableColumns, dbType)
@@ -446,7 +444,7 @@ public fun DataFrame.Companion.getSchemaForAllSqlTables(dbConfig: DatabaseConfig
 public fun DataFrame.Companion.getSchemaForAllSqlTables(connection: Connection): List<DataFrameSchema> {
     val metaData = connection.metaData
     val url = connection.metaData.url
-    val dbType = extractDBTypeFromURL(url)
+    val dbType = extractDBTypeFromUrl(url)
 
     val tableTypes = arrayOf("TABLE")
     // exclude system and other tables without data
diff --git a/plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/GenerateDataSchemaTask.kt b/plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/GenerateDataSchemaTask.kt
@@ -18,8 +18,10 @@ import org.jetbrains.kotlinx.dataframe.impl.codeGen.toStandaloneSnippet
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.urlCodeGenReader
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.urlDfReader
 import java.io.File
+import java.lang.RuntimeException
 import java.net.URL
 import java.nio.file.Paths
+import java.sql.Connection
 import java.sql.DriverManager
 import org.jetbrains.kotlinx.dataframe.io.ArrowFeather
 import org.jetbrains.kotlinx.dataframe.io.CSV
@@ -30,6 +32,7 @@ import org.jetbrains.kotlinx.dataframe.io.TSV
 import org.jetbrains.kotlinx.dataframe.io.getSchemaForSqlQuery
 import org.jetbrains.kotlinx.dataframe.io.getSchemaForSqlTable
 import org.jetbrains.kotlinx.dataframe.io.isURL
+import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
 
 abstract class GenerateDataSchemaTask : DefaultTask() {
 
@@ -84,9 +87,7 @@ abstract class GenerateDataSchemaTask : DefaultTask() {
         if (rawUrl.startsWith("jdbc")) {
             val connection = DriverManager.getConnection(rawUrl, jdbcOptions.user, jdbcOptions.password)
             connection.use {
-                val schema = if(jdbcOptions.sqlQuery.isBlank())
-                    DataFrame.getSchemaForSqlTable(connection,  interfaceName.get())
-                else DataFrame.getSchemaForSqlQuery(connection, jdbcOptions.sqlQuery)
+                val schema = generateSchemaByJdbcOptions(jdbcOptions, connection)
 
                 val codeGenerator = CodeGenerator.create(useFqNames = false)
 
@@ -174,6 +175,23 @@ abstract class GenerateDataSchemaTask : DefaultTask() {
         }
     }
 
+    private fun generateSchemaByJdbcOptions(
+        jdbcOptions: JdbcOptionsDsl,
+        connection: Connection
+    ): DataFrameSchema {
+        logger.debug("Table name: ${jdbcOptions.tableName}")
+        logger.debug("SQL query: ${jdbcOptions.sqlQuery}")
+
+        return if (jdbcOptions.tableName.isNotBlank())
+            DataFrame.getSchemaForSqlTable(connection, jdbcOptions.tableName)
+        else if(jdbcOptions.sqlQuery.isNotBlank())
+            DataFrame.getSchemaForSqlQuery(connection, jdbcOptions.sqlQuery)
+        else throw RuntimeException("Table name: ${jdbcOptions.tableName}, " +
+            "SQL query: ${jdbcOptions.sqlQuery} both are empty! " +
+            "Populate 'tableName' or 'sqlQuery' in jdbcOptions with value to generate schema " +
+            "for SQL table or result of SQL query!")
+    }
+
     private fun stringOf(data: Any): String =
         when (data) {
             is File -> data.absolutePath
diff --git a/plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/SchemaGeneratorExtension.kt b/plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/SchemaGeneratorExtension.kt
@@ -135,5 +135,6 @@ data class JsonOptionsDsl(
 data class JdbcOptionsDsl(
     var user: String = "", // TODO: I'm not sure about the default parameters
     var password: String = "", // TODO: I'm not sure about the default parameters
+    var tableName: String = "",
     var sqlQuery: String = ""
 ) : Serializable
diff --git a/plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/DataSchemaGenerator.kt b/plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/DataSchemaGenerator.kt
@@ -24,9 +24,18 @@ import org.jetbrains.kotlinx.dataframe.impl.codeGen.urlCodeGenReader
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.urlDfReader
 import org.jetbrains.kotlinx.dataframe.io.*
 import java.io.File
+import java.lang.RuntimeException
 import java.net.MalformedURLException
 import java.net.URL
+import java.sql.Connection
 import java.sql.DriverManager
+import org.jetbrains.kotlinx.dataframe.io.db.H2
+import org.jetbrains.kotlinx.dataframe.io.db.MariaDb
+import org.jetbrains.kotlinx.dataframe.io.db.MySql
+import org.jetbrains.kotlinx.dataframe.io.db.PostgreSql
+import org.jetbrains.kotlinx.dataframe.io.db.Sqlite
+import org.jetbrains.kotlinx.dataframe.io.db.driverClassNameFromUrl
+import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
 
 @OptIn(KspExperimental::class)
 class DataSchemaGenerator(
@@ -159,7 +168,8 @@ class DataSchemaGenerator(
         if (importStatement.isJdbc) {
             val url = importStatement.dataSource.pathRepresentation
 
-            if(url.contains("h2")) Class.forName("org.h2.Driver")
+            // Force classloading
+            Class.forName(driverClassNameFromUrl(url))
 
             val connection = DriverManager.getConnection(
                 url,
@@ -168,9 +178,7 @@ class DataSchemaGenerator(
             )
 
             connection.use {
-                val schema = if(importStatement.jdbcOptions.sqlQuery.isBlank())
-                    DataFrame.getSchemaForSqlTable(connection, importStatement.name)
-                else DataFrame.getSchemaForSqlQuery(connection, importStatement.jdbcOptions.sqlQuery)
+                val schema = generateSchemaForImport(importStatement, connection)
 
                 val codeGenerator = CodeGenerator.create(useFqNames = false)
 
@@ -257,5 +265,22 @@ class DataSchemaGenerator(
             it.write(code)
         }
     }
+
+    private fun generateSchemaForImport(
+        importStatement: ImportDataSchemaStatement,
+        connection: Connection
+    ): DataFrameSchema {
+        logger.info("Table name: ${importStatement.jdbcOptions.tableName}")
+        logger.info("SQL query: ${importStatement.jdbcOptions.sqlQuery}")
+
+        return if (importStatement.jdbcOptions.tableName.isNotBlank())
+            DataFrame.getSchemaForSqlTable(connection, importStatement.jdbcOptions.tableName)
+        else if(importStatement.jdbcOptions.sqlQuery.isNotBlank())
+            DataFrame.getSchemaForSqlQuery(connection, importStatement.jdbcOptions.sqlQuery)
+        else throw RuntimeException("Table name: ${importStatement.jdbcOptions.tableName}, " +
+            "SQL query: ${importStatement.jdbcOptions.sqlQuery} both are empty! " +
+            "Populate 'tableName' or 'sqlQuery' in jdbcOptions with value to generate schema " +
+            "for SQL table or result of SQL query!")
+    }
 }
 
diff --git a/plugins/symbol-processor/src/test/kotlin/org/jetbrains/dataframe/ksp/DataFrameJdbcSymbolProcessorTest.kt b/plugins/symbol-processor/src/test/kotlin/org/jetbrains/dataframe/ksp/DataFrameJdbcSymbolProcessorTest.kt
@@ -102,6 +102,7 @@ class DataFrameJdbcSymbolProcessorTest {
                 package test
                 
                 import org.jetbrains.kotlinx.dataframe.annotations.ImportDataSchema
+                import org.jetbrains.kotlinx.dataframe.annotations.JdbcOptions
                 import org.jetbrains.kotlinx.dataframe.api.filter
                 import org.jetbrains.kotlinx.dataframe.DataFrame
                 import org.jetbrains.kotlinx.dataframe.api.cast
@@ -126,11 +127,16 @@ class DataFrameJdbcSymbolProcessorTest {
                         SourceFile.kotlin(
                             "MySources.kt",
                             """
-                @file:ImportDataSchema(name = "Customer", path = "$CONNECTION_URL")
-                
+                @file:ImportDataSchema(
+                     "Customer",
+                     "$CONNECTION_URL",
+                     jdbcOptions = JdbcOptions("", "", tableName = "Customer")
+                )
+
                 package test
                 
                 import org.jetbrains.kotlinx.dataframe.annotations.ImportDataSchema
+                import org.jetbrains.kotlinx.dataframe.annotations.JdbcOptions
                 import org.jetbrains.kotlinx.dataframe.api.filter
                 import org.jetbrains.kotlinx.dataframe.DataFrame
                 import org.jetbrains.kotlinx.dataframe.api.cast
@@ -161,11 +167,16 @@ class DataFrameJdbcSymbolProcessorTest {
                     SourceFile.kotlin(
                         "MySources.kt",
                         """
-                @file:ImportDataSchema(name = "Customer", path = "$CONNECTION_URL")
+                @file:ImportDataSchema(
+                     "Customer",
+                     "$CONNECTION_URL",
+                     jdbcOptions = JdbcOptions("", "", tableName = "Customer")
+                )
                 
                 package test
                 
                 import org.jetbrains.kotlinx.dataframe.annotations.ImportDataSchema
+                import org.jetbrains.kotlinx.dataframe.annotations.JdbcOptions
                 import org.jetbrains.kotlinx.dataframe.api.filter
                 import org.jetbrains.kotlinx.dataframe.DataFrame
                 import org.jetbrains.kotlinx.dataframe.api.cast

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@ public annotation class CsvOptions(`
`52`	`52`	`public annotation class JdbcOptions(`
`53`	`53`	`public val user: String = "", // TODO: I'm not sure about the default parameters`
`54`	`54`	`public val password: String = "", // TODO: I'm not sure about the default parameters)`
	`55`	`+ public val tableName: String = "",`
`55`	`56`	`public val sqlQuery: String = ""`
`56`	`57`	`)`
`57`	`58`