Add JDBC credentials extraction from env variables and improve exception handling (#692)

zaleslaw · web-flow · commit efdbadb546f6 · 2024-05-13T17:56:36.000+02:00
* Add JDBC credentials extraction and improve exception handling

This update introduces the ability to extract JDBC credentials from environment variables enhancing security features. Additionally, improvements have been made to the exception handling when generating schema for SQL table or result of SQL query.

* Add environment variable support to JDBC credentials

Updated the JDBC connection options to interpret the user and password values as keys for system environment variables when extractCredFromEnv is set to true. This change allows more security and flexibility in handling sensitive JDBC credentials.

* Refactor code for readability and enhance jdbcOptions documentation
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
@@ -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 = ""
 )
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
@@ -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 = ""
 )
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
@@ -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
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
@@ -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
diff --git a/plugins/dataframe-gradle-plugin/src/main/kotlin/org/jetbrains/dataframe/gradle/SchemaGeneratorPlugin.kt b/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)
         }
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
@@ -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!"
+        )
+    }
 }
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
@@ -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)
     }

Original file line number	Diff line number	Diff line change
`@@ -132,7 +132,7 @@ class SchemaGeneratorPlugin : Plugin<Project> {`
`132`	`132`	`this.schemaVisibility.set(visibility)`
`133`	`133`	`this.csvOptions.set(schema.csvOptions)`
`134`	`134`	`this.jsonOptions.set(schema.jsonOptions)`
`135`		`- this.jdbcOptions.set(schema.jdbcOptions) // TODO: probably remove`
	`135`	`+ this.jdbcOptions.set(schema.jdbcOptions)`
`136`	`136`	`this.defaultPath.set(defaultPath)`
`137`	`137`	`this.delimiters.set(delimiters)`
`138`	`138`	`}`