diff --git a/dataframe-jdbc/build.gradle.kts b/dataframe-jdbc/build.gradle.kts index 968984f0f3..e7c81d10f8 100644 --- a/dataframe-jdbc/build.gradle.kts +++ b/dataframe-jdbc/build.gradle.kts @@ -33,6 +33,22 @@ dependencies { exclude("org.jetbrains.kotlin", "kotlin-stdlib-jdk8") } testImplementation(libs.hikaricp) + testImplementation(libs.testcontainers) + testImplementation(libs.testcontainers.postgresql) + testImplementation(libs.testcontainers.mysql) + testImplementation(libs.testcontainers.mariadb) + testImplementation(libs.testcontainers.mssqlserver) +} + +buildConfig { + sourceSets.named("test") { + packageName = "org.jetbrains.kotlinx.dataframe.io.testcontainers" + className = "BuildConfig" + buildConfigField("MARIADB_IMAGE", "mariadb:${libs.versions.dockerImage.mariadb.get()}") + buildConfigField("MYSQL_IMAGE", "mysql:${libs.versions.dockerImage.mysql.get()}") + buildConfigField("POSTGRES_IMAGE", "postgres:${libs.versions.dockerImage.postgres.get()}") + buildConfigField("MSSQL_IMAGE", "mcr.microsoft.com/mssql/server:${libs.versions.dockerImage.mssql.get()}") + } } kotlinPublications { @@ -47,3 +63,36 @@ kotlinPublications { tasks.processKDocsMain { dependsOn(tasks.generateBuildConfigClasses) } + +// Implementations of the abstract database tests running against databases in Docker containers +private val testcontainersTests = "org.jetbrains.kotlinx.dataframe.io.testcontainers.*" + +// Implementations of the abstract database tests running against database servers on localhost +private val localDbTests = "org.jetbrains.kotlinx.dataframe.io.local.*LocalTest" + +tasks.test { + filter { + excludeTestsMatching(testcontainersTests) + excludeTestsMatching(localDbTests) + } +} + +tasks.register("testcontainersTest") { + description = "Runs tests that require Docker via Testcontainers." + group = LifecycleBasePlugin.VERIFICATION_GROUP + testClassesDirs = sourceSets.test.get().output.classesDirs + classpath = sourceSets.test.get().runtimeClasspath + filter { + includeTestsMatching(testcontainersTests) + } +} + +tasks.register("localDbTest") { + description = "Runs tests that require database servers running on localhost." + group = LifecycleBasePlugin.VERIFICATION_GROUP + testClassesDirs = sourceSets.test.get().output.classesDirs + classpath = sourceSets.test.get().runtimeClasspath + filter { + includeTestsMatching(localDbTests) + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mariadbTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mariadbTest.kt index 7cb7629ce4..3c5936858c 100644 --- a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mariadbTest.kt +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mariadbTest.kt @@ -1,520 +1,43 @@ package org.jetbrains.kotlinx.dataframe.io.local -import io.kotest.matchers.shouldBe -import org.intellij.lang.annotations.Language -import org.jetbrains.kotlinx.dataframe.DataFrame -import org.jetbrains.kotlinx.dataframe.annotations.DataSchema -import org.jetbrains.kotlinx.dataframe.api.add -import org.jetbrains.kotlinx.dataframe.api.cast -import org.jetbrains.kotlinx.dataframe.api.filter -import org.jetbrains.kotlinx.dataframe.api.select -import org.jetbrains.kotlinx.dataframe.io.inferNullability -import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables -import org.jetbrains.kotlinx.dataframe.io.readSqlQuery -import org.jetbrains.kotlinx.dataframe.io.readSqlTable -import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.jetbrains.kotlinx.dataframe.io.MariadbTestBase +import org.jetbrains.kotlinx.dataframe.io.setUpMariadbTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownMariadbTestData import org.junit.AfterClass import org.junit.BeforeClass -import org.junit.Ignore -import org.junit.Test -import java.math.BigDecimal -import java.math.BigInteger -import java.sql.Blob import java.sql.Connection import java.sql.DriverManager -import java.sql.SQLException -import java.util.Date -import kotlin.reflect.typeOf -import kotlin.time.Instant private const val URL = "jdbc:mariadb://localhost:3306" private const val USER_NAME = "root" private const val PASSWORD = "pass" -private const val TEST_DATABASE_NAME = "testKDFdatabase" -@DataSchema -interface Table1MariaDb { - val id: Int - val bitCol: Boolean - val tinyintCol: Int - val smallintCol: Short? - val mediumintCol: Int - val mediumintUnsignedCol: Int - val integerCol: Int - val intCol: Int - val integerUnsignedCol: Long - val bigintCol: Long - val floatCol: Float - val doubleCol: Double - val decimalCol: BigDecimal - val dateCol: String - val datetimeCol: String - val timestampCol: String - val timeCol: String - val yearCol: String - val varcharCol: String - val charCol: String - val binaryCol: ByteArray - val varbinaryCol: ByteArray - val tinyblobCol: ByteArray - val blobCol: ByteArray - val mediumblobCol: ByteArray - val longblobCol: ByteArray - val textCol: String - val mediumtextCol: String - val longtextCol: String - val enumCol: String - val setCol: Char - val bigintUnsignedCol: BigInteger - val jsonCol: String -} - -@DataSchema -interface Table2MariaDb { - val id: Int - val bitCol: Boolean? - val tinyintCol: Int? - val smallintCol: Int? - val mediumintCol: Int? - val mediumintUnsignedCol: Int? - val integerCol: Int? - val intCol: Int? - val integerUnsignedCol: Long? - val bigintCol: Long? - val floatCol: Float? - val doubleCol: Double? - val decimalCol: Double? - val dateCol: String? - val datetimeCol: String? - val timestampCol: String? - val timeCol: String? - val yearCol: String? - val varcharCol: String? - val charCol: String? - val binaryCol: ByteArray? - val varbinaryCol: ByteArray? - val tinyblobCol: ByteArray? - val blobCol: ByteArray? - val mediumblobCol: ByteArray? - val longblobCol: ByteArray? - val textCol: String? - val mediumtextCol: String? - val longtextCol: String? - val enumCol: String? - val setCol: Char? - val bigintUnsignedCol: BigInteger? - val jsonCol: String? -} - -@DataSchema -interface Table3MariaDb { - val id: Int - val enumCol: String - val setCol: Char? -} +class MariadbLocalTest : MariadbTestBase() { + override val connection: Connection get() = Companion.connection -private const val JSON_STRING = - "{\"details\": {\"foodType\": \"Pizza\", \"menu\": \"https://www.loumalnatis.com/our-menu\"}, \n" + - " \t\"favorites\": [{\"description\": \"Pepperoni deep dish\", \"price\": 18.75}, \n" + - "{\"description\": \"The Lou\", \"price\": 24.75}]}" + override fun connect(database: String?): Connection = openConnection(database) -@Ignore -class MariadbTest { companion object { private lateinit var connection: Connection + private fun openConnection(database: String?): Connection = + DriverManager.getConnection( + if (database == null) URL else "$URL/$database", + USER_NAME, + PASSWORD, + ) + @BeforeClass @JvmStatic fun setUpClass() { - connection = DriverManager.getConnection(URL, USER_NAME, PASSWORD) - - connection.createStatement().use { st -> - // Drop the test database if it exists - val dropDatabaseQuery = "DROP DATABASE IF EXISTS $TEST_DATABASE_NAME" - st.executeUpdate(dropDatabaseQuery) - - // Create the test database - val createDatabaseQuery = "CREATE DATABASE $TEST_DATABASE_NAME" - st.executeUpdate(createDatabaseQuery) - - // Use the newly created database - val useDatabaseQuery = "USE $TEST_DATABASE_NAME" - st.executeUpdate(useDatabaseQuery) - } - - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } - - @Language("SQL") - val createTableQuery = """ - CREATE TABLE IF NOT EXISTS table1 ( - id INT AUTO_INCREMENT PRIMARY KEY, - bitCol BIT NOT NULL, - tinyintCol TINYINT NOT NULL, - smallintCol SMALLINT, - mediumintCol MEDIUMINT NOT NULL, - mediumintUnsignedCol MEDIUMINT UNSIGNED NOT NULL, - integerCol INTEGER NOT NULL, - intCol INT NOT NULL, - integerUnsignedCol INTEGER UNSIGNED NOT NULL, - bigintCol BIGINT NOT NULL, - floatCol FLOAT NOT NULL, - doubleCol DOUBLE NOT NULL, - decimalCol DECIMAL NOT NULL, - dateCol DATE NOT NULL, - datetimeCol DATETIME NOT NULL, - timestampCol TIMESTAMP NOT NULL, - timeCol TIME NOT NULL, - yearCol YEAR NOT NULL, - varcharCol VARCHAR(255) NOT NULL, - charCol CHAR(10) NOT NULL, - binaryCol BINARY(64) NOT NULL, - varbinaryCol VARBINARY(128) NOT NULL, - tinyblobCol TINYBLOB NOT NULL, - blobCol BLOB NOT NULL, - mediumblobCol MEDIUMBLOB NOT NULL , - longblobCol LONGBLOB NOT NULL, - textCol TEXT NOT NULL, - mediumtextCol MEDIUMTEXT NOT NULL, - longtextCol LONGTEXT NOT NULL, - enumCol ENUM('Value1', 'Value2', 'Value3') NOT NULL, - setCol SET('Option1', 'Option2', 'Option3') NOT NULL, - bigintUnsignedCol BIGINT UNSIGNED NOT NULL, - jsonCol JSON NOT NULL - CHECK (JSON_VALID(jsonCol)) - ) - """ - connection.createStatement().execute(createTableQuery.trimIndent()) - - @Language("SQL") - val createTableQuery2 = """ - CREATE TABLE IF NOT EXISTS table2 ( - id INT AUTO_INCREMENT PRIMARY KEY, - bitCol BIT, - tinyintCol TINYINT, - smallintCol SMALLINT, - mediumintCol MEDIUMINT, - mediumintUnsignedCol MEDIUMINT UNSIGNED, - integerCol INTEGER, - intCol INT, - integerUnsignedCol INTEGER UNSIGNED, - bigintCol BIGINT, - floatCol FLOAT, - doubleCol DOUBLE, - decimalCol DECIMAL, - dateCol DATE, - datetimeCol DATETIME, - timestampCol TIMESTAMP, - timeCol TIME, - yearCol YEAR, - varcharCol VARCHAR(255), - charCol CHAR(10), - binaryCol BINARY(64), - varbinaryCol VARBINARY(128), - tinyblobCol TINYBLOB, - blobCol BLOB, - mediumblobCol MEDIUMBLOB, - longblobCol LONGBLOB, - textCol TEXT, - mediumtextCol MEDIUMTEXT, - longtextCol LONGTEXT, - enumCol ENUM('Value1', 'Value2', 'Value3'), - setCol SET('Option1', 'Option2', 'Option3'), - bigintUnsignedCol BIGINT UNSIGNED - ) - """ - connection.createStatement().execute(createTableQuery2.trimIndent()) - - @Language("SQL") - val insertData1 = - """ - INSERT INTO table1 ( - bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, - integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, - timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, - mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol, jsonCol - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """.trimIndent() - - @Language("SQL") - val insertData2 = - """ - INSERT INTO table2 ( - bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, - integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, - timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, - mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """.trimIndent() - - connection.prepareStatement(insertData1).use { st -> - // Insert data into table1 - for (i in 1..3) { - st.setBoolean(1, true) - st.setByte(2, i.toByte()) - st.setShort(3, (i * 10).toShort()) - st.setInt(4, i * 100) - st.setInt(5, i * 100) - st.setInt(6, i * 100) - st.setInt(7, i * 100) - st.setInt(8, i * 100) - st.setInt(9, i * 100) - st.setFloat(10, i * 10.0f) - st.setDouble(11, i * 10.0) - st.setBigDecimal(12, BigDecimal(i * 10)) - st.setDate(13, java.sql.Date(System.currentTimeMillis())) - st.setTimestamp(14, java.sql.Timestamp(System.currentTimeMillis())) - st.setTimestamp(15, java.sql.Timestamp(System.currentTimeMillis())) - st.setTime(16, java.sql.Time(System.currentTimeMillis())) - st.setInt(17, 2023) - st.setString(18, "varcharValue$i") - st.setString(19, "charValue$i") - st.setBytes(20, "binaryValue".toByteArray()) - st.setBytes(21, "varbinaryValue".toByteArray()) - st.setBytes(22, "tinyblobValue".toByteArray()) - st.setBytes(23, "blobValue".toByteArray()) - st.setBytes(24, "mediumblobValue".toByteArray()) - st.setBytes(25, "longblobValue".toByteArray()) - st.setString(26, "textValue$i") - st.setString(27, "mediumtextValue$i") - st.setString(28, "longtextValue$i") - st.setString(29, "Value$i") - st.setString(30, "Option$i") - st.setObject(31, BigInteger.valueOf((i * 1000).toLong())) - st.setString(32, JSON_STRING) - - st.executeUpdate() - } - } - - connection.prepareStatement(insertData2).use { st -> - // Insert data into table2 - for (i in 1..3) { - st.setBoolean(1, false) - st.setByte(2, (i * 2).toByte()) - st.setShort(3, (i * 20).toShort()) - st.setInt(4, i * 200) - st.setInt(5, i * 200) - st.setInt(6, i * 200) - st.setInt(7, i * 200) - st.setInt(8, i * 200) - st.setInt(9, i * 200) - st.setFloat(10, i * 20.0f) - st.setDouble(11, i * 20.0) - st.setBigDecimal(12, BigDecimal(i * 20)) - st.setDate(13, java.sql.Date(System.currentTimeMillis())) - st.setTimestamp(14, java.sql.Timestamp(System.currentTimeMillis())) - st.setTimestamp(15, java.sql.Timestamp(System.currentTimeMillis())) - st.setTime(16, java.sql.Time(System.currentTimeMillis())) - st.setInt(17, 2023) - st.setString(18, "varcharValue$i") - st.setString(19, "charValue$i") - st.setBytes(20, "binaryValue".toByteArray()) - st.setBytes(21, "varbinaryValue".toByteArray()) - st.setBytes(22, "tinyblobValue".toByteArray()) - st.setBytes(23, "blobValue".toByteArray()) - st.setBytes(24, "mediumblobValue".toByteArray()) - st.setBytes(25, "longblobValue".toByteArray()) - st.setString(26, null) - st.setString(27, null) - st.setString(28, "longtextValue$i") - st.setString(29, "Value$i") - st.setString(30, "Option$i") - st.setObject(31, BigInteger.valueOf((i * 2000).toLong())) - st.executeUpdate() - } - } + connection = openConnection(null) + setUpMariadbTestData(connection) } @AfterClass @JvmStatic fun tearDownClass() { - try { - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } - connection.createStatement().use { st -> st.execute("DROP DATABASE IF EXISTS $TEST_DATABASE_NAME") } - connection.close() - } catch (e: SQLException) { - e.printStackTrace() - } - } - } - - @Test - fun `basic test for reading sql tables`() { - val df1 = DataFrame.readSqlTable(connection, "table1").cast() - val result = df1.filter { "id"() == 1 } - result[0][26] shouldBe "textValue1" - val byteArray = "tinyblobValue".toByteArray() - result[0][22] shouldBe byteArray - - val schema = DataFrameSchema.readSqlTable(connection, "table1") - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["textCol"]!!.type shouldBe typeOf() - schema.columns["varbinaryCol"]!!.type shouldBe typeOf() - schema.columns["binaryCol"]!!.type shouldBe typeOf() - schema.columns["longblobCol"]!!.type shouldBe typeOf() - schema.columns["tinyblobCol"]!!.type shouldBe typeOf() - schema.columns["dateCol"]!!.type shouldBe typeOf() - schema.columns["datetimeCol"]!!.type shouldBe typeOf() - schema.columns["timestampCol"]!!.type shouldBe typeOf() - schema.columns["timeCol"]!!.type shouldBe typeOf() - schema.columns["yearCol"]!!.type shouldBe typeOf() - - val df2 = DataFrame.readSqlTable(connection, "table2").cast() - val result2 = df2.filter { "id"() == 1 } - result2[0][26] shouldBe null - - val schema2 = DataFrameSchema.readSqlTable(connection, "table2") - schema2.columns["id"]!!.type shouldBe typeOf() - schema2.columns["textCol"]!!.type shouldBe typeOf() - } - - @Test - fun `read from sql query`() { - @Language("SQL") - val sqlQuery = - """ - SELECT - t1.id, - t1.enumCol, - t2.setCol - FROM table1 t1 - JOIN table2 t2 ON t1.id = t2.id - """.trimIndent() - - val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery).cast() - val result = df.filter { "id"() == 1 } - result[0][2] shouldBe "Option1" - - val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["enumCol"]!!.type shouldBe typeOf() - schema.columns["setCol"]!!.type shouldBe typeOf() - } - - @Test - fun `read from all tables`() { - val dataframes = DataFrame.readAllSqlTables(connection, TEST_DATABASE_NAME, 1000).values.toList() - - val table1Df = dataframes[0].cast() - - table1Df.rowsCount() shouldBe 3 - table1Df.filter { "integerCol"() > 100 }.rowsCount() shouldBe 2 - table1Df[0][11] shouldBe 10.0 - table1Df[0][26] shouldBe "textValue1" - table1Df[0][31] shouldBe BigInteger.valueOf(1000L) - table1Df[0][32] shouldBe JSON_STRING // TODO: https://github.com/Kotlin/dataframe/issues/462 - - val table2Df = dataframes[1].cast() - - table2Df.rowsCount() shouldBe 3 - table2Df.filter { - "integerCol"()?.let { it > 400 } ?: false - }.rowsCount() shouldBe 1 - table2Df[0][11] shouldBe 20.0 - table2Df[0][26] shouldBe null - } - - @Test - fun `reading numeric types`() { - val df1 = DataFrame.readSqlTable(connection, "table1").cast() - - val result = df1.select("tinyintCol") - .add("tinyintCol2") { "tinyintCol"() } - - result[0][1] shouldBe 1 - - val result1 = df1.select("smallintCol") - .add("smallintCol2") { "smallintCol"() } - - result1[0][1] shouldBe 10 - - val result2 = df1.select("mediumintCol") - .add("mediumintCol2") { "mediumintCol"() } - - result2[0][1] shouldBe 100 - - val result3 = df1.select("mediumintUnsignedCol") - .add("mediumintUnsignedCol2") { "mediumintUnsignedCol"() } - - result3[0][1] shouldBe 100 - - val result4 = df1.select("integerUnsignedCol") - .add("integerUnsignedCol2") { "integerUnsignedCol"() } - - result4[0][1] shouldBe 100L - - val result5 = df1.select("bigintCol") - .add("bigintCol2") { "bigintCol"() } - - result5[0][1] shouldBe 100 - - val result5a = df1.select("bigintUnsignedCol") - .add("bigintUnsignedCol2") { "bigintUnsignedCol"() } - - result5a[0][1] shouldBe BigInteger.valueOf(1000) - - val result6 = df1.select("floatCol") - .add("floatCol2") { "floatCol"() } - - result6[0][1] shouldBe 10.0f - - val result7 = df1.select("doubleCol") - .add("doubleCol2") { "doubleCol"() } - - result7[0][1] shouldBe 10.0 - - val result8 = df1.select("decimalCol") - .add("decimalCol2") { "decimalCol"() } - - result8[0][1] shouldBe BigDecimal("10") - - val schema = DataFrameSchema.readSqlTable(connection, "table1") - - schema.columns["tinyintCol"]!!.type shouldBe typeOf() - schema.columns["smallintCol"]!!.type shouldBe typeOf() - schema.columns["mediumintCol"]!!.type shouldBe typeOf() - schema.columns["mediumintUnsignedCol"]!!.type shouldBe typeOf() - schema.columns["integerUnsignedCol"]!!.type shouldBe typeOf() - schema.columns["bigintCol"]!!.type shouldBe typeOf() - schema.columns["bigintUnsignedCol"]!!.type shouldBe typeOf() - schema.columns["floatCol"]!!.type shouldBe typeOf() - schema.columns["doubleCol"]!!.type shouldBe typeOf() - schema.columns["decimalCol"]!!.type shouldBe typeOf() - } - - @Test - fun `infer nullability`() { - inferNullability(connection) - } - - // https://github.com/Kotlin/dataframe/issues/1746 - @Test - fun `readAllSqlTables without catalogue should only return tables from URL database`() { - val secondDb = "testKDFdatabase2" - val testRootConn = DriverManager.getConnection(URL, USER_NAME, PASSWORD) - try { - testRootConn.createStatement().use { stmt -> - stmt.executeUpdate("DROP DATABASE IF EXISTS $secondDb") - stmt.executeUpdate("CREATE DATABASE $secondDb") - } - DriverManager.getConnection("$URL/$secondDb", USER_NAME, PASSWORD).use { conn2 -> - conn2.createStatement().use { stmt -> - stmt.executeUpdate("CREATE TABLE onlyInDb2 (id INT PRIMARY KEY, val VARCHAR(50))") - } - } - - DriverManager.getConnection("$URL/$TEST_DATABASE_NAME", USER_NAME, PASSWORD).use { scopedConn -> - val tableNames = DataFrame.readAllSqlTables(scopedConn).keys - - tableNames.none { "onlyInDb2" in it } shouldBe true - tableNames.any { "table1" in it } shouldBe true - } - } finally { - testRootConn.use { conn -> - conn.createStatement().execute("DROP DATABASE IF EXISTS $secondDb") - } + tearDownMariadbTestData(connection) } } } diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mssqlTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mssqlTest.kt index 7c3b7cdbef..174f0776f5 100644 --- a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mssqlTest.kt +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mssqlTest.kt @@ -1,75 +1,27 @@ package org.jetbrains.kotlinx.dataframe.io.local -import io.kotest.matchers.shouldBe -import org.intellij.lang.annotations.Language -import org.jetbrains.kotlinx.dataframe.DataFrame -import org.jetbrains.kotlinx.dataframe.annotations.DataSchema -import org.jetbrains.kotlinx.dataframe.api.cast -import org.jetbrains.kotlinx.dataframe.api.filter -import org.jetbrains.kotlinx.dataframe.io.inferNullability -import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables -import org.jetbrains.kotlinx.dataframe.io.readSqlQuery -import org.jetbrains.kotlinx.dataframe.io.readSqlTable -import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.jetbrains.kotlinx.dataframe.io.MsSqlTestBase +import org.jetbrains.kotlinx.dataframe.io.setUpMsSqlTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownMsSqlTestData import org.junit.AfterClass import org.junit.BeforeClass -import org.junit.Ignore -import org.junit.Test -import java.math.BigDecimal import java.sql.Connection import java.sql.DriverManager -import java.sql.SQLException -import java.util.Date -import java.util.UUID -import kotlin.reflect.typeOf -import kotlin.time.Instant private const val URL = "jdbc:sqlserver://localhost:1433;encrypt=true;trustServerCertificate=true" private const val USER_NAME = "root" private const val PASSWORD = "pass" -private const val TEST_DATABASE_NAME = "testKDFdatabase" -@DataSchema -interface Table1MSSSQL { - val id: Int - val bigintColumn: Long - val binaryColumn: ByteArray - val bitColumn: Boolean - val charColumn: Char - val dateColumn: Date - val datetime3Column: Instant - val datetime2Column: Instant - val datetimeoffset2Column: String - val decimalColumn: BigDecimal - val floatColumn: Double - val imageColumn: ByteArray? - val intColumn: Int - val moneyColumn: BigDecimal - val ncharColumn: Char - val ntextColumn: String - val numericColumn: BigDecimal - val nvarcharColumn: String - val nvarcharMaxColumn: String - val realColumn: Float - val smalldatetimeColumn: Instant - val smallintColumn: Int - val smallmoneyColumn: BigDecimal - val timeColumn: java.sql.Time - val timestampColumn: Instant - val tinyintColumn: Int - val uniqueidentifierColumn: Char - val varbinaryColumn: ByteArray - val varbinaryMaxColumn: ByteArray - val varcharColumn: String - val varcharMaxColumn: String - val xmlColumn: String - val sqlvariantColumn: String - val geometryColumn: ByteArray - val geographyColumn: ByteArray -} +class MsSqlLocalTest : MsSqlTestBase() { + override val connection: Connection get() = Companion.connection + + override fun connect(database: String?): Connection = + DriverManager.getConnection( + if (database == null) URL else "$URL;databaseName=$database", + USER_NAME, + PASSWORD, + ) -@Ignore -class MSSQLTest { companion object { private lateinit var connection: Connection @@ -77,260 +29,13 @@ class MSSQLTest { @JvmStatic fun setUpClass() { connection = DriverManager.getConnection(URL, USER_NAME, PASSWORD) - - connection.createStatement().use { st -> - // Drop the test database if it exists - val dropDatabaseQuery = "IF DB_ID('$TEST_DATABASE_NAME') IS NOT NULL\n" + - "DROP DATABASE $TEST_DATABASE_NAME" - st.executeUpdate(dropDatabaseQuery) - - // Create the test database - val createDatabaseQuery = "CREATE DATABASE $TEST_DATABASE_NAME" - st.executeUpdate(createDatabaseQuery) - - // Use the newly created database - val useDatabaseQuery = "USE $TEST_DATABASE_NAME" - st.executeUpdate(useDatabaseQuery) - } - - @Language("SQL") - val createTableQuery = """ - CREATE TABLE Table1 ( - id INT NOT NULL IDENTITY PRIMARY KEY, - bigintColumn BIGINT, - binaryColumn BINARY(50), - bitColumn BIT, - charColumn CHAR(10), - dateColumn DATE, - datetime3Column DATETIME2(3), - datetime2Column DATETIME2, - datetimeoffset2Column DATETIMEOFFSET(2), - decimalColumn DECIMAL(10,2), - floatColumn FLOAT, - imageColumn IMAGE, - intColumn INT, - moneyColumn MONEY, - ncharColumn NCHAR(10), - ntextColumn NTEXT, - numericColumn NUMERIC(10,2), - nvarcharColumn NVARCHAR(50), - nvarcharMaxColumn NVARCHAR(MAX), - realColumn REAL, - smalldatetimeColumn SMALLDATETIME, - smallintColumn SMALLINT, - smallmoneyColumn SMALLMONEY, - textColumn TEXT, - timeColumn TIME, - timestampColumn DATETIME2, - tinyintColumn TINYINT, - uniqueidentifierColumn UNIQUEIDENTIFIER, - varbinaryColumn VARBINARY(50), - varbinaryMaxColumn VARBINARY(MAX), - varcharColumn VARCHAR(50), - varcharMaxColumn VARCHAR(MAX), - xmlColumn XML, - sqlvariantColumn SQL_VARIANT, - geometryColumn GEOMETRY, - geographyColumn GEOGRAPHY - ); - """ - - connection.createStatement().execute(createTableQuery.trimIndent()) - - @Language("SQL") - val insertData1 = - """ - INSERT INTO Table1 ( - bigintColumn, binaryColumn, bitColumn, charColumn, dateColumn, datetime3Column, datetime2Column, - datetimeoffset2Column, decimalColumn, floatColumn, imageColumn, intColumn, moneyColumn, ncharColumn, - ntextColumn, numericColumn, nvarcharColumn, nvarcharMaxColumn, realColumn, smalldatetimeColumn, - smallintColumn, smallmoneyColumn, textColumn, timeColumn, timestampColumn, tinyintColumn, - uniqueidentifierColumn, varbinaryColumn, varbinaryMaxColumn, varcharColumn, varcharMaxColumn, - xmlColumn, sqlvariantColumn, geometryColumn, geographyColumn - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """.trimIndent() - - connection.prepareStatement(insertData1).use { st -> - for (i in 1..5) { - st.setLong(1, 123456789012345L) // bigintColumn - st.setBytes(2, byteArrayOf(0x01, 0x23, 0x45, 0x67, 0x67, 0x67, 0x67, 0x67)) // binaryColumn - st.setBoolean(3, true) // bitColumn - st.setString(4, "Sample") // charColumn - st.setDate(5, java.sql.Date(System.currentTimeMillis())) // dateColumn - st.setTimestamp(6, java.sql.Timestamp(System.currentTimeMillis())) // datetime3Column - st.setTimestamp(7, java.sql.Timestamp(System.currentTimeMillis())) // datetime2Column - st.setTimestamp(8, java.sql.Timestamp(System.currentTimeMillis())) // datetimeoffset2Column - st.setBigDecimal(9, BigDecimal("12345.67")) // decimalColumn - st.setFloat(10, 123.45f) // floatColumn - st.setNull(11, java.sql.Types.NULL) // imageColumn (assuming nullable) - st.setInt(12, 123456) // intColumn - st.setBigDecimal(13, BigDecimal("123.45")) // moneyColumn - st.setString(14, "Sample") // ncharColumn - st.setString(15, "Sample$i text") // ntextColumn - st.setBigDecimal(16, BigDecimal("1234.56")) // numericColumn - st.setString(17, "Sample") // nvarcharColumn - st.setString(18, "Sample$i text") // nvarcharMaxColumn - st.setFloat(19, 123.45f) // realColumn - st.setTimestamp(20, java.sql.Timestamp(System.currentTimeMillis())) // smalldatetimeColumn - st.setInt(21, 123) // smallintColumn - st.setBigDecimal(22, BigDecimal("123.45")) // smallmoneyColumn - st.setString(23, "Sample$i text") // textColumn - st.setTime(24, java.sql.Time(System.currentTimeMillis())) // timeColumn - st.setTimestamp(25, java.sql.Timestamp(System.currentTimeMillis())) // timestampColumn - st.setInt(26, 123) // tinyintColumn - // st.setObject(27, null) // udtColumn (assuming nullable) - st.setObject(27, UUID.randomUUID()) // uniqueidentifierColumn - st.setBytes(28, byteArrayOf(0x01, 0x23, 0x45, 0x67, 0x67, 0x67, 0x67, 0x67)) // varbinaryColumn - st.setBytes(29, byteArrayOf(0x01, 0x23, 0x45, 0x67, 0x67, 0x67, 0x67, 0x67)) // varbinaryMaxColumn - st.setString(30, "Sample$i") // varcharColumn - st.setString(31, "Sample$i text") // varcharMaxColumn - st.setString(32, "Sample$i") // xmlColumn - st.setString(33, "SQL_VARIANT") // sqlvariantColumn - st.setBytes( - 34, - @Suppress("ktlint:standard:argument-list-wrapping") - byteArrayOf( - 0xE6.toByte(), 0x10, 0x00, 0x00, 0x01, 0x0C, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, - 0x44, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x05, 0x4C, 0x0, - ), - ) // geometryColumn - st.setString(35, "POINT(1 1)") // geographyColumn - st.executeUpdate() - } - } + setUpMsSqlTestData(connection) } @AfterClass @JvmStatic fun tearDownClass() { - try { - connection.createStatement().use { st -> st.execute("DROP DATABASE IF EXISTS $TEST_DATABASE_NAME") } - connection.close() - } catch (e: SQLException) { - e.printStackTrace() - } - } - } - - @Test - fun `basic test for reading sql tables`() { - val df1 = DataFrame.readSqlTable(connection, "table1", limit = 5).cast() - - val result = df1.filter { "id"() == 1 } - result[0][30] shouldBe "Sample1" - result[0]["bigintColumn"] shouldBe 123456789012345L - result[0]["bitColumn"] shouldBe true - result[0]["intColumn"] shouldBe 123456 - result[0]["ntextColumn"] shouldBe "Sample1 text" - - val schema = DataFrameSchema.readSqlTable(connection, "table1") - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["bigintColumn"]!!.type shouldBe typeOf() - schema.columns["binaryColumn"]!!.type shouldBe typeOf() - schema.columns["bitColumn"]!!.type shouldBe typeOf() - schema.columns["charColumn"]!!.type shouldBe typeOf() - schema.columns["dateColumn"]!!.type shouldBe typeOf() - schema.columns["datetime3Column"]!!.type shouldBe typeOf() - schema.columns["datetime2Column"]!!.type shouldBe typeOf() - schema.columns["datetimeoffset2Column"]!!.type shouldBe typeOf() - schema.columns["decimalColumn"]!!.type shouldBe typeOf() - schema.columns["floatColumn"]!!.type shouldBe typeOf() - schema.columns["imageColumn"]!!.type shouldBe typeOf() - schema.columns["intColumn"]!!.type shouldBe typeOf() - schema.columns["moneyColumn"]!!.type shouldBe typeOf() - schema.columns["ncharColumn"]!!.type shouldBe typeOf() - schema.columns["ntextColumn"]!!.type shouldBe typeOf() - schema.columns["numericColumn"]!!.type shouldBe typeOf() - schema.columns["nvarcharColumn"]!!.type shouldBe typeOf() - schema.columns["nvarcharMaxColumn"]!!.type shouldBe typeOf() - schema.columns["realColumn"]!!.type shouldBe typeOf() - schema.columns["smalldatetimeColumn"]!!.type shouldBe typeOf() - schema.columns["smallintColumn"]!!.type shouldBe typeOf() - schema.columns["smallmoneyColumn"]!!.type shouldBe typeOf() - schema.columns["timeColumn"]!!.type shouldBe typeOf() - schema.columns["timestampColumn"]!!.type shouldBe typeOf() - schema.columns["tinyintColumn"]!!.type shouldBe typeOf() - schema.columns["uniqueidentifierColumn"]!!.type shouldBe typeOf() - schema.columns["varbinaryColumn"]!!.type shouldBe typeOf() - schema.columns["varbinaryMaxColumn"]!!.type shouldBe typeOf() - schema.columns["varcharColumn"]!!.type shouldBe typeOf() - schema.columns["varcharMaxColumn"]!!.type shouldBe typeOf() - schema.columns["xmlColumn"]!!.type shouldBe typeOf() - schema.columns["sqlvariantColumn"]!!.type shouldBe typeOf() - schema.columns["geometryColumn"]!!.type shouldBe typeOf() - schema.columns["geographyColumn"]!!.type shouldBe typeOf() - } - - @Test - fun `read from sql query`() { - @Language("SQL") - val sqlQuery = - """ - SELECT - Table1.id, - Table1.bigintColumn - FROM Table1 - """.trimIndent() - - val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery, limit = 3).cast() - val result = df.filter { "id"() == 1 } - result[0]["bigintColumn"] shouldBe 123456789012345L - - val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["bigintColumn"]!!.type shouldBe typeOf() - } - - @Test - fun `read from all tables`() { - val dataframes = DataFrame.readAllSqlTables(connection, TEST_DATABASE_NAME, 4).values.toList() - - val table1Df = dataframes[0].cast() - - table1Df.rowsCount() shouldBe 4 - table1Df.filter { "id"() > 2 }.rowsCount() shouldBe 2 - table1Df[0]["bigintColumn"] shouldBe 123456789012345L - } - - @Test - fun `infer nullability`() { - inferNullability(connection) - } - - // https://github.com/Kotlin/dataframe/issues/1746 - @Test - fun `readAllSqlTables without catalogue should only return tables from URL database`() { - val secondDb = "testKDFdatabase2" - val testRootConn = DriverManager.getConnection(URL, USER_NAME, PASSWORD) - try { - testRootConn.createStatement().use { stmt -> - stmt.executeUpdate( - "IF DB_ID('$secondDb') IS NOT NULL DROP DATABASE $secondDb", - ) - stmt.executeUpdate("CREATE DATABASE $secondDb") - } - DriverManager.getConnection("$URL;databaseName=$secondDb", USER_NAME, PASSWORD).use { conn2 -> - conn2.createStatement().use { stmt -> - stmt.executeUpdate("CREATE TABLE onlyInDb2 (id INT PRIMARY KEY, val VARCHAR(50))") - } - } - - DriverManager.getConnection( - "$URL;databaseName=$TEST_DATABASE_NAME", - USER_NAME, - PASSWORD, - ).use { scopedConn -> - val tableNames = DataFrame.readAllSqlTables(scopedConn).keys - - tableNames.none { "onlyInDb2" in it } shouldBe true - tableNames.any { "Table1" in it } shouldBe true - } - } finally { - testRootConn.use { conn -> - conn.createStatement().execute( - "IF DB_ID('$secondDb') IS NOT NULL DROP DATABASE $secondDb", - ) - } + tearDownMsSqlTestData(connection) } } } diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mysqlTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mysqlTest.kt index cadbc8580b..62be4ac5c9 100644 --- a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mysqlTest.kt +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/mysqlTest.kt @@ -1,523 +1,43 @@ package org.jetbrains.kotlinx.dataframe.io.local -import io.kotest.matchers.shouldBe -import kotlinx.datetime.LocalDateTime -import org.intellij.lang.annotations.Language -import org.jetbrains.kotlinx.dataframe.DataFrame -import org.jetbrains.kotlinx.dataframe.annotations.DataSchema -import org.jetbrains.kotlinx.dataframe.api.add -import org.jetbrains.kotlinx.dataframe.api.cast -import org.jetbrains.kotlinx.dataframe.api.filter -import org.jetbrains.kotlinx.dataframe.api.select -import org.jetbrains.kotlinx.dataframe.io.inferNullability -import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables -import org.jetbrains.kotlinx.dataframe.io.readSqlQuery -import org.jetbrains.kotlinx.dataframe.io.readSqlTable -import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.jetbrains.kotlinx.dataframe.io.MySqlTestBase +import org.jetbrains.kotlinx.dataframe.io.setUpMySqlTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownMySqlTestData import org.junit.AfterClass import org.junit.BeforeClass -import org.junit.Ignore -import org.junit.Test -import java.math.BigDecimal -import java.math.BigInteger import java.sql.Connection import java.sql.DriverManager -import java.sql.SQLException -import java.util.Date -import kotlin.reflect.typeOf -import kotlin.time.Instant private const val URL = "jdbc:mysql://localhost:3307" private const val USER_NAME = "root" private const val PASSWORD = "pass" -private const val TEST_DATABASE_NAME = "testKDFdatabase" -private const val TIMESTAMP = 1726246245460 -@DataSchema -interface Table1MySql { - val id: Int - val bitCol: Boolean - val tinyintCol: Int - val smallintCol: Int - val mediumintCol: Int - val mediumintUnsignedCol: Int - val integerCol: Int - val intCol: Int - val integerUnsignedCol: Long - val bigintCol: Long - val floatCol: Float - val doubleCol: Double - val decimalCol: BigDecimal - val dateCol: String - val datetimeCol: String - val timestampCol: String - val timeCol: String - val yearCol: String - val varcharCol: String - val charCol: String - val binaryCol: ByteArray - val varbinaryCol: ByteArray - val tinyblobCol: ByteArray - val blobCol: ByteArray - val mediumblobCol: ByteArray - val longblobCol: ByteArray - val textCol: String - val mediumtextCol: String - val longtextCol: String - val enumCol: String - val setCol: Char - val bigintUnsignedCol: BigInteger -} +class MySqlLocalTest : MySqlTestBase() { + override val connection: Connection get() = Companion.connection -@DataSchema -interface Table2MySql { - val id: Int - val bitCol: Boolean? - val tinyintCol: Int? - val smallintCol: Int? - val mediumintCol: Int? - val mediumintUnsignedCol: Int? - val integerCol: Int? - val intCol: Int? - val integerUnsignedCol: Long? - val bigintCol: Long? - val floatCol: Float? - val doubleCol: Double? - val decimalCol: Double? - val dateCol: String? - val datetimeCol: String? - val timestampCol: String? - val timeCol: String? - val yearCol: String? - val varcharCol: String? - val charCol: String? - val binaryCol: ByteArray? - val varbinaryCol: ByteArray? - val tinyblobCol: ByteArray? - val blobCol: ByteArray? - val mediumblobCol: ByteArray? - val longblobCol: ByteArray? - val textCol: String? - val mediumtextCol: String? - val longtextCol: String? - val enumCol: String? - val setCol: Char? - val bigintUnsignedCol: BigInteger? - val jsonCol: String? -} + override fun connect(database: String?): Connection = openConnection(database) -@DataSchema -interface Table3MySql { - val id: Int - val enumCol: String - val setCol: Char? -} - -@Ignore -class MySqlTest { companion object { private lateinit var connection: Connection + private fun openConnection(database: String?): Connection = + DriverManager.getConnection( + if (database == null) URL else "$URL/$database", + USER_NAME, + PASSWORD, + ) + @BeforeClass @JvmStatic fun setUpClass() { - connection = DriverManager.getConnection(URL, USER_NAME, PASSWORD) - - connection.createStatement().use { st -> - // Drop the test database if it exists - val dropDatabaseQuery = "DROP DATABASE IF EXISTS $TEST_DATABASE_NAME" - st.executeUpdate(dropDatabaseQuery) - - // Create the test database - val createDatabaseQuery = "CREATE DATABASE $TEST_DATABASE_NAME" - st.executeUpdate(createDatabaseQuery) - - // Use the newly created database - val useDatabaseQuery = "USE $TEST_DATABASE_NAME" - st.executeUpdate(useDatabaseQuery) - } - - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } - - @Language("SQL") - val createTableQuery = """ - CREATE TABLE IF NOT EXISTS table1 ( - id INT AUTO_INCREMENT PRIMARY KEY, - bitCol BIT NOT NULL, - tinyintCol TINYINT NOT NULL, - smallintCol SMALLINT NOT NULL, - mediumintCol MEDIUMINT NOT NULL, - mediumintUnsignedCol MEDIUMINT UNSIGNED NOT NULL, - integerCol INTEGER NOT NULL, - intCol INT NOT NULL, - integerUnsignedCol INTEGER UNSIGNED NOT NULL, - bigintCol BIGINT NOT NULL, - floatCol FLOAT NOT NULL, - doubleCol DOUBLE NOT NULL, - decimalCol DECIMAL NOT NULL, - dateCol DATE NOT NULL, - datetimeCol DATETIME NOT NULL, - timestampCol TIMESTAMP NOT NULL, - timeCol TIME NOT NULL, - yearCol YEAR NOT NULL, - varcharCol VARCHAR(255) NOT NULL, - charCol CHAR(10) NOT NULL, - binaryCol BINARY(64) NOT NULL, - varbinaryCol VARBINARY(128) NOT NULL, - tinyblobCol TINYBLOB NOT NULL, - blobCol BLOB NOT NULL, - mediumblobCol MEDIUMBLOB NOT NULL , - longblobCol LONGBLOB NOT NULL, - textCol TEXT NOT NULL, - mediumtextCol MEDIUMTEXT NOT NULL, - longtextCol LONGTEXT NOT NULL, - enumCol ENUM('Value1', 'Value2', 'Value3') NOT NULL, - setCol SET('Option1', 'Option2', 'Option3') NOT NULL, - bigintUnsignedCol BIGINT UNSIGNED NOT NULL, - location GEOMETRY, - data JSON - CHECK (JSON_VALID(data)) - ) - """ - - connection.createStatement().execute(createTableQuery.trimIndent()) - - @Language("SQL") - val createTableQuery2 = """ - CREATE TABLE IF NOT EXISTS table2 ( - id INT AUTO_INCREMENT PRIMARY KEY, - bitCol BIT, - tinyintCol TINYINT, - smallintCol SMALLINT, - mediumintCol MEDIUMINT, - mediumintUnsignedCol MEDIUMINT UNSIGNED, - integerCol INTEGER, - intCol INT, - integerUnsignedCol INTEGER UNSIGNED, - bigintCol BIGINT, - floatCol FLOAT, - doubleCol DOUBLE, - decimalCol DECIMAL, - dateCol DATE, - datetimeCol DATETIME, - timestampCol TIMESTAMP, - timeCol TIME, - yearCol YEAR, - varcharCol VARCHAR(255), - charCol CHAR(10), - binaryCol BINARY(64), - varbinaryCol VARBINARY(128), - tinyblobCol TINYBLOB, - blobCol BLOB, - mediumblobCol MEDIUMBLOB, - longblobCol LONGBLOB, - textCol TEXT, - mediumtextCol MEDIUMTEXT, - longtextCol LONGTEXT, - enumCol ENUM('Value1', 'Value2', 'Value3'), - setCol SET('Option1', 'Option2', 'Option3'), - bigintUnsignedCol BIGINT UNSIGNED, - location GEOMETRY, - data JSON - CHECK (JSON_VALID(data)) - ) - """ - - connection.createStatement().execute(createTableQuery2.trimIndent()) - - @Language("SQL") - val insertData1 = - """ - INSERT INTO table1 ( - bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, - integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, - timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, - mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol, location, data - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ST_GeomFromText('POINT(1 1)'), ?) - """.trimIndent() - - @Language("SQL") - val insertData2 = - """ - INSERT INTO table2 ( - bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, - integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, - timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, - mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol, location, data - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ST_GeomFromText('POINT(1 1)'), ?) - """.trimIndent() - - connection.prepareStatement(insertData1).use { st -> - // Insert data into table1 - for (i in 1..3) { - st.setBoolean(1, true) - st.setByte(2, i.toByte()) - st.setShort(3, (i * 10).toShort()) - st.setInt(4, i * 100) - st.setInt(5, i * 100) - st.setInt(6, i * 100) - st.setInt(7, i * 100) - st.setInt(8, i * 100) - st.setInt(9, i * 100) - st.setFloat(10, i * 10.0f) - st.setDouble(11, i * 10.0) - st.setBigDecimal(12, BigDecimal(i * 10)) - st.setDate(13, java.sql.Date(TIMESTAMP)) - st.setTimestamp(14, java.sql.Timestamp(TIMESTAMP)) - st.setTimestamp(15, java.sql.Timestamp(TIMESTAMP)) - st.setTime(16, java.sql.Time(TIMESTAMP)) - st.setInt(17, 2023) - st.setString(18, "varcharValue$i") - st.setString(19, "charValue$i") - st.setBytes(20, "binaryValue".toByteArray()) - st.setBytes(21, "varbinaryValue".toByteArray()) - st.setBytes(22, "tinyblobValue".toByteArray()) - st.setBytes(23, "blobValue".toByteArray()) - st.setBytes(24, "mediumblobValue".toByteArray()) - st.setBytes(25, "longblobValue".toByteArray()) - st.setString(26, "textValue$i") - st.setString(27, "mediumtextValue$i") - st.setString(28, "longtextValue$i") - st.setString(29, "Value$i") - st.setString(30, "Option$i") - st.setObject(31, BigInteger.valueOf((i * 1000).toLong())) - st.setString(32, "{\"key\": \"value\"}") - st.executeUpdate() - } - } - - connection.prepareStatement(insertData2).use { st -> - // Insert data into table2 - for (i in 1..3) { - st.setBoolean(1, false) - st.setByte(2, (i * 2).toByte()) - st.setShort(3, (i * 20).toShort()) - st.setInt(4, i * 200) - st.setInt(5, i * 200) - st.setInt(6, i * 200) - st.setInt(7, i * 200) - st.setInt(8, i * 200) - st.setInt(9, i * 200) - st.setFloat(10, i * 20.0f) - st.setDouble(11, i * 20.0) - st.setBigDecimal(12, BigDecimal(i * 20)) - st.setDate(13, java.sql.Date(TIMESTAMP)) - st.setTimestamp(14, java.sql.Timestamp(TIMESTAMP)) - st.setTimestamp(15, java.sql.Timestamp(TIMESTAMP)) - st.setTime(16, java.sql.Time(TIMESTAMP)) - st.setInt(17, 2023) - st.setString(18, "varcharValue$i") - st.setString(19, "charValue$i") - st.setBytes(20, "binaryValue".toByteArray()) - st.setBytes(21, "varbinaryValue".toByteArray()) - st.setBytes(22, "tinyblobValue".toByteArray()) - st.setBytes(23, "blobValue".toByteArray()) - st.setBytes(24, "mediumblobValue".toByteArray()) - st.setBytes(25, "longblobValue".toByteArray()) - st.setString(26, null) - st.setString(27, null) - st.setString(28, "longtextValue$i") - st.setString(29, "Value$i") - st.setString(30, "Option$i") - st.setObject(31, BigInteger.valueOf((i * 2000).toLong())) - st.setString(32, "{\"key\": \"value\"}") - st.executeUpdate() - } - } + connection = openConnection(null) + setUpMySqlTestData(connection) } @AfterClass @JvmStatic fun tearDownClass() { - try { - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } - connection.createStatement().use { st -> st.execute("DROP DATABASE IF EXISTS $TEST_DATABASE_NAME") } - connection.close() - } catch (e: SQLException) { - e.printStackTrace() - } - } - } - - @Test - fun `basic test for reading sql tables`() { - val df1 = DataFrame.readSqlTable(connection, "table1").cast() - val result = df1.filter { "id"() == 1 } - result[0][26] shouldBe "textValue1" - result[0][22] shouldBe "tinyblobValue".toByteArray() - - val schema = DataFrameSchema.readSqlTable(connection, "table1") - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["textCol"]!!.type shouldBe typeOf() - schema.columns["dateCol"]!!.type shouldBe typeOf() - schema.columns["datetimeCol"]!!.type shouldBe typeOf() - schema.columns["timestampCol"]!!.type shouldBe typeOf() - schema.columns["timeCol"]!!.type shouldBe typeOf() - schema.columns["yearCol"]!!.type shouldBe typeOf() - schema.columns["textCol"]!!.type shouldBe typeOf() - schema.columns["varbinaryCol"]!!.type shouldBe typeOf() - schema.columns["binaryCol"]!!.type shouldBe typeOf() - schema.columns["longblobCol"]!!.type shouldBe typeOf() - schema.columns["tinyblobCol"]!!.type shouldBe typeOf() - - val df2 = DataFrame.readSqlTable(connection, "table2").cast() - val result2 = df2.filter { "id"() == 1 } - result2[0][26] shouldBe null - - val schema2 = DataFrameSchema.readSqlTable(connection, "table2") - schema2.columns["id"]!!.type shouldBe typeOf() - schema2.columns["textCol"]!!.type shouldBe typeOf() - } - - @Test - fun `read from sql query`() { - @Language("SQL") - val sqlQuery = - """ - SELECT - t1.id, - t1.enumCol, - t2.setCol - FROM table1 t1 - JOIN table2 t2 ON t1.id = t2.id - """.trimIndent() - - val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery).cast() - val result = df.filter { "id"() == 1 } - result[0][2] shouldBe "Option1" - - val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["enumCol"]!!.type shouldBe typeOf() - schema.columns["setCol"]!!.type shouldBe typeOf() - } - - @Test - fun `read from all tables`() { - val dataframes = DataFrame.readAllSqlTables(connection).values.toList() - - val table1Df = dataframes[0].cast() - - table1Df.rowsCount() shouldBe 3 - table1Df.filter { "integerCol"() > 100 }.rowsCount() shouldBe 2 - table1Df[0][11] shouldBe 10.0 - table1Df[0][26] shouldBe "textValue1" - - val table2Df = dataframes[1].cast() - - table2Df.rowsCount() shouldBe 3 - table2Df.filter { - "integerCol"()?.let { it > 400 } ?: false - }.rowsCount() shouldBe 1 - table2Df[0][11] shouldBe 20.0 - table2Df[0][26] shouldBe null - } - - @Test - fun `reading numeric types`() { - val df1 = DataFrame.readSqlTable(connection, "table1").cast() - - val result = df1.select("tinyintCol").add("tinyintCol2") { "tinyintCol"() } - - result[0][1] shouldBe 1.toByte() - - val result1 = df1.select("smallintCol") - .add("smallintCol2") { "smallintCol"() } - - result1[0][1] shouldBe 10.toShort() - - val result2 = df1.select("mediumintCol") - .add("mediumintCol2") { "mediumintCol"() } - - result2[0][1] shouldBe 100 - - val result3 = df1.select("mediumintUnsignedCol") - .add("mediumintUnsignedCol2") { "mediumintUnsignedCol"() } - - result3[0][1] shouldBe 100 - - val result4 = df1.select("integerUnsignedCol") - .add("integerUnsignedCol2") { "integerUnsignedCol"() } - - result4[0][1] shouldBe 100L - - val result5 = df1.select("bigintCol") - .add("bigintCol2") { "bigintCol"() } - - result5[0][1] shouldBe 100 - - val result5a = df1.select("bigintUnsignedCol") - .add("bigintUnsignedCol2") { "bigintUnsignedCol"() } - - result5a[0][1] shouldBe BigInteger.valueOf(1000) - - val result6 = df1.select("floatCol") - .add("floatCol2") { "floatCol"() } - - result6[0][1] shouldBe 10.0f - - val result7 = df1.select("doubleCol") - .add("doubleCol2") { "doubleCol"() } - - result7[0][1] shouldBe 10.0 - - val result8 = df1.select("decimalCol") - .add("decimalCol2") { "decimalCol"() } - - result8[0][1] shouldBe BigDecimal("10") - - val schema = DataFrameSchema.readSqlTable(connection, "table1") - - schema.columns["tinyintCol"]!!.type shouldBe typeOf() - schema.columns["smallintCol"]!!.type shouldBe typeOf() - schema.columns["mediumintCol"]!!.type shouldBe typeOf() - schema.columns["mediumintUnsignedCol"]!!.type shouldBe typeOf() - schema.columns["integerUnsignedCol"]!!.type shouldBe typeOf() - schema.columns["bigintCol"]!!.type shouldBe typeOf() - schema.columns["bigintUnsignedCol"]!!.type shouldBe typeOf() - schema.columns["floatCol"]!!.type shouldBe typeOf() - schema.columns["doubleCol"]!!.type shouldBe typeOf() - schema.columns["decimalCol"]!!.type shouldBe typeOf() - // TODO: all unsigned types - // TODO: new mapping system based on class names - // validation after mapping in getObject - // getObject(i+1, type) catch getObject catch getString - // add direct mapping to getString and other methods - } - - @Test - fun `infer nullability`() { - inferNullability(connection) - } - - // https://github.com/Kotlin/dataframe/issues/1746 - @Test - fun `readAllSqlTables without catalogue should only return tables from URL database`() { - val secondDb = "testKDFdatabase2" - val testRootConn = DriverManager.getConnection(URL, USER_NAME, PASSWORD) - try { - testRootConn.createStatement().use { stmt -> - stmt.executeUpdate("DROP DATABASE IF EXISTS $secondDb") - stmt.executeUpdate("CREATE DATABASE $secondDb") - } - DriverManager.getConnection("$URL/$secondDb", USER_NAME, PASSWORD).use { conn2 -> - conn2.createStatement().use { stmt -> - stmt.executeUpdate("CREATE TABLE onlyInDb2 (id INT PRIMARY KEY, val VARCHAR(50))") - } - } - - DriverManager.getConnection("$URL/$TEST_DATABASE_NAME", USER_NAME, PASSWORD).use { scopedConn -> - val tableNames = DataFrame.readAllSqlTables(scopedConn).keys - - tableNames.none { "onlyInDb2" in it } shouldBe true - tableNames.any { "table1" in it } shouldBe true - } - } finally { - testRootConn.use { conn -> - conn.createStatement().execute("DROP DATABASE IF EXISTS $secondDb") - } + tearDownMySqlTestData(connection) } } } diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresConnectionUrlTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresConnectionUrlTest.kt index 7f1fa4213b..869295544f 100644 --- a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresConnectionUrlTest.kt +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresConnectionUrlTest.kt @@ -1,108 +1,15 @@ package org.jetbrains.kotlinx.dataframe.io.local -import io.kotest.assertions.throwables.shouldThrow -import io.kotest.matchers.shouldBe -import org.jetbrains.kotlinx.dataframe.DataFrame -import org.jetbrains.kotlinx.dataframe.api.cast -import org.jetbrains.kotlinx.dataframe.api.filter -import org.jetbrains.kotlinx.dataframe.io.DbConnectionConfig -import org.jetbrains.kotlinx.dataframe.io.readDataFrame -import org.jetbrains.kotlinx.dataframe.io.readSqlTable -import org.junit.Ignore -import org.junit.Test -import java.sql.DriverManager +import org.jetbrains.kotlinx.dataframe.io.PostgresConnectionUrlTestBase -private const val URL_WITH_LOGIN_PASSWORD = "jdbc:postgresql://localhost:5432/test?" + - "user=postgres&password=pass&connectTimeout=10&tcpKeepAlive=true" +private const val URL = "jdbc:postgresql://localhost:5432/test" +private const val USER_NAME = "postgres" +private const val PASSWORD = "pass" -private const val URL_NO_LOGIN_PASSWORD = "jdbc:postgresql://localhost:5432/test?connectTimeout=10&tcpKeepAlive=true" +class PostgresConnectionUrlLocalTest : PostgresConnectionUrlTestBase() { + override val baseUrl: String get() = URL -private const val URL_WITH_PASSWORD = - "jdbc:postgresql://localhost:5432/test?password=pass&connectTimeout=10&tcpKeepAlive=true" + override val userName: String get() = USER_NAME -private const val URL_WITH_LOGIN = - "jdbc:postgresql://localhost:5432/test?user=postgres&connectTimeout=10&tcpKeepAlive=true" - -private const val TABLE_NAME = "table1" - -@Ignore -class PostgresConnectionUrlTest { - @Test - fun `read from table with login and password in connection URL`() { - DriverManager.getConnection(URL_WITH_LOGIN_PASSWORD).use { connection -> - createTestData(connection) - - val df1 = DataFrame.readSqlTable(connection, TABLE_NAME).cast() - val result1 = df1.filter { "id"() == 1 } - - result1[0][2] shouldBe 11 - - val df2 = connection.readDataFrame(TABLE_NAME).cast() - val result2 = df2.filter { "id"() == 1 } - - result2[0][2] shouldBe 11 - - clearTestData(connection) - } - } - - @Test - fun `read from table with login and password in connection URL for DBConfig`() { - DriverManager.getConnection(URL_WITH_LOGIN_PASSWORD).use { connection -> - createTestData(connection) - - val dbConfig = DbConnectionConfig(URL_WITH_LOGIN_PASSWORD) - val df1 = DataFrame.readSqlTable(dbConfig = dbConfig, TABLE_NAME).cast() - val result1 = df1.filter { "id"() == 1 } - - result1[0][2] shouldBe 11 - - val df2 = dbConfig.readDataFrame(TABLE_NAME).cast() - val result2 = df2.filter { "id"() == 1 } - - result2[0][2] shouldBe 11 - - clearTestData(connection) - } - } - - @Test - fun `read from table without login and password`() { - val dbConfig = DbConnectionConfig(URL_NO_LOGIN_PASSWORD) - - shouldThrow { - testReadFromTable(dbConfig) - } - } - - @Test - fun `read from table with password only`() { - val dbConfig = DbConnectionConfig(URL_WITH_PASSWORD) - - shouldThrow { - testReadFromTable(dbConfig) - } - } - - @Test - fun `read from table with login only`() { - val dbConfig = DbConnectionConfig(URL_WITH_LOGIN) - - shouldThrow { - testReadFromTable(dbConfig) - } - } - - private fun testReadFromTable(dbConfig: DbConnectionConfig) { - DriverManager.getConnection(URL_WITH_LOGIN_PASSWORD).use { connection -> - createTestData(connection) - - val df2 = dbConfig.readDataFrame(TABLE_NAME).cast() - val result2 = df2.filter { "id"() == 1 } - - result2[0][2] shouldBe 11 - - clearTestData(connection) - } - } + override val password: String get() = PASSWORD } diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresTest.kt index 0bcb107ef7..e110a0f049 100644 --- a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresTest.kt +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/local/postgresTest.kt @@ -1,422 +1,34 @@ package org.jetbrains.kotlinx.dataframe.io.local -import io.kotest.matchers.shouldBe -import org.intellij.lang.annotations.Language -import org.jetbrains.kotlinx.dataframe.DataFrame -import org.jetbrains.kotlinx.dataframe.annotations.DataSchema -import org.jetbrains.kotlinx.dataframe.api.add -import org.jetbrains.kotlinx.dataframe.api.cast -import org.jetbrains.kotlinx.dataframe.api.filter -import org.jetbrains.kotlinx.dataframe.api.select -import org.jetbrains.kotlinx.dataframe.io.inferNullability -import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables -import org.jetbrains.kotlinx.dataframe.io.readSqlQuery -import org.jetbrains.kotlinx.dataframe.io.readSqlTable -import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.jetbrains.kotlinx.dataframe.io.PostgresTestBase +import org.jetbrains.kotlinx.dataframe.io.createPostgresTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownPostgresTestData import org.junit.AfterClass import org.junit.BeforeClass -import org.junit.Ignore -import org.junit.Test -import org.postgresql.geometric.PGbox -import org.postgresql.geometric.PGcircle -import org.postgresql.geometric.PGline -import org.postgresql.geometric.PGlseg -import org.postgresql.geometric.PGpath -import org.postgresql.geometric.PGpoint -import org.postgresql.geometric.PGpolygon -import org.postgresql.util.PGInterval -import org.postgresql.util.PGmoney -import org.postgresql.util.PGobject -import java.math.BigDecimal import java.sql.Connection -import java.sql.Date import java.sql.DriverManager -import java.sql.SQLException -import java.sql.Time -import java.sql.Timestamp -import java.sql.Types -import java.util.UUID -import kotlin.reflect.typeOf -private const val BASIC_URL = "jdbc:postgresql://localhost:5432/test" +private const val URL = "jdbc:postgresql://localhost:5432/test" private const val USER_NAME = "postgres" private const val PASSWORD = "pass" -@DataSchema -interface Table1 { - val id: Int - val bigintcol: Long - val smallintcol: Int - val bigserialcol: Long - val booleancol: Boolean - val boxcol: PGbox - val byteacol: ByteArray - val charactercol: String - val characterncol: String - val charcol: String - val circlecol: PGcircle - val datecol: java.sql.Date - val doublecol: Double - val integercol: Int? - val intervalcol: String - val jsoncol: String - val jsonbcol: String -} - -@DataSchema -interface Table2 { - val id: Int - val linecol: PGline - val lsegcol: PGlseg - val macaddrcol: String - val moneycol: PGmoney - val numericcol: BigDecimal - val pathcol: PGpath - val pointcol: PGpoint - val polygoncol: PGpolygon - val realcol: Float - val smallintcol: Int - val smallserialcol: Int - val serialcol: Int - val textcol: String? - val timecol: String - val timewithzonecol: String - val timestampcol: String - val timestampwithzonecol: String - val uuidcol: String - val xmlcol: String -} - -@DataSchema -interface ViewTable { - val id: Int - val bigintcol: Long - val linecol: String - val textCol: String? -} - -internal fun createTestData(connection: Connection) { - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } - - val createTableStatement = """ - CREATE TABLE IF NOT EXISTS table1 ( - id serial PRIMARY KEY, - bigintCol bigint not null, - smallintCol smallint not null, - bigserialCol bigserial not null, - booleanCol boolean not null, - boxCol box not null, - byteaCol bytea not null, - characterCol character not null, - characterNCol character(10) not null, - charCol char not null, - circleCol circle not null, - dateCol date not null, - doubleCol double precision not null, - integerCol integer, - intervalCol interval not null, - jsonCol json not null, - jsonbCol jsonb not null, - intArrayCol integer[], - doubleArrayCol double precision array, - dateArrayCol date array, - textArrayCol text array, - booleanArrayCol boolean array - ) - """ - connection.createStatement().execute(createTableStatement.trimIndent()) - - val createTableQuery = """ - CREATE TABLE IF NOT EXISTS table2 ( - id serial PRIMARY KEY, - lineCol line not null, - lsegCol lseg not null, - macaddrCol macaddr not null, - moneyCol money not null, - numericCol numeric not null, - pathCol path not null, - pointCol point not null, - polygonCol polygon not null, - realCol real not null, - smallintCol smallint not null, - smallserialCol smallserial not null, - serialCol serial not null, - textCol text, - timeCol time not null, - timeWithZoneCol time with time zone not null, - timestampCol timestamp not null, - timestampWithZoneCol timestamp with time zone not null, - uuidCol uuid not null, - xmlCol xml not null - ) - """ - connection.createStatement().execute(createTableQuery.trimIndent()) +class PostgresLocalTest : PostgresTestBase() { + override val connection: Connection get() = Companion.connection - @Language("SQL") - val insertData1 = """ - INSERT INTO table1 ( - bigintCol, smallintCol, bigserialCol, booleanCol, - boxCol, byteaCol, characterCol, characterNCol, charCol, - circleCol, dateCol, doubleCol, - integerCol, intervalCol, jsonCol, jsonbCol, intArrayCol, - doubleArrayCol, dateArrayCol, textArrayCol, booleanArrayCol - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """ - - @Language("SQL") - val insertData2 = """ - INSERT INTO table2 ( - lineCol, lsegCol, macaddrCol, moneyCol, numericCol, - pathCol, pointCol, polygonCol, realCol, smallintCol, - smallserialCol, serialCol, textCol, timeCol, - timeWithZoneCol, timestampCol, timestampWithZoneCol, - uuidCol, xmlCol - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """ - - val intArray = connection.createArrayOf("INTEGER", arrayOf(1, 2, 3)) - val doubleArray = connection.createArrayOf("DOUBLE", arrayOf(1.1, 2.2, 3.3)) - val dateArray = connection.createArrayOf("DATE", arrayOf(Date.valueOf("2023-08-01"), Date.valueOf("2023-08-02"))) - val textArray = connection.createArrayOf("TEXT", arrayOf("Hello", "World")) - val booleanArray = connection.createArrayOf("BOOLEAN", arrayOf(true, false, true)) - - connection.prepareStatement(insertData1).use { st -> - // Insert data into table1 - for (i in 1..3) { - st.setLong(1, i * 1000L) - st.setShort(2, 11.toShort()) - st.setLong(3, 1000000000L + i) - st.setBoolean(4, i % 2 == 1) - st.setObject(5, PGbox("(1,1),(2,2)")) - st.setBytes(6, byteArrayOf(1, 2, 3)) - st.setString(7, "A") - st.setString(8, "Hello") - st.setString(9, "A") - st.setObject(10, PGcircle("<(1,2),3>")) - st.setDate(11, Date.valueOf("2023-08-01")) - st.setDouble(12, 12.34) - st.setInt(13, 12345 * i) - st.setObject(14, PGInterval("1 year")) - - val jsonbObject = PGobject() - jsonbObject.type = "jsonb" - jsonbObject.value = "{\"key\": \"value\"}" - - st.setObject(15, jsonbObject) - st.setObject(16, jsonbObject) - st.setArray(17, intArray) - st.setArray(18, doubleArray) - st.setArray(19, dateArray) - st.setArray(20, textArray) - st.setArray(21, booleanArray) - st.executeUpdate() - } - } - - connection.prepareStatement(insertData2).use { st -> - // Insert data into table2 - for (i in 1..3) { - st.setObject(1, PGline("{1,2,3}")) - st.setObject(2, PGlseg("[(-1,0),(1,0)]")) - - val macaddrObject = PGobject() - macaddrObject.type = "macaddr" - macaddrObject.value = "00:00:00:00:00:0$i" - - st.setObject(3, macaddrObject) - st.setBigDecimal(4, BigDecimal("123.45")) - st.setBigDecimal(5, BigDecimal("12.34")) - st.setObject(6, PGpath("((1,2),(3,$i))")) - st.setObject(7, PGpoint("(1,2)")) - st.setObject(8, PGpolygon("((1,1),(2,2),(3,3))")) - st.setFloat(9, 12.34f) - st.setShort(10, (i * 100).toShort()) - st.setInt(11, 1000 + i) - st.setInt(12, 1000000 + i) - st.setString(13, null) - st.setTime(14, Time.valueOf("12:34:56")) - - st.setTimestamp(15, Timestamp(System.currentTimeMillis())) - st.setTimestamp(16, Timestamp(System.currentTimeMillis())) - st.setTimestamp(17, Timestamp(System.currentTimeMillis())) - - st.setObject(18, UUID.randomUUID(), Types.OTHER) - val xmlObject = PGobject() - xmlObject.type = "xml" - xmlObject.value = "data" - - st.setObject(19, xmlObject) - st.executeUpdate() - } - } -} - -internal fun clearTestData(connection: Connection) { - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } - connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } -} - -@Ignore -class PostgresTest { companion object { private lateinit var connection: Connection @BeforeClass @JvmStatic fun setUpClass() { - connection = DriverManager.getConnection(BASIC_URL, USER_NAME, PASSWORD) - createTestData(connection) + connection = DriverManager.getConnection(URL, USER_NAME, PASSWORD) + createPostgresTestData(connection) } @AfterClass @JvmStatic fun tearDownClass() { - try { - clearTestData(connection) - connection.close() - } catch (e: SQLException) { - e.printStackTrace() - } + tearDownPostgresTestData(connection) } } - - @Test - fun `read from tables`() { - val tableName1 = "table1" - val df1 = DataFrame.readSqlTable(connection, tableName1).cast() - val result = df1.filter { "id"() == 1 } - - result[0][2] shouldBe 11 - result[0][13] shouldBe 12345 - result[0][17] shouldBe arrayOf(1, 2, 3) - result[0][18] shouldBe arrayOf(1.1, 2.2, 3.3) - result[0][19] shouldBe arrayOf(Date.valueOf("2023-08-01"), Date.valueOf("2023-08-02")) - result[0][20] shouldBe arrayOf("Hello", "World") - result[0][21] shouldBe arrayOf(true, false, true) - - val schema = DataFrameSchema.readSqlTable(connection, tableName1) - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["integercol"]!!.type shouldBe typeOf() - schema.columns["smallintcol"]!!.type shouldBe typeOf() - schema.columns["circlecol"]!!.type shouldBe typeOf() - schema.columns["intarraycol"]!!.type.classifier shouldBe kotlin.Array::class - schema.columns["doublearraycol"]!!.type.classifier shouldBe kotlin.Array::class - schema.columns["datearraycol"]!!.type.classifier shouldBe kotlin.Array::class - schema.columns["textarraycol"]!!.type.classifier shouldBe kotlin.Array::class - schema.columns["booleanarraycol"]!!.type.classifier shouldBe kotlin.Array::class - - val tableName2 = "table2" - val df2 = DataFrame.readSqlTable(connection, tableName2).cast() - val result2 = df2.filter { "id"() == 1 } - result2[0][11] shouldBe 1001 - result2[0][13] shouldBe null - - val schema2 = DataFrameSchema.readSqlTable(connection, tableName2) - schema2.columns["id"]!!.type shouldBe typeOf() - schema2.columns["pathcol"]!!.type shouldBe typeOf() - schema2.columns["textcol"]!!.type shouldBe typeOf() - schema2.columns["linecol"]!!.type shouldBe typeOf() - } - - @Test - fun `read from sql query`() { - @Language("SQL") - val sqlQuery = - """ - SELECT - t1.id, - t1.bigintCol, - t2.lineCol, - t2.textCol - FROM table1 t1 - JOIN table2 t2 ON t1.id = t2.id - """.trimIndent() - - val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery).cast() - val result = df.filter { "id"() == 1 } - result[0][3] shouldBe null - - val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) - schema.columns["id"]!!.type shouldBe typeOf() - schema.columns["bigintcol"]!!.type shouldBe typeOf() - schema.columns["textcol"]!!.type shouldBe typeOf() - } - - @Test - fun `read from all tables`() { - val dataframes = DataFrame.readAllSqlTables(connection).values.toList() - - val table1Df = dataframes[0].cast() - - table1Df.rowsCount() shouldBe 3 - table1Df.filter { "integercol"()?.let { it > 12345 } ?: false }.rowsCount() shouldBe 2 - table1Df[0][1] shouldBe 1000L - table1Df[0][2] shouldBe 11 - - val table2Df = dataframes[1].cast() - - table2Df.rowsCount() shouldBe 3 - table2Df.filter { - "pathcol"() == PGpath("((1,2),(3,1))") - }.rowsCount() shouldBe 1 - table2Df[0][11] shouldBe 1001 - } - - @Test - fun `read columns of different types to check type mapping`() { - val tableName1 = "table1" - val df1 = DataFrame.readSqlTable(connection, tableName1).cast() - val result = df1.select("smallintcol") - .add("smallintcol2") { "smallintcol"() } - result[0][1] shouldBe 11 - - val result1 = df1.select("bigserialcol") - .add("bigserialcol2") { "bigserialcol"() } - result1[0][1] shouldBe 1000000001L - - val result2 = df1.select("doublecol") - .add("doublecol2") { "doublecol"() } - result2[0][1] shouldBe 12.34 - - val tableName2 = "table2" - val df2 = DataFrame.readSqlTable(connection, tableName2).cast() - - val result3 = df2.select("moneycol") - .add("moneycol2") { "moneycol"() } - (result3[0][1] as PGmoney).`val` shouldBe 123.45 - - val result4 = df2.select("numericcol") - .add("numericcol2") { "numericcol"() } - result4[0][1] shouldBe BigDecimal("12.34") - - val result5 = df2.select("realcol") - .add("realcol2") { "realcol"() } - result5[0][1] shouldBe 12.34f - - val result7 = df2.select("smallserialcol") - .add("smallserialcol2") { "smallserialcol"() } - result7[0][1] shouldBe 1001 - - val result8 = df2.select("serialcol") - .add("serialcol2") { "serialcol"() } - result8[0][1] shouldBe 1000001 - - val schema = DataFrameSchema.readSqlTable(connection, tableName1) - schema.columns["smallintcol"]!!.type shouldBe typeOf() - schema.columns["bigserialcol"]!!.type shouldBe typeOf() - schema.columns["doublecol"]!!.type shouldBe typeOf() - - val schema1 = DataFrameSchema.readSqlTable(connection, tableName2) - schema1.columns["moneycol"]!!.type shouldBe typeOf() - schema1.columns["numericcol"]!!.type shouldBe typeOf() - schema1.columns["realcol"]!!.type shouldBe typeOf() - schema1.columns["smallserialcol"]!!.type shouldBe typeOf() - schema1.columns["serialcol"]!!.type shouldBe typeOf() - } - - @Test - fun `infer nullability`() { - inferNullability(connection) - } } diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mariadbTestBase.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mariadbTestBase.kt new file mode 100644 index 0000000000..9c58f10371 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mariadbTestBase.kt @@ -0,0 +1,506 @@ +package org.jetbrains.kotlinx.dataframe.io + +import io.kotest.matchers.shouldBe +import org.intellij.lang.annotations.Language +import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema +import org.jetbrains.kotlinx.dataframe.api.add +import org.jetbrains.kotlinx.dataframe.api.cast +import org.jetbrains.kotlinx.dataframe.api.filter +import org.jetbrains.kotlinx.dataframe.api.select +import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.junit.Test +import java.math.BigDecimal +import java.math.BigInteger +import java.sql.Blob +import java.sql.Connection +import java.sql.SQLException +import java.util.Date +import kotlin.reflect.typeOf +import kotlin.time.Instant +import java.sql.Time as SqlTime +import java.sql.Timestamp as SqlTimestamp + +internal const val MARIADB_TEST_DATABASE_NAME = "testKDFdatabase" + +@DataSchema +interface Table1MariaDb { + val id: Int + val bitCol: Boolean + val tinyintCol: Int + val smallintCol: Short? + val mediumintCol: Int + val mediumintUnsignedCol: Int + val integerCol: Int + val intCol: Int + val integerUnsignedCol: Long + val bigintCol: Long + val floatCol: Float + val doubleCol: Double + val decimalCol: BigDecimal + val dateCol: String + val datetimeCol: String + val timestampCol: String + val timeCol: String + val yearCol: String + val varcharCol: String + val charCol: String + val binaryCol: ByteArray + val varbinaryCol: ByteArray + val tinyblobCol: ByteArray + val blobCol: ByteArray + val mediumblobCol: ByteArray + val longblobCol: ByteArray + val textCol: String + val mediumtextCol: String + val longtextCol: String + val enumCol: String + val setCol: Char + val bigintUnsignedCol: BigInteger + val jsonCol: String +} + +@DataSchema +interface Table2MariaDb { + val id: Int + val bitCol: Boolean? + val tinyintCol: Int? + val smallintCol: Int? + val mediumintCol: Int? + val mediumintUnsignedCol: Int? + val integerCol: Int? + val intCol: Int? + val integerUnsignedCol: Long? + val bigintCol: Long? + val floatCol: Float? + val doubleCol: Double? + val decimalCol: Double? + val dateCol: String? + val datetimeCol: String? + val timestampCol: String? + val timeCol: String? + val yearCol: String? + val varcharCol: String? + val charCol: String? + val binaryCol: ByteArray? + val varbinaryCol: ByteArray? + val tinyblobCol: ByteArray? + val blobCol: ByteArray? + val mediumblobCol: ByteArray? + val longblobCol: ByteArray? + val textCol: String? + val mediumtextCol: String? + val longtextCol: String? + val enumCol: String? + val setCol: Char? + val bigintUnsignedCol: BigInteger? + val jsonCol: String? +} + +@DataSchema +interface Table3MariaDb { + val id: Int + val enumCol: String + val setCol: Char? +} + +private const val JSON_STRING = + "{\"details\": {\"foodType\": \"Pizza\", \"menu\": \"https://www.loumalnatis.com/our-menu\"}, \n" + + " \t\"favorites\": [{\"description\": \"Pepperoni deep dish\", \"price\": 18.75}, \n" + + "{\"description\": \"The Lou\", \"price\": 24.75}]}" + +internal fun setUpMariadbTestData(connection: Connection) { + connection.createStatement().use { st -> + // Drop the test database if it exists + val dropDatabaseQuery = "DROP DATABASE IF EXISTS $MARIADB_TEST_DATABASE_NAME" + st.executeUpdate(dropDatabaseQuery) + + // Create the test database + val createDatabaseQuery = "CREATE DATABASE $MARIADB_TEST_DATABASE_NAME" + st.executeUpdate(createDatabaseQuery) + + // Use the newly created database + val useDatabaseQuery = "USE $MARIADB_TEST_DATABASE_NAME" + st.executeUpdate(useDatabaseQuery) + } + + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } + + @Language("SQL") + val createTableQuery = """ + CREATE TABLE IF NOT EXISTS table1 ( + id INT AUTO_INCREMENT PRIMARY KEY, + bitCol BIT NOT NULL, + tinyintCol TINYINT NOT NULL, + smallintCol SMALLINT, + mediumintCol MEDIUMINT NOT NULL, + mediumintUnsignedCol MEDIUMINT UNSIGNED NOT NULL, + integerCol INTEGER NOT NULL, + intCol INT NOT NULL, + integerUnsignedCol INTEGER UNSIGNED NOT NULL, + bigintCol BIGINT NOT NULL, + floatCol FLOAT NOT NULL, + doubleCol DOUBLE NOT NULL, + decimalCol DECIMAL NOT NULL, + dateCol DATE NOT NULL, + datetimeCol DATETIME NOT NULL, + timestampCol TIMESTAMP NOT NULL, + timeCol TIME NOT NULL, + yearCol YEAR NOT NULL, + varcharCol VARCHAR(255) NOT NULL, + charCol CHAR(10) NOT NULL, + binaryCol BINARY(64) NOT NULL, + varbinaryCol VARBINARY(128) NOT NULL, + tinyblobCol TINYBLOB NOT NULL, + blobCol BLOB NOT NULL, + mediumblobCol MEDIUMBLOB NOT NULL , + longblobCol LONGBLOB NOT NULL, + textCol TEXT NOT NULL, + mediumtextCol MEDIUMTEXT NOT NULL, + longtextCol LONGTEXT NOT NULL, + enumCol ENUM('Value1', 'Value2', 'Value3') NOT NULL, + setCol SET('Option1', 'Option2', 'Option3') NOT NULL, + bigintUnsignedCol BIGINT UNSIGNED NOT NULL, + jsonCol JSON NOT NULL + CHECK (JSON_VALID(jsonCol)) + ) + """ + connection.createStatement().execute(createTableQuery.trimIndent()) + + @Language("SQL") + val createTableQuery2 = """ + CREATE TABLE IF NOT EXISTS table2 ( + id INT AUTO_INCREMENT PRIMARY KEY, + bitCol BIT, + tinyintCol TINYINT, + smallintCol SMALLINT, + mediumintCol MEDIUMINT, + mediumintUnsignedCol MEDIUMINT UNSIGNED, + integerCol INTEGER, + intCol INT, + integerUnsignedCol INTEGER UNSIGNED, + bigintCol BIGINT, + floatCol FLOAT, + doubleCol DOUBLE, + decimalCol DECIMAL, + dateCol DATE, + datetimeCol DATETIME, + timestampCol TIMESTAMP, + timeCol TIME, + yearCol YEAR, + varcharCol VARCHAR(255), + charCol CHAR(10), + binaryCol BINARY(64), + varbinaryCol VARBINARY(128), + tinyblobCol TINYBLOB, + blobCol BLOB, + mediumblobCol MEDIUMBLOB, + longblobCol LONGBLOB, + textCol TEXT, + mediumtextCol MEDIUMTEXT, + longtextCol LONGTEXT, + enumCol ENUM('Value1', 'Value2', 'Value3'), + setCol SET('Option1', 'Option2', 'Option3'), + bigintUnsignedCol BIGINT UNSIGNED + ) + """ + connection.createStatement().execute(createTableQuery2.trimIndent()) + + @Language("SQL") + val insertData1 = + """ + INSERT INTO table1 ( + bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, + integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, + timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, + mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol, jsonCol + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """.trimIndent() + + @Language("SQL") + val insertData2 = + """ + INSERT INTO table2 ( + bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, + integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, + timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, + mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """.trimIndent() + + connection.prepareStatement(insertData1).use { st -> + // Insert data into table1 + for (i in 1..3) { + st.setBoolean(1, true) + st.setByte(2, i.toByte()) + st.setShort(3, (i * 10).toShort()) + st.setInt(4, i * 100) + st.setInt(5, i * 100) + st.setInt(6, i * 100) + st.setInt(7, i * 100) + st.setInt(8, i * 100) + st.setInt(9, i * 100) + st.setFloat(10, i * 10.0f) + st.setDouble(11, i * 10.0) + st.setBigDecimal(12, BigDecimal(i * 10)) + st.setDate(13, java.sql.Date(System.currentTimeMillis())) + st.setTimestamp(14, SqlTimestamp(System.currentTimeMillis())) + st.setTimestamp(15, SqlTimestamp(System.currentTimeMillis())) + st.setTime(16, SqlTime(System.currentTimeMillis())) + st.setInt(17, 2023) + st.setString(18, "varcharValue$i") + st.setString(19, "charValue$i") + st.setBytes(20, "binaryValue".toByteArray()) + st.setBytes(21, "varbinaryValue".toByteArray()) + st.setBytes(22, "tinyblobValue".toByteArray()) + st.setBytes(23, "blobValue".toByteArray()) + st.setBytes(24, "mediumblobValue".toByteArray()) + st.setBytes(25, "longblobValue".toByteArray()) + st.setString(26, "textValue$i") + st.setString(27, "mediumtextValue$i") + st.setString(28, "longtextValue$i") + st.setString(29, "Value$i") + st.setString(30, "Option$i") + st.setObject(31, BigInteger.valueOf((i * 1000).toLong())) + st.setString(32, JSON_STRING) + + st.executeUpdate() + } + } + + connection.prepareStatement(insertData2).use { st -> + // Insert data into table2 + for (i in 1..3) { + st.setBoolean(1, false) + st.setByte(2, (i * 2).toByte()) + st.setShort(3, (i * 20).toShort()) + st.setInt(4, i * 200) + st.setInt(5, i * 200) + st.setInt(6, i * 200) + st.setInt(7, i * 200) + st.setInt(8, i * 200) + st.setInt(9, i * 200) + st.setFloat(10, i * 20.0f) + st.setDouble(11, i * 20.0) + st.setBigDecimal(12, BigDecimal(i * 20)) + st.setDate(13, java.sql.Date(System.currentTimeMillis())) + st.setTimestamp(14, SqlTimestamp(System.currentTimeMillis())) + st.setTimestamp(15, SqlTimestamp(System.currentTimeMillis())) + st.setTime(16, SqlTime(System.currentTimeMillis())) + st.setInt(17, 2023) + st.setString(18, "varcharValue$i") + st.setString(19, "charValue$i") + st.setBytes(20, "binaryValue".toByteArray()) + st.setBytes(21, "varbinaryValue".toByteArray()) + st.setBytes(22, "tinyblobValue".toByteArray()) + st.setBytes(23, "blobValue".toByteArray()) + st.setBytes(24, "mediumblobValue".toByteArray()) + st.setBytes(25, "longblobValue".toByteArray()) + st.setString(26, null) + st.setString(27, null) + st.setString(28, "longtextValue$i") + st.setString(29, "Value$i") + st.setString(30, "Option$i") + st.setObject(31, BigInteger.valueOf((i * 2000).toLong())) + st.executeUpdate() + } + } +} + +internal fun tearDownMariadbTestData(connection: Connection) { + try { + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } + connection.createStatement().use { st -> + st.execute("DROP DATABASE IF EXISTS $MARIADB_TEST_DATABASE_NAME") + } + connection.close() + } catch (e: SQLException) { + e.printStackTrace() + } +} + +abstract class MariadbTestBase { + protected abstract val connection: Connection + + protected abstract fun connect(database: String? = null): Connection + + @Test + fun `basic test for reading sql tables`() { + val df1 = DataFrame.readSqlTable(connection, "table1").cast() + val result = df1.filter { "id"() == 1 } + result[0][26] shouldBe "textValue1" + val byteArray = "tinyblobValue".toByteArray() + result[0][22] shouldBe byteArray + + val schema = DataFrameSchema.readSqlTable(connection, "table1") + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["textCol"]!!.type shouldBe typeOf() + schema.columns["varbinaryCol"]!!.type shouldBe typeOf() + schema.columns["binaryCol"]!!.type shouldBe typeOf() + schema.columns["longblobCol"]!!.type shouldBe typeOf() + schema.columns["tinyblobCol"]!!.type shouldBe typeOf() + schema.columns["dateCol"]!!.type shouldBe typeOf() + schema.columns["datetimeCol"]!!.type shouldBe typeOf() + schema.columns["timestampCol"]!!.type shouldBe typeOf() + schema.columns["timeCol"]!!.type shouldBe typeOf() + schema.columns["yearCol"]!!.type shouldBe typeOf() + + val df2 = DataFrame.readSqlTable(connection, "table2").cast() + val result2 = df2.filter { "id"() == 1 } + result2[0][26] shouldBe null + + val schema2 = DataFrameSchema.readSqlTable(connection, "table2") + schema2.columns["id"]!!.type shouldBe typeOf() + schema2.columns["textCol"]!!.type shouldBe typeOf() + } + + @Test + fun `read from sql query`() { + @Language("SQL") + val sqlQuery = + """ + SELECT + t1.id, + t1.enumCol, + t2.setCol + FROM table1 t1 + JOIN table2 t2 ON t1.id = t2.id + """.trimIndent() + + val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery).cast() + val result = df.filter { "id"() == 1 } + result[0][2] shouldBe "Option1" + + val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["enumCol"]!!.type shouldBe typeOf() + schema.columns["setCol"]!!.type shouldBe typeOf() + } + + @Test + fun `read from all tables`() { + val dataframes = DataFrame.readAllSqlTables(connection, MARIADB_TEST_DATABASE_NAME, 1000).values.toList() + + val table1Df = dataframes[0].cast() + + table1Df.rowsCount() shouldBe 3 + table1Df.filter { "integerCol"() > 100 }.rowsCount() shouldBe 2 + table1Df[0][11] shouldBe 10.0 + table1Df[0][26] shouldBe "textValue1" + table1Df[0][31] shouldBe BigInteger.valueOf(1000L) + table1Df[0][32] shouldBe JSON_STRING // TODO: https://github.com/Kotlin/dataframe/issues/462 + + val table2Df = dataframes[1].cast() + + table2Df.rowsCount() shouldBe 3 + table2Df.filter { + "integerCol"()?.let { it > 400 } ?: false + }.rowsCount() shouldBe 1 + table2Df[0][11] shouldBe 20.0 + table2Df[0][26] shouldBe null + } + + @Test + fun `reading numeric types`() { + val df1 = DataFrame.readSqlTable(connection, "table1").cast() + + val result = df1.select("tinyintCol") + .add("tinyintCol2") { "tinyintCol"() } + + result[0][1] shouldBe 1 + + val result1 = df1.select("smallintCol") + .add("smallintCol2") { "smallintCol"() } + + result1[0][1] shouldBe 10 + + val result2 = df1.select("mediumintCol") + .add("mediumintCol2") { "mediumintCol"() } + + result2[0][1] shouldBe 100 + + val result3 = df1.select("mediumintUnsignedCol") + .add("mediumintUnsignedCol2") { "mediumintUnsignedCol"() } + + result3[0][1] shouldBe 100 + + val result4 = df1.select("integerUnsignedCol") + .add("integerUnsignedCol2") { "integerUnsignedCol"() } + + result4[0][1] shouldBe 100L + + val result5 = df1.select("bigintCol") + .add("bigintCol2") { "bigintCol"() } + + result5[0][1] shouldBe 100 + + val result5a = df1.select("bigintUnsignedCol") + .add("bigintUnsignedCol2") { "bigintUnsignedCol"() } + + result5a[0][1] shouldBe BigInteger.valueOf(1000) + + val result6 = df1.select("floatCol") + .add("floatCol2") { "floatCol"() } + + result6[0][1] shouldBe 10.0f + + val result7 = df1.select("doubleCol") + .add("doubleCol2") { "doubleCol"() } + + result7[0][1] shouldBe 10.0 + + val result8 = df1.select("decimalCol") + .add("decimalCol2") { "decimalCol"() } + + result8[0][1] shouldBe BigDecimal("10") + + val schema = DataFrameSchema.readSqlTable(connection, "table1") + + schema.columns["tinyintCol"]!!.type shouldBe typeOf() + schema.columns["smallintCol"]!!.type shouldBe typeOf() + schema.columns["mediumintCol"]!!.type shouldBe typeOf() + schema.columns["mediumintUnsignedCol"]!!.type shouldBe typeOf() + schema.columns["integerUnsignedCol"]!!.type shouldBe typeOf() + schema.columns["bigintCol"]!!.type shouldBe typeOf() + schema.columns["bigintUnsignedCol"]!!.type shouldBe typeOf() + schema.columns["floatCol"]!!.type shouldBe typeOf() + schema.columns["doubleCol"]!!.type shouldBe typeOf() + schema.columns["decimalCol"]!!.type shouldBe typeOf() + } + + @Test + fun `infer nullability`() { + inferNullability(connection) + } + + // https://github.com/Kotlin/dataframe/issues/1746 + @Test + fun `readAllSqlTables without catalogue should only return tables from URL database`() { + val secondDb = "testKDFdatabase2" + val testRootConn = connect() + try { + testRootConn.createStatement().use { stmt -> + stmt.executeUpdate("DROP DATABASE IF EXISTS $secondDb") + stmt.executeUpdate("CREATE DATABASE $secondDb") + } + connect(secondDb).use { conn2 -> + conn2.createStatement().use { stmt -> + stmt.executeUpdate("CREATE TABLE onlyInDb2 (id INT PRIMARY KEY, val VARCHAR(50))") + } + } + + connect(MARIADB_TEST_DATABASE_NAME).use { scopedConn -> + val tableNames = DataFrame.readAllSqlTables(scopedConn).keys + + tableNames.none { "onlyInDb2" in it } shouldBe true + tableNames.any { "table1" in it } shouldBe true + } + } finally { + testRootConn.use { conn -> + conn.createStatement().execute("DROP DATABASE IF EXISTS $secondDb") + } + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mssqlTestBase.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mssqlTestBase.kt new file mode 100644 index 0000000000..ace3cfbf6e --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mssqlTestBase.kt @@ -0,0 +1,316 @@ +package org.jetbrains.kotlinx.dataframe.io + +import io.kotest.matchers.shouldBe +import org.intellij.lang.annotations.Language +import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema +import org.jetbrains.kotlinx.dataframe.api.cast +import org.jetbrains.kotlinx.dataframe.api.filter +import org.jetbrains.kotlinx.dataframe.io.inferNullability +import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables +import org.jetbrains.kotlinx.dataframe.io.readSqlQuery +import org.jetbrains.kotlinx.dataframe.io.readSqlTable +import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.junit.Test +import java.math.BigDecimal +import java.sql.Connection +import java.util.Date +import java.util.UUID +import kotlin.reflect.typeOf +import kotlin.time.Instant + +internal const val MSSQL_TEST_DATABASE_NAME = "testKDFdatabase" + +@DataSchema +interface Table1MSSSQL { + val id: Int + val bigintColumn: Long + val binaryColumn: ByteArray + val bitColumn: Boolean + val charColumn: Char + val dateColumn: Date + val datetime3Column: Instant + val datetime2Column: Instant + val datetimeoffset2Column: String + val decimalColumn: BigDecimal + val floatColumn: Double + val imageColumn: ByteArray? + val intColumn: Int + val moneyColumn: BigDecimal + val ncharColumn: Char + val ntextColumn: String + val numericColumn: BigDecimal + val nvarcharColumn: String + val nvarcharMaxColumn: String + val realColumn: Float + val smalldatetimeColumn: Instant + val smallintColumn: Int + val smallmoneyColumn: BigDecimal + val timeColumn: java.sql.Time + val timestampColumn: Instant + val tinyintColumn: Int + val uniqueidentifierColumn: Char + val varbinaryColumn: ByteArray + val varbinaryMaxColumn: ByteArray + val varcharColumn: String + val varcharMaxColumn: String + val xmlColumn: String + val sqlvariantColumn: String + val geometryColumn: ByteArray + val geographyColumn: ByteArray +} + +internal fun setUpMsSqlTestData(connection: Connection) { + connection.createStatement().use { st -> + // Drop the test database if it exists + val dropDatabaseQuery = "IF DB_ID('$MSSQL_TEST_DATABASE_NAME') IS NOT NULL\n" + + "DROP DATABASE $MSSQL_TEST_DATABASE_NAME" + st.executeUpdate(dropDatabaseQuery) + + // Create the test database + val createDatabaseQuery = "CREATE DATABASE $MSSQL_TEST_DATABASE_NAME" + st.executeUpdate(createDatabaseQuery) + + // Use the newly created database + val useDatabaseQuery = "USE $MSSQL_TEST_DATABASE_NAME" + st.executeUpdate(useDatabaseQuery) + } + + @Language("SQL") + val createTableQuery = """ + CREATE TABLE Table1 ( + id INT NOT NULL IDENTITY PRIMARY KEY, + bigintColumn BIGINT, + binaryColumn BINARY(50), + bitColumn BIT, + charColumn CHAR(10), + dateColumn DATE, + datetime3Column DATETIME2(3), + datetime2Column DATETIME2, + datetimeoffset2Column DATETIMEOFFSET(2), + decimalColumn DECIMAL(10,2), + floatColumn FLOAT, + imageColumn IMAGE, + intColumn INT, + moneyColumn MONEY, + ncharColumn NCHAR(10), + ntextColumn NTEXT, + numericColumn NUMERIC(10,2), + nvarcharColumn NVARCHAR(50), + nvarcharMaxColumn NVARCHAR(MAX), + realColumn REAL, + smalldatetimeColumn SMALLDATETIME, + smallintColumn SMALLINT, + smallmoneyColumn SMALLMONEY, + textColumn TEXT, + timeColumn TIME, + timestampColumn DATETIME2, + tinyintColumn TINYINT, + uniqueidentifierColumn UNIQUEIDENTIFIER, + varbinaryColumn VARBINARY(50), + varbinaryMaxColumn VARBINARY(MAX), + varcharColumn VARCHAR(50), + varcharMaxColumn VARCHAR(MAX), + xmlColumn XML, + sqlvariantColumn SQL_VARIANT, + geometryColumn GEOMETRY, + geographyColumn GEOGRAPHY + ); + """ + + connection.createStatement().execute(createTableQuery.trimIndent()) + + @Language("SQL") + val insertData1 = + """ + INSERT INTO Table1 ( + bigintColumn, binaryColumn, bitColumn, charColumn, dateColumn, datetime3Column, datetime2Column, + datetimeoffset2Column, decimalColumn, floatColumn, imageColumn, intColumn, moneyColumn, ncharColumn, + ntextColumn, numericColumn, nvarcharColumn, nvarcharMaxColumn, realColumn, smalldatetimeColumn, + smallintColumn, smallmoneyColumn, textColumn, timeColumn, timestampColumn, tinyintColumn, + uniqueidentifierColumn, varbinaryColumn, varbinaryMaxColumn, varcharColumn, varcharMaxColumn, + xmlColumn, sqlvariantColumn, geometryColumn, geographyColumn + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """.trimIndent() + + connection.prepareStatement(insertData1).use { st -> + for (i in 1..5) { + st.setLong(1, 123456789012345L) // bigintColumn + st.setBytes(2, byteArrayOf(0x01, 0x23, 0x45, 0x67, 0x67, 0x67, 0x67, 0x67)) // binaryColumn + st.setBoolean(3, true) // bitColumn + st.setString(4, "Sample") // charColumn + st.setDate(5, java.sql.Date(System.currentTimeMillis())) // dateColumn + st.setTimestamp(6, java.sql.Timestamp(System.currentTimeMillis())) // datetime3Column + st.setTimestamp(7, java.sql.Timestamp(System.currentTimeMillis())) // datetime2Column + st.setTimestamp(8, java.sql.Timestamp(System.currentTimeMillis())) // datetimeoffset2Column + st.setBigDecimal(9, BigDecimal("12345.67")) // decimalColumn + st.setFloat(10, 123.45f) // floatColumn + st.setNull(11, java.sql.Types.NULL) // imageColumn (assuming nullable) + st.setInt(12, 123456) // intColumn + st.setBigDecimal(13, BigDecimal("123.45")) // moneyColumn + st.setString(14, "Sample") // ncharColumn + st.setString(15, "Sample$i text") // ntextColumn + st.setBigDecimal(16, BigDecimal("1234.56")) // numericColumn + st.setString(17, "Sample") // nvarcharColumn + st.setString(18, "Sample$i text") // nvarcharMaxColumn + st.setFloat(19, 123.45f) // realColumn + st.setTimestamp(20, java.sql.Timestamp(System.currentTimeMillis())) // smalldatetimeColumn + st.setInt(21, 123) // smallintColumn + st.setBigDecimal(22, BigDecimal("123.45")) // smallmoneyColumn + st.setString(23, "Sample$i text") // textColumn + st.setTime(24, java.sql.Time(System.currentTimeMillis())) // timeColumn + st.setTimestamp(25, java.sql.Timestamp(System.currentTimeMillis())) // timestampColumn + st.setInt(26, 123) // tinyintColumn + // st.setObject(27, null) // udtColumn (assuming nullable) + st.setObject(27, UUID.randomUUID()) // uniqueidentifierColumn + st.setBytes(28, byteArrayOf(0x01, 0x23, 0x45, 0x67, 0x67, 0x67, 0x67, 0x67)) // varbinaryColumn + st.setBytes(29, byteArrayOf(0x01, 0x23, 0x45, 0x67, 0x67, 0x67, 0x67, 0x67)) // varbinaryMaxColumn + st.setString(30, "Sample$i") // varcharColumn + st.setString(31, "Sample$i text") // varcharMaxColumn + st.setString(32, "Sample$i") // xmlColumn + st.setString(33, "SQL_VARIANT") // sqlvariantColumn + st.setBytes( + 34, + @Suppress("ktlint:standard:argument-list-wrapping") + byteArrayOf( + 0xE6.toByte(), 0x10, 0x00, 0x00, 0x01, 0x0C, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, + 0x44, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x09, 0x05, 0x4C, 0x0, + ), + ) // geometryColumn + st.setString(35, "POINT(1 1)") // geographyColumn + st.executeUpdate() + } + } +} + +internal fun tearDownMsSqlTestData(connection: Connection) { + connection.createStatement().use { st -> + st.execute("USE master") + st.execute("DROP DATABASE IF EXISTS $MSSQL_TEST_DATABASE_NAME") + } + connection.close() +} + +abstract class MsSqlTestBase { + protected abstract val connection: Connection + + protected abstract fun connect(database: String? = null): Connection + + @Test + fun `basic test for reading sql tables`() { + val df1 = DataFrame.readSqlTable(connection, "table1", limit = 5).cast() + + val result = df1.filter { "id"() == 1 } + result[0][30] shouldBe "Sample1" + result[0]["bigintColumn"] shouldBe 123456789012345L + result[0]["bitColumn"] shouldBe true + result[0]["intColumn"] shouldBe 123456 + result[0]["ntextColumn"] shouldBe "Sample1 text" + + val schema = DataFrameSchema.readSqlTable(connection, "table1") + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["bigintColumn"]!!.type shouldBe typeOf() + schema.columns["binaryColumn"]!!.type shouldBe typeOf() + schema.columns["bitColumn"]!!.type shouldBe typeOf() + schema.columns["charColumn"]!!.type shouldBe typeOf() + schema.columns["dateColumn"]!!.type shouldBe typeOf() + schema.columns["datetime3Column"]!!.type shouldBe typeOf() + schema.columns["datetime2Column"]!!.type shouldBe typeOf() + schema.columns["datetimeoffset2Column"]!!.type shouldBe typeOf() + schema.columns["decimalColumn"]!!.type shouldBe typeOf() + schema.columns["floatColumn"]!!.type shouldBe typeOf() + schema.columns["imageColumn"]!!.type shouldBe typeOf() + schema.columns["intColumn"]!!.type shouldBe typeOf() + schema.columns["moneyColumn"]!!.type shouldBe typeOf() + schema.columns["ncharColumn"]!!.type shouldBe typeOf() + schema.columns["ntextColumn"]!!.type shouldBe typeOf() + schema.columns["numericColumn"]!!.type shouldBe typeOf() + schema.columns["nvarcharColumn"]!!.type shouldBe typeOf() + schema.columns["nvarcharMaxColumn"]!!.type shouldBe typeOf() + schema.columns["realColumn"]!!.type shouldBe typeOf() + schema.columns["smalldatetimeColumn"]!!.type shouldBe typeOf() + schema.columns["smallintColumn"]!!.type shouldBe typeOf() + schema.columns["smallmoneyColumn"]!!.type shouldBe typeOf() + schema.columns["timeColumn"]!!.type shouldBe typeOf() + schema.columns["timestampColumn"]!!.type shouldBe typeOf() + schema.columns["tinyintColumn"]!!.type shouldBe typeOf() + schema.columns["uniqueidentifierColumn"]!!.type shouldBe typeOf() + schema.columns["varbinaryColumn"]!!.type shouldBe typeOf() + schema.columns["varbinaryMaxColumn"]!!.type shouldBe typeOf() + schema.columns["varcharColumn"]!!.type shouldBe typeOf() + schema.columns["varcharMaxColumn"]!!.type shouldBe typeOf() + schema.columns["xmlColumn"]!!.type shouldBe typeOf() + schema.columns["sqlvariantColumn"]!!.type shouldBe typeOf() + schema.columns["geometryColumn"]!!.type shouldBe typeOf() + schema.columns["geographyColumn"]!!.type shouldBe typeOf() + } + + @Test + fun `read from sql query`() { + @Language("SQL") + val sqlQuery = + """ + SELECT + Table1.id, + Table1.bigintColumn + FROM Table1 + """.trimIndent() + + val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery, limit = 3).cast() + val result = df.filter { "id"() == 1 } + result[0]["bigintColumn"] shouldBe 123456789012345L + + val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["bigintColumn"]!!.type shouldBe typeOf() + } + + @Test + fun `read from all tables`() { + val dataframes = DataFrame.readAllSqlTables(connection, MSSQL_TEST_DATABASE_NAME, 4).values.toList() + + val table1Df = dataframes[0].cast() + + table1Df.rowsCount() shouldBe 4 + table1Df.filter { "id"() > 2 }.rowsCount() shouldBe 2 + table1Df[0]["bigintColumn"] shouldBe 123456789012345L + } + + @Test + fun `infer nullability`() { + inferNullability(connection) + } + + // https://github.com/Kotlin/dataframe/issues/1746 + @Test + fun `readAllSqlTables without catalogue should only return tables from URL database`() { + val secondDb = "testKDFdatabase2" + val testRootConn = connect() + try { + testRootConn.createStatement().use { stmt -> + stmt.executeUpdate( + "IF DB_ID('$secondDb') IS NOT NULL DROP DATABASE $secondDb", + ) + stmt.executeUpdate("CREATE DATABASE $secondDb") + } + connect(secondDb).use { conn2 -> + conn2.createStatement().use { stmt -> + stmt.executeUpdate("CREATE TABLE onlyInDb2 (id INT PRIMARY KEY, val VARCHAR(50))") + } + } + + connect(MSSQL_TEST_DATABASE_NAME).use { scopedConn -> + val tableNames = DataFrame.readAllSqlTables(scopedConn).keys + + tableNames.none { "onlyInDb2" in it } shouldBe true + tableNames.any { "Table1" in it } shouldBe true + } + } finally { + testRootConn.use { conn -> + conn.createStatement().execute( + "IF DB_ID('$secondDb') IS NOT NULL DROP DATABASE $secondDb", + ) + } + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mysqlTestBase.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mysqlTestBase.kt new file mode 100644 index 0000000000..94fdf51bb5 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/mysqlTestBase.kt @@ -0,0 +1,510 @@ +package org.jetbrains.kotlinx.dataframe.io + +import io.kotest.matchers.shouldBe +import kotlinx.datetime.LocalDateTime +import org.intellij.lang.annotations.Language +import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema +import org.jetbrains.kotlinx.dataframe.api.add +import org.jetbrains.kotlinx.dataframe.api.cast +import org.jetbrains.kotlinx.dataframe.api.filter +import org.jetbrains.kotlinx.dataframe.api.select +import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.junit.Test +import java.math.BigDecimal +import java.math.BigInteger +import java.sql.Connection +import java.sql.SQLException +import java.util.Date +import kotlin.reflect.typeOf +import kotlin.time.Instant +import java.sql.Time as SqlTime +import java.sql.Timestamp as SqlTimestamp + +internal const val MYSQL_TEST_DATABASE_NAME = "testKDFdatabase" + +private const val TIMESTAMP = 1726246245460 + +@DataSchema +interface Table1MySql { + val id: Int + val bitCol: Boolean + val tinyintCol: Int + val smallintCol: Int + val mediumintCol: Int + val mediumintUnsignedCol: Int + val integerCol: Int + val intCol: Int + val integerUnsignedCol: Long + val bigintCol: Long + val floatCol: Float + val doubleCol: Double + val decimalCol: BigDecimal + val dateCol: String + val datetimeCol: String + val timestampCol: String + val timeCol: String + val yearCol: String + val varcharCol: String + val charCol: String + val binaryCol: ByteArray + val varbinaryCol: ByteArray + val tinyblobCol: ByteArray + val blobCol: ByteArray + val mediumblobCol: ByteArray + val longblobCol: ByteArray + val textCol: String + val mediumtextCol: String + val longtextCol: String + val enumCol: String + val setCol: Char + val bigintUnsignedCol: BigInteger +} + +@DataSchema +interface Table2MySql { + val id: Int + val bitCol: Boolean? + val tinyintCol: Int? + val smallintCol: Int? + val mediumintCol: Int? + val mediumintUnsignedCol: Int? + val integerCol: Int? + val intCol: Int? + val integerUnsignedCol: Long? + val bigintCol: Long? + val floatCol: Float? + val doubleCol: Double? + val decimalCol: Double? + val dateCol: String? + val datetimeCol: String? + val timestampCol: String? + val timeCol: String? + val yearCol: String? + val varcharCol: String? + val charCol: String? + val binaryCol: ByteArray? + val varbinaryCol: ByteArray? + val tinyblobCol: ByteArray? + val blobCol: ByteArray? + val mediumblobCol: ByteArray? + val longblobCol: ByteArray? + val textCol: String? + val mediumtextCol: String? + val longtextCol: String? + val enumCol: String? + val setCol: Char? + val bigintUnsignedCol: BigInteger? + val jsonCol: String? +} + +@DataSchema +interface Table3MySql { + val id: Int + val enumCol: String + val setCol: Char? +} + +internal fun setUpMySqlTestData(connection: Connection) { + connection.createStatement().use { st -> + // Drop the test database if it exists + val dropDatabaseQuery = "DROP DATABASE IF EXISTS $MYSQL_TEST_DATABASE_NAME" + st.executeUpdate(dropDatabaseQuery) + + // Create the test database + val createDatabaseQuery = "CREATE DATABASE $MYSQL_TEST_DATABASE_NAME" + st.executeUpdate(createDatabaseQuery) + + // Use the newly created database + val useDatabaseQuery = "USE $MYSQL_TEST_DATABASE_NAME" + st.executeUpdate(useDatabaseQuery) + } + + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } + + @Language("SQL") + val createTableQuery = """ + CREATE TABLE IF NOT EXISTS table1 ( + id INT AUTO_INCREMENT PRIMARY KEY, + bitCol BIT NOT NULL, + tinyintCol TINYINT NOT NULL, + smallintCol SMALLINT NOT NULL, + mediumintCol MEDIUMINT NOT NULL, + mediumintUnsignedCol MEDIUMINT UNSIGNED NOT NULL, + integerCol INTEGER NOT NULL, + intCol INT NOT NULL, + integerUnsignedCol INTEGER UNSIGNED NOT NULL, + bigintCol BIGINT NOT NULL, + floatCol FLOAT NOT NULL, + doubleCol DOUBLE NOT NULL, + decimalCol DECIMAL NOT NULL, + dateCol DATE NOT NULL, + datetimeCol DATETIME NOT NULL, + timestampCol TIMESTAMP NOT NULL, + timeCol TIME NOT NULL, + yearCol YEAR NOT NULL, + varcharCol VARCHAR(255) NOT NULL, + charCol CHAR(10) NOT NULL, + binaryCol BINARY(64) NOT NULL, + varbinaryCol VARBINARY(128) NOT NULL, + tinyblobCol TINYBLOB NOT NULL, + blobCol BLOB NOT NULL, + mediumblobCol MEDIUMBLOB NOT NULL , + longblobCol LONGBLOB NOT NULL, + textCol TEXT NOT NULL, + mediumtextCol MEDIUMTEXT NOT NULL, + longtextCol LONGTEXT NOT NULL, + enumCol ENUM('Value1', 'Value2', 'Value3') NOT NULL, + setCol SET('Option1', 'Option2', 'Option3') NOT NULL, + bigintUnsignedCol BIGINT UNSIGNED NOT NULL, + location GEOMETRY, + data JSON + CHECK (JSON_VALID(data)) + ) + """ + + connection.createStatement().execute(createTableQuery.trimIndent()) + + @Language("SQL") + val createTableQuery2 = """ + CREATE TABLE IF NOT EXISTS table2 ( + id INT AUTO_INCREMENT PRIMARY KEY, + bitCol BIT, + tinyintCol TINYINT, + smallintCol SMALLINT, + mediumintCol MEDIUMINT, + mediumintUnsignedCol MEDIUMINT UNSIGNED, + integerCol INTEGER, + intCol INT, + integerUnsignedCol INTEGER UNSIGNED, + bigintCol BIGINT, + floatCol FLOAT, + doubleCol DOUBLE, + decimalCol DECIMAL, + dateCol DATE, + datetimeCol DATETIME, + timestampCol TIMESTAMP, + timeCol TIME, + yearCol YEAR, + varcharCol VARCHAR(255), + charCol CHAR(10), + binaryCol BINARY(64), + varbinaryCol VARBINARY(128), + tinyblobCol TINYBLOB, + blobCol BLOB, + mediumblobCol MEDIUMBLOB, + longblobCol LONGBLOB, + textCol TEXT, + mediumtextCol MEDIUMTEXT, + longtextCol LONGTEXT, + enumCol ENUM('Value1', 'Value2', 'Value3'), + setCol SET('Option1', 'Option2', 'Option3'), + bigintUnsignedCol BIGINT UNSIGNED, + location GEOMETRY, + data JSON + CHECK (JSON_VALID(data)) + ) + """ + + connection.createStatement().execute(createTableQuery2.trimIndent()) + + @Language("SQL") + val insertData1 = + """ + INSERT INTO table1 ( + bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, + integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, + timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, + mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol, location, data + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ST_GeomFromText('POINT(1 1)'), ?) + """.trimIndent() + + @Language("SQL") + val insertData2 = + """ + INSERT INTO table2 ( + bitCol, tinyintCol, smallintCol, mediumintCol, mediumintUnsignedCol, integerCol, intCol, + integerUnsignedCol, bigintCol, floatCol, doubleCol, decimalCol, dateCol, datetimeCol, timestampCol, + timeCol, yearCol, varcharCol, charCol, binaryCol, varbinaryCol, tinyblobCol, blobCol, + mediumblobCol, longblobCol, textCol, mediumtextCol, longtextCol, enumCol, setCol, bigintUnsignedCol, location, data + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ST_GeomFromText('POINT(1 1)'), ?) + """.trimIndent() + + connection.prepareStatement(insertData1).use { st -> + // Insert data into table1 + for (i in 1..3) { + st.setBoolean(1, true) + st.setByte(2, i.toByte()) + st.setShort(3, (i * 10).toShort()) + st.setInt(4, i * 100) + st.setInt(5, i * 100) + st.setInt(6, i * 100) + st.setInt(7, i * 100) + st.setInt(8, i * 100) + st.setInt(9, i * 100) + st.setFloat(10, i * 10.0f) + st.setDouble(11, i * 10.0) + st.setBigDecimal(12, BigDecimal(i * 10)) + st.setDate(13, java.sql.Date(TIMESTAMP)) + st.setTimestamp(14, SqlTimestamp(TIMESTAMP)) + st.setTimestamp(15, SqlTimestamp(TIMESTAMP)) + st.setTime(16, SqlTime(TIMESTAMP)) + st.setInt(17, 2023) + st.setString(18, "varcharValue$i") + st.setString(19, "charValue$i") + st.setBytes(20, "binaryValue".toByteArray()) + st.setBytes(21, "varbinaryValue".toByteArray()) + st.setBytes(22, "tinyblobValue".toByteArray()) + st.setBytes(23, "blobValue".toByteArray()) + st.setBytes(24, "mediumblobValue".toByteArray()) + st.setBytes(25, "longblobValue".toByteArray()) + st.setString(26, "textValue$i") + st.setString(27, "mediumtextValue$i") + st.setString(28, "longtextValue$i") + st.setString(29, "Value$i") + st.setString(30, "Option$i") + st.setObject(31, BigInteger.valueOf((i * 1000).toLong())) + st.setString(32, "{\"key\": \"value\"}") + st.executeUpdate() + } + } + + connection.prepareStatement(insertData2).use { st -> + // Insert data into table2 + for (i in 1..3) { + st.setBoolean(1, false) + st.setByte(2, (i * 2).toByte()) + st.setShort(3, (i * 20).toShort()) + st.setInt(4, i * 200) + st.setInt(5, i * 200) + st.setInt(6, i * 200) + st.setInt(7, i * 200) + st.setInt(8, i * 200) + st.setInt(9, i * 200) + st.setFloat(10, i * 20.0f) + st.setDouble(11, i * 20.0) + st.setBigDecimal(12, BigDecimal(i * 20)) + st.setDate(13, java.sql.Date(TIMESTAMP)) + st.setTimestamp(14, SqlTimestamp(TIMESTAMP)) + st.setTimestamp(15, SqlTimestamp(TIMESTAMP)) + st.setTime(16, SqlTime(TIMESTAMP)) + st.setInt(17, 2023) + st.setString(18, "varcharValue$i") + st.setString(19, "charValue$i") + st.setBytes(20, "binaryValue".toByteArray()) + st.setBytes(21, "varbinaryValue".toByteArray()) + st.setBytes(22, "tinyblobValue".toByteArray()) + st.setBytes(23, "blobValue".toByteArray()) + st.setBytes(24, "mediumblobValue".toByteArray()) + st.setBytes(25, "longblobValue".toByteArray()) + st.setString(26, null) + st.setString(27, null) + st.setString(28, "longtextValue$i") + st.setString(29, "Value$i") + st.setString(30, "Option$i") + st.setObject(31, BigInteger.valueOf((i * 2000).toLong())) + st.setString(32, "{\"key\": \"value\"}") + st.executeUpdate() + } + } +} + +internal fun tearDownMySqlTestData(connection: Connection) { + try { + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } + connection.createStatement().use { st -> + st.execute("DROP DATABASE IF EXISTS $MYSQL_TEST_DATABASE_NAME") + } + connection.close() + } catch (e: SQLException) { + e.printStackTrace() + } +} + +abstract class MySqlTestBase { + protected abstract val connection: Connection + + protected abstract fun connect(database: String? = null): Connection + + @Test + fun `basic test for reading sql tables`() { + val df1 = DataFrame.readSqlTable(connection, "table1").cast() + val result = df1.filter { "id"() == 1 } + result[0][26] shouldBe "textValue1" + result[0][22] shouldBe "tinyblobValue".toByteArray() + + val schema = DataFrameSchema.readSqlTable(connection, "table1") + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["textCol"]!!.type shouldBe typeOf() + schema.columns["dateCol"]!!.type shouldBe typeOf() + schema.columns["datetimeCol"]!!.type shouldBe typeOf() + schema.columns["timestampCol"]!!.type shouldBe typeOf() + schema.columns["timeCol"]!!.type shouldBe typeOf() + schema.columns["yearCol"]!!.type shouldBe typeOf() + schema.columns["textCol"]!!.type shouldBe typeOf() + schema.columns["varbinaryCol"]!!.type shouldBe typeOf() + schema.columns["binaryCol"]!!.type shouldBe typeOf() + schema.columns["longblobCol"]!!.type shouldBe typeOf() + schema.columns["tinyblobCol"]!!.type shouldBe typeOf() + + val df2 = DataFrame.readSqlTable(connection, "table2").cast() + val result2 = df2.filter { "id"() == 1 } + result2[0][26] shouldBe null + + val schema2 = DataFrameSchema.readSqlTable(connection, "table2") + schema2.columns["id"]!!.type shouldBe typeOf() + schema2.columns["textCol"]!!.type shouldBe typeOf() + } + + @Test + fun `read from sql query`() { + @Language("SQL") + val sqlQuery = + """ + SELECT + t1.id, + t1.enumCol, + t2.setCol + FROM table1 t1 + JOIN table2 t2 ON t1.id = t2.id + """.trimIndent() + + val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery).cast() + val result = df.filter { "id"() == 1 } + result[0][2] shouldBe "Option1" + + val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["enumCol"]!!.type shouldBe typeOf() + schema.columns["setCol"]!!.type shouldBe typeOf() + } + + @Test + fun `read from all tables`() { + val dataframes = DataFrame.readAllSqlTables(connection).values.toList() + + val table1Df = dataframes[0].cast() + + table1Df.rowsCount() shouldBe 3 + table1Df.filter { "integerCol"() > 100 }.rowsCount() shouldBe 2 + table1Df[0][11] shouldBe 10.0 + table1Df[0][26] shouldBe "textValue1" + + val table2Df = dataframes[1].cast() + + table2Df.rowsCount() shouldBe 3 + table2Df.filter { + "integerCol"()?.let { it > 400 } ?: false + }.rowsCount() shouldBe 1 + table2Df[0][11] shouldBe 20.0 + table2Df[0][26] shouldBe null + } + + @Test + fun `reading numeric types`() { + val df1 = DataFrame.readSqlTable(connection, "table1").cast() + + val result = df1.select("tinyintCol").add("tinyintCol2") { "tinyintCol"() } + + result[0][1] shouldBe 1.toByte() + + val result1 = df1.select("smallintCol") + .add("smallintCol2") { "smallintCol"() } + + result1[0][1] shouldBe 10.toShort() + + val result2 = df1.select("mediumintCol") + .add("mediumintCol2") { "mediumintCol"() } + + result2[0][1] shouldBe 100 + + val result3 = df1.select("mediumintUnsignedCol") + .add("mediumintUnsignedCol2") { "mediumintUnsignedCol"() } + + result3[0][1] shouldBe 100 + + val result4 = df1.select("integerUnsignedCol") + .add("integerUnsignedCol2") { "integerUnsignedCol"() } + + result4[0][1] shouldBe 100L + + val result5 = df1.select("bigintCol") + .add("bigintCol2") { "bigintCol"() } + + result5[0][1] shouldBe 100 + + val result5a = df1.select("bigintUnsignedCol") + .add("bigintUnsignedCol2") { "bigintUnsignedCol"() } + + result5a[0][1] shouldBe BigInteger.valueOf(1000) + + val result6 = df1.select("floatCol") + .add("floatCol2") { "floatCol"() } + + result6[0][1] shouldBe 10.0f + + val result7 = df1.select("doubleCol") + .add("doubleCol2") { "doubleCol"() } + + result7[0][1] shouldBe 10.0 + + val result8 = df1.select("decimalCol") + .add("decimalCol2") { "decimalCol"() } + + result8[0][1] shouldBe BigDecimal("10") + + val schema = DataFrameSchema.readSqlTable(connection, "table1") + + schema.columns["tinyintCol"]!!.type shouldBe typeOf() + schema.columns["smallintCol"]!!.type shouldBe typeOf() + schema.columns["mediumintCol"]!!.type shouldBe typeOf() + schema.columns["mediumintUnsignedCol"]!!.type shouldBe typeOf() + schema.columns["integerUnsignedCol"]!!.type shouldBe typeOf() + schema.columns["bigintCol"]!!.type shouldBe typeOf() + schema.columns["bigintUnsignedCol"]!!.type shouldBe typeOf() + schema.columns["floatCol"]!!.type shouldBe typeOf() + schema.columns["doubleCol"]!!.type shouldBe typeOf() + schema.columns["decimalCol"]!!.type shouldBe typeOf() + // TODO: all unsigned types + // TODO: new mapping system based on class names + // validation after mapping in getObject + // getObject(i+1, type) catch getObject catch getString + // add direct mapping to getString and other methods + } + + @Test + fun `infer nullability`() { + inferNullability(connection) + } + + // https://github.com/Kotlin/dataframe/issues/1746 + @Test + fun `readAllSqlTables without catalogue should only return tables from URL database`() { + val secondDb = "testKDFdatabase2" + val testRootConn = connect() + try { + testRootConn.createStatement().use { stmt -> + stmt.executeUpdate("DROP DATABASE IF EXISTS $secondDb") + stmt.executeUpdate("CREATE DATABASE $secondDb") + } + connect(secondDb).use { conn2 -> + conn2.createStatement().use { stmt -> + stmt.executeUpdate("CREATE TABLE onlyInDb2 (id INT PRIMARY KEY, val VARCHAR(50))") + } + } + + connect(MYSQL_TEST_DATABASE_NAME).use { scopedConn -> + val tableNames = DataFrame.readAllSqlTables(scopedConn).keys + + tableNames.none { "onlyInDb2" in it } shouldBe true + tableNames.any { "table1" in it } shouldBe true + } + } finally { + testRootConn.use { conn -> + conn.createStatement().execute("DROP DATABASE IF EXISTS $secondDb") + } + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/postgresConnectionUrlTestBase.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/postgresConnectionUrlTestBase.kt new file mode 100644 index 0000000000..23b316f686 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/postgresConnectionUrlTestBase.kt @@ -0,0 +1,107 @@ +package org.jetbrains.kotlinx.dataframe.io + +import io.kotest.assertions.throwables.shouldThrow +import io.kotest.matchers.shouldBe +import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.api.cast +import org.jetbrains.kotlinx.dataframe.api.filter +import org.junit.Test +import org.postgresql.util.PSQLException +import java.sql.DriverManager + +private const val URL_PARAMS = "connectTimeout=10&tcpKeepAlive=true" + +private const val TABLE_NAME = "table1" + +abstract class PostgresConnectionUrlTestBase { + protected abstract val baseUrl: String + + protected abstract val userName: String + + protected abstract val password: String + + private val urlWithLoginPassword: String + get() = "$baseUrl?user=$userName&password=$password&$URL_PARAMS" + private val urlNoLoginPassword: String get() = "$baseUrl?$URL_PARAMS" + private val urlWithPassword: String get() = "$baseUrl?password=$password&$URL_PARAMS" + private val urlWithLogin: String get() = "$baseUrl?user=$userName&$URL_PARAMS" + + @Test + fun `read from table with login and password in connection URL`() { + DriverManager.getConnection(urlWithLoginPassword).use { connection -> + createPostgresTestData(connection) + + val df1 = DataFrame.readSqlTable(connection, TABLE_NAME).cast() + val result1 = df1.filter { "id"() == 1 } + + result1[0][2] shouldBe 11 + + val df2 = connection.readDataFrame(TABLE_NAME).cast() + val result2 = df2.filter { "id"() == 1 } + + result2[0][2] shouldBe 11 + + clearPostgresTestData(connection) + } + } + + @Test + fun `read from table with login and password in connection URL for DBConfig`() { + DriverManager.getConnection(urlWithLoginPassword).use { connection -> + createPostgresTestData(connection) + + val dbConfig = DbConnectionConfig(urlWithLoginPassword) + val df1 = DataFrame.readSqlTable(dbConfig = dbConfig, TABLE_NAME).cast() + val result1 = df1.filter { "id"() == 1 } + + result1[0][2] shouldBe 11 + + val df2 = dbConfig.readDataFrame(TABLE_NAME).cast() + val result2 = df2.filter { "id"() == 1 } + + result2[0][2] shouldBe 11 + + clearPostgresTestData(connection) + } + } + + @Test + fun `read from table without login and password`() { + val dbConfig = DbConnectionConfig(urlNoLoginPassword) + + shouldThrow { + testReadFromTable(dbConfig) + } + } + + @Test + fun `read from table with password only`() { + val dbConfig = DbConnectionConfig(urlWithPassword) + + shouldThrow { + testReadFromTable(dbConfig) + } + } + + @Test + fun `read from table with login only`() { + val dbConfig = DbConnectionConfig(urlWithLogin) + + shouldThrow { + testReadFromTable(dbConfig) + } + } + + private fun testReadFromTable(dbConfig: DbConnectionConfig) { + DriverManager.getConnection(urlWithLoginPassword).use { connection -> + createPostgresTestData(connection) + + val df2 = dbConfig.readDataFrame(TABLE_NAME).cast() + val result2 = df2.filter { "id"() == 1 } + + result2[0][2] shouldBe 11 + + clearPostgresTestData(connection) + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/postgresTestBase.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/postgresTestBase.kt new file mode 100644 index 0000000000..6bdf5ef5e3 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/postgresTestBase.kt @@ -0,0 +1,402 @@ +package org.jetbrains.kotlinx.dataframe.io + +import io.kotest.matchers.shouldBe +import org.intellij.lang.annotations.Language +import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema +import org.jetbrains.kotlinx.dataframe.api.add +import org.jetbrains.kotlinx.dataframe.api.cast +import org.jetbrains.kotlinx.dataframe.api.filter +import org.jetbrains.kotlinx.dataframe.api.select +import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema +import org.junit.Test +import org.postgresql.geometric.PGbox +import org.postgresql.geometric.PGcircle +import org.postgresql.geometric.PGline +import org.postgresql.geometric.PGlseg +import org.postgresql.geometric.PGpath +import org.postgresql.geometric.PGpoint +import org.postgresql.geometric.PGpolygon +import org.postgresql.util.PGInterval +import org.postgresql.util.PGmoney +import org.postgresql.util.PGobject +import java.math.BigDecimal +import java.sql.Connection +import java.sql.Date +import java.sql.SQLException +import java.sql.Types +import java.util.UUID +import kotlin.reflect.typeOf +import java.sql.Time as SqlTime +import java.sql.Timestamp as SqlTimestamp + +@DataSchema +interface Table1Postgres { + val id: Int + val bigintcol: Long + val smallintcol: Int + val bigserialcol: Long + val booleancol: Boolean + val boxcol: PGbox + val byteacol: ByteArray + val charactercol: String + val characterncol: String + val charcol: String + val circlecol: PGcircle + val datecol: java.sql.Date + val doublecol: Double + val integercol: Int? + val intervalcol: String + val jsoncol: String + val jsonbcol: String +} + +@DataSchema +interface Table2Postgres { + val id: Int + val linecol: PGline + val lsegcol: PGlseg + val macaddrcol: String + val moneycol: PGmoney + val numericcol: BigDecimal + val pathcol: PGpath + val pointcol: PGpoint + val polygoncol: PGpolygon + val realcol: Float + val smallintcol: Int + val smallserialcol: Int + val serialcol: Int + val textcol: String? + val timecol: String + val timewithzonecol: String + val timestampcol: String + val timestampwithzonecol: String + val uuidcol: String + val xmlcol: String +} + +@DataSchema +interface ViewTablePostgres { + val id: Int + val bigintcol: Long + val linecol: String + val textCol: String? +} + +internal fun createPostgresTestData(connection: Connection) { + connection.createStatement().use { st -> + st.execute("SET lc_monetary TO 'C'") + st.execute("SET client_encoding TO 'UTF8'") + } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } + + val createTableStatement = """ + CREATE TABLE IF NOT EXISTS table1 ( + id serial PRIMARY KEY, + bigintCol bigint not null, + smallintCol smallint not null, + bigserialCol bigserial not null, + booleanCol boolean not null, + boxCol box not null, + byteaCol bytea not null, + characterCol character not null, + characterNCol character(10) not null, + charCol char not null, + circleCol circle not null, + dateCol date not null, + doubleCol double precision not null, + integerCol integer, + intervalCol interval not null, + jsonCol json not null, + jsonbCol jsonb not null, + intArrayCol integer[], + doubleArrayCol double precision array, + dateArrayCol date array, + textArrayCol text array, + booleanArrayCol boolean array + ) + """ + connection.createStatement().execute(createTableStatement.trimIndent()) + + val createTableQuery = """ + CREATE TABLE IF NOT EXISTS table2 ( + id serial PRIMARY KEY, + lineCol line not null, + lsegCol lseg not null, + macaddrCol macaddr not null, + moneyCol money not null, + numericCol numeric not null, + pathCol path not null, + pointCol point not null, + polygonCol polygon not null, + realCol real not null, + smallintCol smallint not null, + smallserialCol smallserial not null, + serialCol serial not null, + textCol text, + timeCol time not null, + timeWithZoneCol time with time zone not null, + timestampCol timestamp not null, + timestampWithZoneCol timestamp with time zone not null, + uuidCol uuid not null, + xmlCol xml not null + ) + """ + connection.createStatement().execute(createTableQuery.trimIndent()) + + @Language("SQL") + val insertData1 = """ + INSERT INTO table1 ( + bigintCol, smallintCol, bigserialCol, booleanCol, + boxCol, byteaCol, characterCol, characterNCol, charCol, + circleCol, dateCol, doubleCol, + integerCol, intervalCol, jsonCol, jsonbCol, intArrayCol, + doubleArrayCol, dateArrayCol, textArrayCol, booleanArrayCol + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """ + + @Language("SQL") + val insertData2 = """ + INSERT INTO table2 ( + lineCol, lsegCol, macaddrCol, moneyCol, numericCol, + pathCol, pointCol, polygonCol, realCol, smallintCol, + smallserialCol, serialCol, textCol, timeCol, + timeWithZoneCol, timestampCol, timestampWithZoneCol, + uuidCol, xmlCol + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """ + + val intArray = connection.createArrayOf("INTEGER", arrayOf(1, 2, 3)) + val doubleArray = connection.createArrayOf("DOUBLE", arrayOf(1.1, 2.2, 3.3)) + val dateArray = connection.createArrayOf("DATE", arrayOf(Date.valueOf("2023-08-01"), Date.valueOf("2023-08-02"))) + val textArray = connection.createArrayOf("TEXT", arrayOf("Hello", "World")) + val booleanArray = connection.createArrayOf("BOOLEAN", arrayOf(true, false, true)) + + connection.prepareStatement(insertData1).use { st -> + // Insert data into table1 + for (i in 1..3) { + st.setLong(1, i * 1000L) + st.setShort(2, 11.toShort()) + st.setLong(3, 1000000000L + i) + st.setBoolean(4, i % 2 == 1) + st.setObject(5, PGbox("(1,1),(2,2)")) + st.setBytes(6, byteArrayOf(1, 2, 3)) + st.setString(7, "A") + st.setString(8, "Hello") + st.setString(9, "A") + st.setObject(10, PGcircle("<(1,2),3>")) + st.setDate(11, Date.valueOf("2023-08-01")) + st.setDouble(12, 12.34) + st.setInt(13, 12345 * i) + st.setObject(14, PGInterval("1 year")) + + val jsonbObject = PGobject() + jsonbObject.type = "jsonb" + jsonbObject.value = "{\"key\": \"value\"}" + + st.setObject(15, jsonbObject) + st.setObject(16, jsonbObject) + st.setArray(17, intArray) + st.setArray(18, doubleArray) + st.setArray(19, dateArray) + st.setArray(20, textArray) + st.setArray(21, booleanArray) + st.executeUpdate() + } + } + + connection.prepareStatement(insertData2).use { st -> + // Insert data into table2 + for (i in 1..3) { + st.setObject(1, PGline("{1,2,3}")) + st.setObject(2, PGlseg("[(-1,0),(1,0)]")) + + val macaddrObject = PGobject() + macaddrObject.type = "macaddr" + macaddrObject.value = "00:00:00:00:00:0$i" + + st.setObject(3, macaddrObject) + st.setBigDecimal(4, BigDecimal("123.45")) + st.setBigDecimal(5, BigDecimal("12.34")) + st.setObject(6, PGpath("((1,2),(3,$i))")) + st.setObject(7, PGpoint("(1,2)")) + st.setObject(8, PGpolygon("((1,1),(2,2),(3,3))")) + st.setFloat(9, 12.34f) + st.setShort(10, (i * 100).toShort()) + st.setInt(11, 1000 + i) + st.setInt(12, 1000000 + i) + st.setString(13, null) + st.setTime(14, SqlTime.valueOf("12:34:56")) + + st.setTimestamp(15, SqlTimestamp(System.currentTimeMillis())) + st.setTimestamp(16, SqlTimestamp(System.currentTimeMillis())) + st.setTimestamp(17, SqlTimestamp(System.currentTimeMillis())) + + st.setObject(18, UUID.randomUUID(), Types.OTHER) + val xmlObject = PGobject() + xmlObject.type = "xml" + xmlObject.value = "data" + + st.setObject(19, xmlObject) + st.executeUpdate() + } + } +} + +internal fun clearPostgresTestData(connection: Connection) { + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table1") } + connection.createStatement().use { st -> st.execute("DROP TABLE IF EXISTS table2") } +} + +internal fun tearDownPostgresTestData(connection: Connection) { + try { + clearPostgresTestData(connection) + connection.close() + } catch (e: SQLException) { + e.printStackTrace() + } +} + +abstract class PostgresTestBase { + protected abstract val connection: Connection + + @Test + fun `read from tables`() { + val tableName1 = "table1" + val df1 = DataFrame.readSqlTable(connection, tableName1).cast() + val result = df1.filter { "id"() == 1 } + + result[0][2] shouldBe 11 + result[0][13] shouldBe 12345 + result[0][17] shouldBe arrayOf(1, 2, 3) + result[0][18] shouldBe arrayOf(1.1, 2.2, 3.3) + result[0][19] shouldBe arrayOf(Date.valueOf("2023-08-01"), Date.valueOf("2023-08-02")) + result[0][20] shouldBe arrayOf("Hello", "World") + result[0][21] shouldBe arrayOf(true, false, true) + + val schema = DataFrameSchema.readSqlTable(connection, tableName1) + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["integercol"]!!.type shouldBe typeOf() + schema.columns["smallintcol"]!!.type shouldBe typeOf() + schema.columns["circlecol"]!!.type shouldBe typeOf() + schema.columns["intarraycol"]!!.type.classifier shouldBe Array::class + schema.columns["doublearraycol"]!!.type.classifier shouldBe Array::class + schema.columns["datearraycol"]!!.type.classifier shouldBe Array::class + schema.columns["textarraycol"]!!.type.classifier shouldBe Array::class + schema.columns["booleanarraycol"]!!.type.classifier shouldBe Array::class + + val tableName2 = "table2" + val df2 = DataFrame.readSqlTable(connection, tableName2).cast() + val result2 = df2.filter { "id"() == 1 } + result2[0][11] shouldBe 1001 + result2[0][13] shouldBe null + + val schema2 = DataFrameSchema.readSqlTable(connection, tableName2) + schema2.columns["id"]!!.type shouldBe typeOf() + schema2.columns["pathcol"]!!.type shouldBe typeOf() + schema2.columns["textcol"]!!.type shouldBe typeOf() + schema2.columns["linecol"]!!.type shouldBe typeOf() + } + + @Test + fun `read from sql query`() { + @Language("SQL") + val sqlQuery = + """ + SELECT + t1.id, + t1.bigintCol, + t2.lineCol, + t2.textCol + FROM table1 t1 + JOIN table2 t2 ON t1.id = t2.id + """.trimIndent() + + val df = DataFrame.readSqlQuery(connection, sqlQuery = sqlQuery).cast() + val result = df.filter { "id"() == 1 } + result[0][3] shouldBe null + + val schema = DataFrameSchema.readSqlQuery(connection, sqlQuery = sqlQuery) + schema.columns["id"]!!.type shouldBe typeOf() + schema.columns["bigintcol"]!!.type shouldBe typeOf() + schema.columns["textcol"]!!.type shouldBe typeOf() + } + + @Test + fun `read from all tables`() { + val dataframes = DataFrame.readAllSqlTables(connection).values.toList() + + val table1Df = dataframes[0].cast() + + table1Df.rowsCount() shouldBe 3 + table1Df.filter { "integercol"()?.let { it > 12345 } ?: false }.rowsCount() shouldBe 2 + table1Df[0][1] shouldBe 1000L + table1Df[0][2] shouldBe 11 + + val table2Df = dataframes[1].cast() + + table2Df.rowsCount() shouldBe 3 + table2Df.filter { + "pathcol"() == PGpath("((1,2),(3,1))") + }.rowsCount() shouldBe 1 + table2Df[0][11] shouldBe 1001 + } + + @Test + fun `read columns of different types to check type mapping`() { + val tableName1 = "table1" + val df1 = DataFrame.readSqlTable(connection, tableName1).cast() + val result = df1.select("smallintcol") + .add("smallintcol2") { "smallintcol"() } + result[0][1] shouldBe 11 + + val result1 = df1.select("bigserialcol") + .add("bigserialcol2") { "bigserialcol"() } + result1[0][1] shouldBe 1000000001L + + val result2 = df1.select("doublecol") + .add("doublecol2") { "doublecol"() } + result2[0][1] shouldBe 12.34 + + val tableName2 = "table2" + val df2 = DataFrame.readSqlTable(connection, tableName2).cast() + + val result3 = df2.select("moneycol") + .add("moneycol2") { "moneycol"() } + (result3[0][1] as PGmoney).`val` shouldBe 123.45 + + val result4 = df2.select("numericcol") + .add("numericcol2") { "numericcol"() } + result4[0][1] shouldBe BigDecimal("12.34") + + val result5 = df2.select("realcol") + .add("realcol2") { "realcol"() } + result5[0][1] shouldBe 12.34f + + val result7 = df2.select("smallserialcol") + .add("smallserialcol2") { "smallserialcol"() } + result7[0][1] shouldBe 1001 + + val result8 = df2.select("serialcol") + .add("serialcol2") { "serialcol"() } + result8[0][1] shouldBe 1000001 + + val schema = DataFrameSchema.readSqlTable(connection, tableName1) + schema.columns["smallintcol"]!!.type shouldBe typeOf() + schema.columns["bigserialcol"]!!.type shouldBe typeOf() + schema.columns["doublecol"]!!.type shouldBe typeOf() + + val schema1 = DataFrameSchema.readSqlTable(connection, tableName2) + schema1.columns["moneycol"]!!.type shouldBe typeOf() + schema1.columns["numericcol"]!!.type shouldBe typeOf() + schema1.columns["realcol"]!!.type shouldBe typeOf() + schema1.columns["smallserialcol"]!!.type shouldBe typeOf() + schema1.columns["serialcol"]!!.type shouldBe typeOf() + } + + @Test + fun `infer nullability`() { + inferNullability(connection) + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mariadbTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mariadbTest.kt new file mode 100644 index 0000000000..e28959403b --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mariadbTest.kt @@ -0,0 +1,53 @@ +package org.jetbrains.kotlinx.dataframe.io.testcontainers + +import org.jetbrains.kotlinx.dataframe.io.MariadbTestBase +import org.jetbrains.kotlinx.dataframe.io.setUpMariadbTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownMariadbTestData +import org.junit.AfterClass +import org.junit.BeforeClass +import org.testcontainers.mariadb.MariaDBContainer +import java.sql.Connection +import java.sql.DriverManager + +private const val USER_NAME = "root" +private const val PASSWORD = "pass" + +class MariadbContainerTest : MariadbTestBase() { + override val connection: Connection get() = Companion.connection + + override fun connect(database: String?): Connection = openConnection(database) + + companion object { + private val mariadb: MariaDBContainer = MariaDBContainer(BuildConfig.MARIADB_IMAGE).apply { + withUsername(USER_NAME) + withPassword(PASSWORD) + } + + private lateinit var connection: Connection + + private val rootUrl: String + get() = "jdbc:mariadb://${mariadb.host}:${mariadb.firstMappedPort}" + + private fun openConnection(database: String?): Connection = + DriverManager.getConnection( + if (database == null) rootUrl else "$rootUrl/$database", + USER_NAME, + PASSWORD, + ) + + @BeforeClass + @JvmStatic + fun setUpClass() { + mariadb.start() + connection = openConnection(null) + setUpMariadbTestData(connection) + } + + @AfterClass + @JvmStatic + fun tearDownClass() { + tearDownMariadbTestData(connection) + mariadb.stop() + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mssqlTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mssqlTest.kt new file mode 100644 index 0000000000..4a92ca9ac7 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mssqlTest.kt @@ -0,0 +1,51 @@ +package org.jetbrains.kotlinx.dataframe.io.testcontainers + +import org.jetbrains.kotlinx.dataframe.io.MsSqlTestBase +import org.jetbrains.kotlinx.dataframe.io.setUpMsSqlTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownMsSqlTestData +import org.junit.AfterClass +import org.junit.BeforeClass +import org.testcontainers.mssqlserver.MSSQLServerContainer +import java.sql.Connection +import java.sql.DriverManager + +private const val USER_NAME = "sa" +private const val PASSWORD = "A_Str0ng_Required_Password" + +class MsSqlContainerTest : MsSqlTestBase() { + override val connection: Connection get() = Companion.connection + + override fun connect(database: String?): Connection = + DriverManager.getConnection( + if (database == null) rootUrl else "$rootUrl;databaseName=$database", + USER_NAME, + PASSWORD, + ) + + companion object { + private val mssql: MSSQLServerContainer = MSSQLServerContainer(BuildConfig.MSSQL_IMAGE).apply { + acceptLicense() + withPassword(PASSWORD) + } + + private lateinit var connection: Connection + + private val rootUrl: String + get() = "jdbc:sqlserver://${mssql.host}:${mssql.firstMappedPort};encrypt=true;trustServerCertificate=true" + + @BeforeClass + @JvmStatic + fun setUpClass() { + mssql.start() + connection = DriverManager.getConnection(rootUrl, USER_NAME, PASSWORD) + setUpMsSqlTestData(connection) + } + + @AfterClass + @JvmStatic + fun tearDownClass() { + tearDownMsSqlTestData(connection) + mssql.stop() + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mysqlTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mysqlTest.kt new file mode 100644 index 0000000000..9dca340e14 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/mysqlTest.kt @@ -0,0 +1,53 @@ +package org.jetbrains.kotlinx.dataframe.io.testcontainers + +import org.jetbrains.kotlinx.dataframe.io.MySqlTestBase +import org.jetbrains.kotlinx.dataframe.io.setUpMySqlTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownMySqlTestData +import org.junit.AfterClass +import org.junit.BeforeClass +import org.testcontainers.mysql.MySQLContainer +import java.sql.Connection +import java.sql.DriverManager + +private const val USER_NAME = "root" +private const val PASSWORD = "pass" + +class MySqlContainerTest : MySqlTestBase() { + override val connection: Connection get() = Companion.connection + + override fun connect(database: String?): Connection = openConnection(database) + + companion object { + private val mysql: MySQLContainer = MySQLContainer(BuildConfig.MYSQL_IMAGE).apply { + withUsername(USER_NAME) + withPassword(PASSWORD) + } + + private lateinit var connection: Connection + + private val rootUrl: String + get() = "jdbc:mysql://${mysql.host}:${mysql.firstMappedPort}" + + private fun openConnection(database: String?): Connection = + DriverManager.getConnection( + if (database == null) rootUrl else "$rootUrl/$database", + USER_NAME, + PASSWORD, + ) + + @BeforeClass + @JvmStatic + fun setUpClass() { + mysql.start() + connection = openConnection(null) + setUpMySqlTestData(connection) + } + + @AfterClass + @JvmStatic + fun tearDownClass() { + tearDownMySqlTestData(connection) + mysql.stop() + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/postgresConnectionUrlTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/postgresConnectionUrlTest.kt new file mode 100644 index 0000000000..4606a250b6 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/postgresConnectionUrlTest.kt @@ -0,0 +1,39 @@ +package org.jetbrains.kotlinx.dataframe.io.testcontainers + +import org.jetbrains.kotlinx.dataframe.io.PostgresConnectionUrlTestBase +import org.junit.AfterClass +import org.junit.BeforeClass +import org.testcontainers.postgresql.PostgreSQLContainer + +private const val USER_NAME = "postgres" +private const val PASSWORD = "pass" +private const val DATABASE_NAME = "test" + +class PostgresConnectionUrlContainerTest : PostgresConnectionUrlTestBase() { + override val baseUrl: String + get() = "jdbc:postgresql://${postgres.host}:${postgres.firstMappedPort}/$DATABASE_NAME" + + override val userName: String get() = USER_NAME + + override val password: String get() = PASSWORD + + companion object { + private val postgres: PostgreSQLContainer = PostgreSQLContainer(BuildConfig.POSTGRES_IMAGE).apply { + withDatabaseName(DATABASE_NAME) + withUsername(USER_NAME) + withPassword(PASSWORD) + } + + @BeforeClass + @JvmStatic + fun setUpClass() { + postgres.start() + } + + @AfterClass + @JvmStatic + fun tearDownClass() { + postgres.stop() + } + } +} diff --git a/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/postgresTest.kt b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/postgresTest.kt new file mode 100644 index 0000000000..bed0675064 --- /dev/null +++ b/dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/testcontainers/postgresTest.kt @@ -0,0 +1,43 @@ +package org.jetbrains.kotlinx.dataframe.io.testcontainers + +import org.jetbrains.kotlinx.dataframe.io.PostgresTestBase +import org.jetbrains.kotlinx.dataframe.io.createPostgresTestData +import org.jetbrains.kotlinx.dataframe.io.tearDownPostgresTestData +import org.junit.AfterClass +import org.junit.BeforeClass +import org.testcontainers.postgresql.PostgreSQLContainer +import java.sql.Connection +import java.sql.DriverManager + +private const val USER_NAME = "postgres" +private const val PASSWORD = "pass" +private const val DATABASE_NAME = "test" + +class PostgresContainerTest : PostgresTestBase() { + override val connection: Connection get() = Companion.connection + + companion object { + private val postgres: PostgreSQLContainer = PostgreSQLContainer(BuildConfig.POSTGRES_IMAGE).apply { + withDatabaseName(DATABASE_NAME) + withUsername(USER_NAME) + withPassword(PASSWORD) + } + + private lateinit var connection: Connection + + @BeforeClass + @JvmStatic + fun setUpClass() { + postgres.start() + connection = DriverManager.getConnection(postgres.jdbcUrl, USER_NAME, PASSWORD) + createPostgresTestData(connection) + } + + @AfterClass + @JvmStatic + fun tearDownClass() { + tearDownPostgresTestData(connection) + postgres.stop() + } + } +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 1d74e7794c..477b1008e3 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -62,6 +62,14 @@ ktor = "3.5.1" kotlin-compile-testing = "0.7.1" # unused for now hikari = "7.1.0" duckdb = "1.5.4.0" +testcontainers = "2.0.5" + +# Docker image tags of the databases used in the Testcontainers tests +dockerImage-mariadb = "12.3.2" # available tags: https://hub.docker.com/_/mariadb/tags +dockerImage-mysql = "9.7" # available tags: https://hub.docker.com/_/mysql/tags +dockerImage-postgres = "18-alpine" # available tags: https://hub.docker.com/_/postgres/tags +dockerImage-mssql = "2022-CU22-ubuntu-22.04" # available tags: https://mcr.microsoft.com/en-us/artifact/mar/mssql/server + buildconfig = "6.0.10" benchmark = "0.4.17" @@ -120,6 +128,12 @@ mysql = { group = "com.mysql", name = "mysql-connector-j", version.ref = "mysql" postgresql = { group = "org.postgresql", name = "postgresql", version.ref = "postgresql" } sqlite = { group = "org.xerial", name = "sqlite-jdbc", version.ref = "sqlite" } +testcontainers = { group = "org.testcontainers", name = "testcontainers", version.ref = "testcontainers" } +testcontainers-postgresql = { group = "org.testcontainers", name = "testcontainers-postgresql", version.ref = "testcontainers" } +testcontainers-mysql = { group = "org.testcontainers", name = "testcontainers-mysql", version.ref = "testcontainers" } +testcontainers-mariadb = { group = "org.testcontainers", name = "testcontainers-mariadb", version.ref = "testcontainers" } +testcontainers-mssqlserver = { group = "org.testcontainers", name = "testcontainers-mssqlserver", version.ref = "testcontainers" } + kandy = { group = "org.jetbrains.kotlinx", name = "kandy-lets-plot", version.ref = "kandy" } kandy-geo = { group = "org.jetbrains.kotlinx", name = "kandy-geo", version.ref = "kandy" } kandy-samples-utils = { group = "org.jetbrains.kotlinx", name = "kandy-samples-utils", version.ref = "kandy" }