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 index 0d5ce169c..428f753ca 100644 --- 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 index 08941e0e7..1664c6234 100644 --- 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 index 33c9d2b3d..791b541cb 100644 --- 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,6 +174,7 @@ abstract class GenerateDataSchemaTask : DefaultTask() { } } + // TODO: copy pasted from symbol-processor: DataSchemaGenerator, should be refactored somehow private fun generateSchemaByJdbcOptions( jdbcOptions: JdbcOptionsDsl, connection: Connection, @@ -181,20 +182,45 @@ abstract class GenerateDataSchemaTask : DefaultTask() { 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 index 5f424bd6a..6b3610e04 100644 --- 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 = 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 index 18756bdeb..d3efad8f3 100644 --- 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 { 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 index bc98fa3f4..dd1aec2d1 100644 --- 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 index 1a9ad70f7..09ddfdfc6 100644 --- 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() + df.filter { it[Customer::age] != null && it[Customer::age]!! > 30 } + + val df1 = DataFrame.readSqlTable(connection, tableName, 1).cast() + df1.filter { it[Customer::age] != null && it[Customer::age]!! > 30 } + + val dbConfig = DatabaseConfiguration(url = "$CONNECTION_URL") + val df2 = DataFrame.readSqlTable(dbConfig, tableName).cast() + df2.filter { it[Customer::age] != null && it[Customer::age]!! > 30 } + + val df3 = DataFrame.readSqlTable(dbConfig, tableName, 1).cast() + df3.filter { it[Customer::age] != null && it[Customer::age]!! > 30 } + + } + } + """.trimIndent() + ) + ) + ) + ) + result.successfulCompilation shouldBe true + } + private fun KotlinCompileTestingCompilationResult.inspectLines(f: (List) -> Unit) { inspectLines(generatedFile, f) }