Add JDBC credentials extraction from env variables and improve exception handling

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/ImportDataSchema.kt

@@ -49,9 +49,21 @@ public annotation class CsvOptions(
     public val delimiter: Char,
 )
 
+/**
+ * An annotation class that represents options for JDBC connection.
+ *
+ * @property [user] The username for the JDBC connection. Default value is an empty string.
+ * If [extractCredFromEnv] is true, the [user] value will be interpreted as key for system environment variable.
+ * @property [password] The password for the JDBC connection. Default value is an empty string.
+ * If [extractCredFromEnv] is true, the [password] value will be interpreted as key for system environment variable.
+ * @property [extractCredFromEnv] Whether to extract the JDBC credentials from environment variables. Default value is false.
+ * @property [tableName] The name of the table for the JDBC connection. Default value is an empty string.
+ * @property [sqlQuery] The SQL query to be executed in the JDBC connection. Default value is an empty string.
+ */
 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 user: String = "",
+    public val password: String = "",
+    public val extractCredFromEnv: Boolean = false,
     public val tableName: String = "",
     public val sqlQuery: String = ""
 )

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/ImportDataSchema.kt

@@ -49,9 +49,21 @@ public annotation class CsvOptions(
     public val delimiter: Char,
 )
 
+/**
+ * An annotation class that represents options for JDBC connection.
+ *
+ * @property [user] The username for the JDBC connection. Default value is an empty string.
+ * If [extractCredFromEnv] is true, the [user] value will be interpreted as key for system environment variable.
+ * @property [password] The password for the JDBC connection. Default value is an empty string.
+ * If [extractCredFromEnv] is true, the [password] value will be interpreted as key for system environment variable.
+ * @property [extractCredFromEnv] Whether to extract the JDBC credentials from environment variables. Default value is false.
+ * @property [tableName] The name of the table for the JDBC connection. Default value is an empty string.
+ * @property [sqlQuery] The SQL query to be executed in the JDBC connection. Default value is an empty string.
+ */
 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 user: String = "",
+    public val password: String = "",
+    public val extractCredFromEnv: Boolean = false,
     public val tableName: String = "",
     public val sqlQuery: String = ""
 )

plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/GenerateDataSchemaTask.kt

@@ -174,27 +174,53 @@ abstract class GenerateDataSchemaTask : DefaultTask() {
         }
     }
 
+    // TODO: copy pasted from symbol-processor: DataSchemaGenerator, should be refactored somehow
     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!"
-            )
+        val tableName = jdbcOptions.tableName
+        val sqlQuery = jdbcOptions.sqlQuery
+
+        return when {
+            isTableNameNotBlankAndQueryBlank(tableName, sqlQuery) -> generateSchemaForTable(connection, tableName)
+            isQueryNotBlankAndTableBlank(tableName, sqlQuery) -> generateSchemaForQuery(connection, sqlQuery)
+            areBothNotBlank(tableName, sqlQuery) -> throwBothFieldsFilledException(tableName, sqlQuery)
+            else -> throwBothFieldsEmptyException(tableName, sqlQuery)
         }
     }
 
+    private fun isTableNameNotBlankAndQueryBlank(tableName: String, sqlQuery: String) =
+        tableName.isNotBlank() && sqlQuery.isBlank()
+
+    private fun isQueryNotBlankAndTableBlank(tableName: String, sqlQuery: String) =
+        sqlQuery.isNotBlank() && tableName.isBlank()
+
+    private fun areBothNotBlank(tableName: String, sqlQuery: String) = sqlQuery.isNotBlank() && tableName.isNotBlank()
+
+    private fun generateSchemaForTable(connection: Connection, tableName: String) =
+        DataFrame.getSchemaForSqlTable(connection, tableName)
+
+    private fun generateSchemaForQuery(connection: Connection, sqlQuery: String) =
+        DataFrame.getSchemaForSqlQuery(connection, sqlQuery)
+
+    private fun throwBothFieldsFilledException(tableName: String, sqlQuery: String): Nothing {
+        throw RuntimeException(
+            "Table name '$tableName' and SQL query '$sqlQuery' both are filled! " +
+                "Clear 'tableName' or 'sqlQuery' properties in jdbcOptions with value to generate schema for SQL table or result of SQL query!"
+        )
+    }
+
+    private fun throwBothFieldsEmptyException(tableName: String, sqlQuery: String): Nothing {
+        throw RuntimeException(
+            "Table name '$tableName' and SQL query '$sqlQuery' both are empty! " +
+                "Populate 'tableName' or 'sqlQuery' properties 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

plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/SchemaGeneratorExtension.kt

@@ -132,9 +132,17 @@ data class JsonOptionsDsl(
     var keyValuePaths: List<JsonPath> = emptyList(),
 ) : Serializable
 
+/**
+ * Represents the configuration options for JDBC data source.
+ *
+ * @property [user] The username used to authenticate with the database. Default is an empty string.
+ * @property [password] The password used to authenticate with the database. Default is an empty string.
+ * @property [tableName] The name of the table to generate schema for. Default is an empty string.
+ * @property [sqlQuery] The SQL query used to generate schema. Default is an empty string.
+ */
 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 user: String = "",
+    var password: String = "",
     var tableName: String = "",
     var sqlQuery: String = ""
 ) : Serializable

plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/SchemaGeneratorPlugin.kt

@@ -132,7 +132,7 @@ class SchemaGeneratorPlugin : Plugin<Project> {
             this.schemaVisibility.set(visibility)
             this.csvOptions.set(schema.csvOptions)
             this.jsonOptions.set(schema.jsonOptions)
-            this.jdbcOptions.set(schema.jdbcOptions) // TODO: probably remove
+            this.jdbcOptions.set(schema.jdbcOptions)
             this.defaultPath.set(defaultPath)
             this.delimiters.set(delimiters)
         }

plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/DataSchemaGenerator.kt

@@ -172,10 +172,19 @@ class DataSchemaGenerator(
             // Force classloading
             Class.forName(driverClassNameFromUrl(url))
 
+            var userName = importStatement.jdbcOptions.user
+            var password = importStatement.jdbcOptions.password
+
+            // treat the passed userName and password parameters as env variables
+            if (importStatement.jdbcOptions.extractCredFromEnv) {
+                userName = System.getenv(userName) ?: userName
+                password = System.getenv(password) ?: password
+            }
+
             val connection = DriverManager.getConnection(
                 url,
-                importStatement.jdbcOptions.user,
-                importStatement.jdbcOptions.password
+                userName,
+                password
             )
 
             connection.use {
@@ -271,22 +280,47 @@ class DataSchemaGenerator(
 
     private fun generateSchemaForImport(
         importStatement: ImportDataSchemaStatement,
-        connection: Connection,
+        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!"
-            )
+        val tableName = importStatement.jdbcOptions.tableName
+        val sqlQuery = importStatement.jdbcOptions.sqlQuery
+
+        return when {
+            isTableNameNotBlankAndQueryBlank(tableName, sqlQuery) -> generateSchemaForTable(connection, tableName)
+            isQueryNotBlankAndTableBlank(tableName, sqlQuery) -> generateSchemaForQuery(connection, sqlQuery)
+            areBothNotBlank(tableName, sqlQuery) -> throwBothFieldsFilledException(tableName, sqlQuery)
+            else -> throwBothFieldsEmptyException(tableName, sqlQuery)
         }
     }
+
+    private fun isTableNameNotBlankAndQueryBlank(tableName: String, sqlQuery: String) =
+        tableName.isNotBlank() && sqlQuery.isBlank()
+
+    private fun isQueryNotBlankAndTableBlank(tableName: String, sqlQuery: String) =
+        sqlQuery.isNotBlank() && tableName.isBlank()
+
+    private fun areBothNotBlank(tableName: String, sqlQuery: String) = sqlQuery.isNotBlank() && tableName.isNotBlank()
+
+    private fun generateSchemaForTable(connection: Connection, tableName: String) =
+        DataFrame.getSchemaForSqlTable(connection, tableName)
+
+    private fun generateSchemaForQuery(connection: Connection, sqlQuery: String) =
+        DataFrame.getSchemaForSqlQuery(connection, sqlQuery)
+
+    private fun throwBothFieldsFilledException(tableName: String, sqlQuery: String): Nothing {
+        throw RuntimeException(
+            "Table name '$tableName' and SQL query '$sqlQuery' both are filled! " +
+                "Clear 'tableName' or 'sqlQuery' properties in jdbcOptions with value to generate schema for SQL table or result of SQL query!"
+        )
+    }
+
+    private fun throwBothFieldsEmptyException(tableName: String, sqlQuery: String): Nothing {
+        throw RuntimeException(
+            "Table name '$tableName' and SQL query '$sqlQuery' both are empty! " +
+                "Populate 'tableName' or 'sqlQuery' properties in jdbcOptions with value to generate schema for SQL table or result of SQL query!"
+        )
+    }
 }

plugins/symbol-processor/src/test/kotlin/org/jetbrains/dataframe/ksp/DataFrameJdbcSymbolProcessorTest.kt

@@ -212,6 +212,62 @@ class DataFrameJdbcSymbolProcessorTest {
         result.successfulCompilation shouldBe true
     }
 
+    /**
+     * Test code is copied from test above.
+     */
+    @Test
+    fun `schema extracted via readFromDB method is resolved with db credentials from env variables`() {
+        val result = KspCompilationTestRunner.compile(
+            TestCompilationParameters(
+                sources = listOf(
+                    SourceFile.kotlin(
+                        "MySources.kt",
+                        """
+                @file:ImportDataSchema(
+                     "Customer",
+                     "$CONNECTION_URL",
+                     jdbcOptions = JdbcOptions("", "", extractCredFromEnv = true, 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
+                import java.sql.Connection
+                import java.sql.DriverManager
+                import java.sql.SQLException
+                import org.jetbrains.kotlinx.dataframe.io.readSqlTable
+                import org.jetbrains.kotlinx.dataframe.io.DatabaseConfiguration
+
+                fun main() {    
+                    val tableName = "Customer"
+                    DriverManager.getConnection("$CONNECTION_URL").use { connection ->
+                        val df = DataFrame.readSqlTable(connection, tableName).cast<Customer>()
+                        df.filter { it[Customer::age] != null && it[Customer::age]!! > 30 }
+
+                        val df1 = DataFrame.readSqlTable(connection, tableName, 1).cast<Customer>()
+                        df1.filter { it[Customer::age] != null && it[Customer::age]!! > 30 }
+
+                        val dbConfig = DatabaseConfiguration(url = "$CONNECTION_URL")
+                        val df2 = DataFrame.readSqlTable(dbConfig, tableName).cast<Customer>()
+                        df2.filter { it[Customer::age] != null && it[Customer::age]!! > 30 }
+
+                        val df3 = DataFrame.readSqlTable(dbConfig, tableName, 1).cast<Customer>()
+                        df3.filter { it[Customer::age] != null && it[Customer::age]!! > 30 }
+
+                    }
+                }
+                        """.trimIndent()
+                    )
+                )
+            )
+        )
+        result.successfulCompilation shouldBe true
+    }
+
     private fun KotlinCompileTestingCompilationResult.inspectLines(f: (List<String>) -> Unit) {
         inspectLines(generatedFile, f)
     }