generated from TBD54566975/tbd-project-template
-
Notifications
You must be signed in to change notification settings - Fork 8
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: provide Kotlin API for db connections
- Loading branch information
Showing
4 changed files
with
80 additions
and
122 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
78 changes: 32 additions & 46 deletions
78
kotlin-runtime/ftl-runtime/src/main/kotlin/xyz/block/ftl/database/Db.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,66 +1,52 @@ | ||
| package xyz.block.ftl.database | ||
|
|
||
| import com.fasterxml.jackson.databind.ObjectMapper | ||
| import com.fasterxml.jackson.databind.exc.MismatchedInputException | ||
| import com.fasterxml.jackson.dataformat.yaml.YAMLFactory | ||
| import com.fasterxml.jackson.module.kotlin.registerKotlinModule | ||
| import xyz.block.ftl.logging.Logging | ||
| import java.nio.file.Files | ||
| import java.nio.file.Paths | ||
| import java.net.URI | ||
| import java.sql.Connection | ||
| import java.sql.DriverManager | ||
|
|
||
| private const val DEFAULT_DB_YAML = "db.yaml" | ||
|
|
||
| data class DatabaseConfig( | ||
| val databases: Map<String, Database>, | ||
| ) | ||
|
|
||
| data class Database( | ||
| val jdbcUrl: String, | ||
| val username: String, | ||
| val password: String, | ||
| ) | ||
| const val DSN_VAR = "FTL_POSTGRES_DSN" | ||
|
|
||
| /** | ||
| * `Db` is a simple wrapper around the JDBC driver manager that provides a connection to the database specified | ||
| * by the DB YAML file (`db.yaml` unless otherwise specified). The YAML must be provided in the `src/main/resources` | ||
| * directory of an FTL module. | ||
| * | ||
| * Example YAML file: | ||
| * databases: | ||
| * example1: | ||
| * jdbcUrl: jdbc:postgresql://<host>:<port>/<dbname> | ||
| * username: <user> | ||
| * password: <password> | ||
| * example2: | ||
| * ... | ||
| * by the FTL_POSTGRES_DSN environment variable. | ||
| */ | ||
| object Db { | ||
| private val logger = Logging.logger(Db::class) | ||
| private val mapper = ObjectMapper(YAMLFactory()).registerKotlinModule() | ||
| fun Conn(id: String? = null, configFile: String = DEFAULT_DB_YAML): Connection { | ||
| require(configFile.endsWith(".yaml") || configFile.endsWith(".yml")) | ||
|
|
||
| fun conn(): Connection { | ||
| return try { | ||
| val config = Files.newBufferedReader(Paths.get("classes", configFile)).use { | ||
| mapper.readValue(it, DatabaseConfig::class.java) | ||
| } | ||
| val dsn = System.getenv(DSN_VAR) | ||
| require(dsn != null) { "$DSN_VAR environment variable not set" } | ||
|
|
||
| val db = if (id != null) { | ||
| config.databases[id] ?: throw IllegalArgumentException("Could not find config for database $id") | ||
| } else { | ||
| require(config.databases.size == 1) { "Multiple databases found, please specify a name" } | ||
| config.databases.values.single() | ||
| } | ||
|
|
||
| DriverManager.getConnection(db.jdbcUrl, db.username, db.password) | ||
| } catch (e: MismatchedInputException) { | ||
| logger.error("Could not read $configFile file", e) | ||
| throw e | ||
| DriverManager.getConnection(dsnToJdbcUrl(dsn)) | ||
| } catch (e: Exception) { | ||
| logger.error("Could not connect to database", e) | ||
| throw e | ||
| } | ||
| } | ||
|
|
||
| private fun dsnToJdbcUrl(dsn: String): String { | ||
| val uri = URI(dsn) | ||
| val scheme = uri.scheme ?: throw IllegalArgumentException("Missing scheme in DSN.") | ||
| val userInfo = uri.userInfo?.split(":") ?: throw IllegalArgumentException("Missing userInfo in DSN.") | ||
| val user = userInfo.firstOrNull() ?: throw IllegalArgumentException("Missing user in userInfo.") | ||
| val password = if (userInfo.size > 1) userInfo[1] else "" | ||
| val host = uri.host ?: throw IllegalArgumentException("Missing host in DSN.") | ||
| val port = if (uri.port != -1) uri.port.toString() else throw IllegalArgumentException("Missing port in DSN.") | ||
| val database = uri.path.trimStart('/') | ||
| val parameters = uri.query?.replace("&", "?") ?: "" | ||
|
|
||
| val jdbcScheme = when (scheme) { | ||
| "mysql" -> "jdbc:mysql" | ||
| "postgresql" -> "jdbc:postgresql" | ||
| else -> throw IllegalArgumentException("Unsupported scheme: $scheme") | ||
| } | ||
|
|
||
| val jdbcUrl = "$jdbcScheme://$host:$port/$database?$parameters" | ||
| return if (user.isNotBlank() && password.isNotBlank()) { | ||
| "$jdbcUrl&user=$user&password=$password" | ||
| } else { | ||
| jdbcUrl | ||
| } | ||
| } | ||
| } |