Skip to content

Commit 7a895e2

Browse files
authored
Added this test for other databases (#954)
Refactor inferNullability test logic Extracted common inferNullability test logic to `commonTestScenarios.kt` for reusability. Removed redundant code from individual test files and added necessary imports to support the new structure.
1 parent c4bb29c commit 7a895e2

File tree

10 files changed

+172
-368
lines changed

10 files changed

+172
-368
lines changed
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
package org.jetbrains.kotlinx.dataframe.io
2+
3+
import io.kotest.matchers.shouldBe
4+
import org.intellij.lang.annotations.Language
5+
import org.jetbrains.kotlinx.dataframe.DataFrame
6+
import org.jetbrains.kotlinx.dataframe.api.schema
7+
import org.jetbrains.kotlinx.dataframe.io.db.MsSql
8+
import java.sql.Connection
9+
import java.sql.ResultSet
10+
import kotlin.reflect.typeOf
11+
12+
internal fun inferNullability(connection: Connection) {
13+
// prepare tables and data
14+
@Language("SQL")
15+
val createTestTable1Query = """
16+
CREATE TABLE TestTable1 (
17+
id INT PRIMARY KEY,
18+
name VARCHAR(50),
19+
surname VARCHAR(50),
20+
age INT NOT NULL
21+
)
22+
"""
23+
24+
connection.createStatement().execute(createTestTable1Query)
25+
26+
connection.createStatement()
27+
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (1, 'John', 'Crawford', 40)")
28+
connection.createStatement()
29+
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (2, 'Alice', 'Smith', 25)")
30+
connection.createStatement()
31+
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (3, 'Bob', 'Johnson', 47)")
32+
connection.createStatement()
33+
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (4, 'Sam', NULL, 15)")
34+
35+
// start testing `readSqlTable` method
36+
37+
// with default inferNullability: Boolean = true
38+
val tableName = "TestTable1"
39+
val df = DataFrame.readSqlTable(connection, tableName)
40+
df.schema().columns["id"]!!.type shouldBe typeOf<Int>()
41+
df.schema().columns["name"]!!.type shouldBe typeOf<String>()
42+
df.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
43+
df.schema().columns["age"]!!.type shouldBe typeOf<Int>()
44+
45+
val dataSchema = DataFrame.getSchemaForSqlTable(connection, tableName)
46+
dataSchema.columns.size shouldBe 4
47+
dataSchema.columns["id"]!!.type shouldBe typeOf<Int>()
48+
dataSchema.columns["name"]!!.type shouldBe typeOf<String?>()
49+
dataSchema.columns["surname"]!!.type shouldBe typeOf<String?>()
50+
dataSchema.columns["age"]!!.type shouldBe typeOf<Int>()
51+
52+
// with inferNullability: Boolean = false
53+
val df1 = DataFrame.readSqlTable(connection, tableName, inferNullability = false)
54+
df1.schema().columns["id"]!!.type shouldBe typeOf<Int>()
55+
56+
// this column changed a type because it doesn't contain nulls
57+
df1.schema().columns["name"]!!.type shouldBe typeOf<String?>()
58+
df1.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
59+
df1.schema().columns["age"]!!.type shouldBe typeOf<Int>()
60+
61+
// end testing `readSqlTable` method
62+
63+
// start testing `readSQLQuery` method
64+
65+
// ith default inferNullability: Boolean = true
66+
@Language("SQL")
67+
val sqlQuery =
68+
"""
69+
SELECT name, surname, age FROM TestTable1
70+
""".trimIndent()
71+
72+
val df2 = DataFrame.readSqlQuery(connection, sqlQuery)
73+
df2.schema().columns["name"]!!.type shouldBe typeOf<String>()
74+
df2.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
75+
df2.schema().columns["age"]!!.type shouldBe typeOf<Int>()
76+
77+
val dataSchema2 = DataFrame.getSchemaForSqlQuery(connection, sqlQuery)
78+
dataSchema2.columns.size shouldBe 3
79+
dataSchema2.columns["name"]!!.type shouldBe typeOf<String?>()
80+
dataSchema2.columns["surname"]!!.type shouldBe typeOf<String?>()
81+
dataSchema2.columns["age"]!!.type shouldBe typeOf<Int>()
82+
83+
// with inferNullability: Boolean = false
84+
val df3 = DataFrame.readSqlQuery(connection, sqlQuery, inferNullability = false)
85+
// this column changed a type because it doesn't contain nulls
86+
df3.schema().columns["name"]!!.type shouldBe typeOf<String?>()
87+
df3.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
88+
df3.schema().columns["age"]!!.type shouldBe typeOf<Int>()
89+
90+
// end testing `readSQLQuery` method
91+
92+
// start testing `readResultSet` method
93+
94+
connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE).use { st ->
95+
@Language("SQL")
96+
val selectStatement = "SELECT * FROM TestTable1"
97+
98+
st.executeQuery(selectStatement).use { rs ->
99+
// ith default inferNullability: Boolean = true
100+
val df4 = DataFrame.readResultSet(rs, MsSql)
101+
df4.schema().columns["id"]!!.type shouldBe typeOf<Int>()
102+
df4.schema().columns["name"]!!.type shouldBe typeOf<String>()
103+
df4.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
104+
df4.schema().columns["age"]!!.type shouldBe typeOf<Int>()
105+
106+
rs.beforeFirst()
107+
108+
val dataSchema3 = DataFrame.getSchemaForResultSet(rs, MsSql)
109+
dataSchema3.columns.size shouldBe 4
110+
dataSchema3.columns["id"]!!.type shouldBe typeOf<Int>()
111+
dataSchema3.columns["name"]!!.type shouldBe typeOf<String?>()
112+
dataSchema3.columns["surname"]!!.type shouldBe typeOf<String?>()
113+
dataSchema3.columns["age"]!!.type shouldBe typeOf<Int>()
114+
115+
// with inferNullability: Boolean = false
116+
rs.beforeFirst()
117+
118+
val df5 = DataFrame.readResultSet(rs, MsSql, inferNullability = false)
119+
df5.schema().columns["id"]!!.type shouldBe typeOf<Int>()
120+
121+
// this column changed a type because it doesn't contain nulls
122+
df5.schema().columns["name"]!!.type shouldBe typeOf<String?>()
123+
df5.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
124+
df5.schema().columns["age"]!!.type shouldBe typeOf<Int>()
125+
}
126+
}
127+
// end testing `readResultSet` method
128+
129+
connection.createStatement().execute("DROP TABLE TestTable1")
130+
}

dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/h2/h2Test.kt

Lines changed: 2 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
1010
import org.jetbrains.kotlinx.dataframe.api.add
1111
import org.jetbrains.kotlinx.dataframe.api.cast
1212
import org.jetbrains.kotlinx.dataframe.api.filter
13-
import org.jetbrains.kotlinx.dataframe.api.schema
1413
import org.jetbrains.kotlinx.dataframe.api.select
1514
import org.jetbrains.kotlinx.dataframe.io.DbConnectionConfig
1615
import org.jetbrains.kotlinx.dataframe.io.db.H2
@@ -20,6 +19,7 @@ import org.jetbrains.kotlinx.dataframe.io.getSchemaForAllSqlTables
2019
import org.jetbrains.kotlinx.dataframe.io.getSchemaForResultSet
2120
import org.jetbrains.kotlinx.dataframe.io.getSchemaForSqlQuery
2221
import org.jetbrains.kotlinx.dataframe.io.getSchemaForSqlTable
22+
import org.jetbrains.kotlinx.dataframe.io.inferNullability
2323
import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables
2424
import org.jetbrains.kotlinx.dataframe.io.readDataFrame
2525
import org.jetbrains.kotlinx.dataframe.io.readResultSet
@@ -841,128 +841,9 @@ class JdbcTest {
841841
saleDataSchema1.columns["amount"]!!.type shouldBe typeOf<BigDecimal>()
842842
}
843843

844-
// TODO: add the same test for each particular database and refactor the scenario to the common test case
845-
// https://github.com/Kotlin/dataframe/issues/688
846844
@Test
847845
fun `infer nullability`() {
848-
// prepare tables and data
849-
@Language("SQL")
850-
val createTestTable1Query = """
851-
CREATE TABLE TestTable1 (
852-
id INT PRIMARY KEY,
853-
name VARCHAR(50),
854-
surname VARCHAR(50),
855-
age INT NOT NULL
856-
)
857-
"""
858-
859-
connection.createStatement().execute(createTestTable1Query)
860-
861-
connection.createStatement()
862-
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (1, 'John', 'Crawford', 40)")
863-
connection.createStatement()
864-
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (2, 'Alice', 'Smith', 25)")
865-
connection.createStatement()
866-
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (3, 'Bob', 'Johnson', 47)")
867-
connection.createStatement()
868-
.execute("INSERT INTO TestTable1 (id, name, surname, age) VALUES (4, 'Sam', NULL, 15)")
869-
870-
// start testing `readSqlTable` method
871-
872-
// with default inferNullability: Boolean = true
873-
val tableName = "TestTable1"
874-
val df = DataFrame.readSqlTable(connection, tableName)
875-
df.schema().columns["id"]!!.type shouldBe typeOf<Int>()
876-
df.schema().columns["name"]!!.type shouldBe typeOf<String>()
877-
df.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
878-
df.schema().columns["age"]!!.type shouldBe typeOf<Int>()
879-
880-
val dataSchema = DataFrame.getSchemaForSqlTable(connection, tableName)
881-
dataSchema.columns.size shouldBe 4
882-
dataSchema.columns["id"]!!.type shouldBe typeOf<Int>()
883-
dataSchema.columns["name"]!!.type shouldBe typeOf<String?>()
884-
dataSchema.columns["surname"]!!.type shouldBe typeOf<String?>()
885-
dataSchema.columns["age"]!!.type shouldBe typeOf<Int>()
886-
887-
// with inferNullability: Boolean = false
888-
val df1 = DataFrame.readSqlTable(connection, tableName, inferNullability = false)
889-
df1.schema().columns["id"]!!.type shouldBe typeOf<Int>()
890-
891-
// this column changed a type because it doesn't contain nulls
892-
df1.schema().columns["name"]!!.type shouldBe typeOf<String?>()
893-
df1.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
894-
df1.schema().columns["age"]!!.type shouldBe typeOf<Int>()
895-
896-
// end testing `readSqlTable` method
897-
898-
// start testing `readSQLQuery` method
899-
900-
// ith default inferNullability: Boolean = true
901-
@Language("SQL")
902-
val sqlQuery =
903-
"""
904-
SELECT name, surname, age FROM TestTable1
905-
""".trimIndent()
906-
907-
val df2 = DataFrame.readSqlQuery(connection, sqlQuery)
908-
df2.schema().columns["name"]!!.type shouldBe typeOf<String>()
909-
df2.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
910-
df2.schema().columns["age"]!!.type shouldBe typeOf<Int>()
911-
912-
val dataSchema2 = DataFrame.getSchemaForSqlQuery(connection, sqlQuery)
913-
dataSchema2.columns.size shouldBe 3
914-
dataSchema2.columns["name"]!!.type shouldBe typeOf<String?>()
915-
dataSchema2.columns["surname"]!!.type shouldBe typeOf<String?>()
916-
dataSchema2.columns["age"]!!.type shouldBe typeOf<Int>()
917-
918-
// with inferNullability: Boolean = false
919-
val df3 = DataFrame.readSqlQuery(connection, sqlQuery, inferNullability = false)
920-
921-
// this column changed a type because it doesn't contain nulls
922-
df3.schema().columns["name"]!!.type shouldBe typeOf<String?>()
923-
df3.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
924-
df3.schema().columns["age"]!!.type shouldBe typeOf<Int>()
925-
926-
// end testing `readSQLQuery` method
927-
928-
// start testing `readResultSet` method
929-
930-
connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE).use { st ->
931-
@Language("SQL")
932-
val selectStatement = "SELECT * FROM TestTable1"
933-
934-
st.executeQuery(selectStatement).use { rs ->
935-
// ith default inferNullability: Boolean = true
936-
val df4 = DataFrame.readResultSet(rs, H2(MySql))
937-
df4.schema().columns["id"]!!.type shouldBe typeOf<Int>()
938-
df4.schema().columns["name"]!!.type shouldBe typeOf<String>()
939-
df4.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
940-
df4.schema().columns["age"]!!.type shouldBe typeOf<Int>()
941-
942-
rs.beforeFirst()
943-
944-
val dataSchema3 = DataFrame.getSchemaForResultSet(rs, H2(MySql))
945-
dataSchema3.columns.size shouldBe 4
946-
dataSchema3.columns["id"]!!.type shouldBe typeOf<Int>()
947-
dataSchema3.columns["name"]!!.type shouldBe typeOf<String?>()
948-
dataSchema3.columns["surname"]!!.type shouldBe typeOf<String?>()
949-
dataSchema3.columns["age"]!!.type shouldBe typeOf<Int>()
950-
951-
// with inferNullability: Boolean = false
952-
rs.beforeFirst()
953-
954-
val df5 = DataFrame.readResultSet(rs, H2(MySql), inferNullability = false)
955-
df5.schema().columns["id"]!!.type shouldBe typeOf<Int>()
956-
957-
// this column changed a type because it doesn't contain nulls
958-
df5.schema().columns["name"]!!.type shouldBe typeOf<String?>()
959-
df5.schema().columns["surname"]!!.type shouldBe typeOf<String?>()
960-
df5.schema().columns["age"]!!.type shouldBe typeOf<Int>()
961-
}
962-
}
963-
// end testing `readResultSet` method
964-
965-
connection.createStatement().execute("DROP TABLE TestTable1")
846+
inferNullability(connection)
966847
}
967848

968849
@Test

dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/h2/mariadbH2Test.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.api.filter
1010
import org.jetbrains.kotlinx.dataframe.api.select
1111
import org.jetbrains.kotlinx.dataframe.io.getSchemaForSqlQuery
1212
import org.jetbrains.kotlinx.dataframe.io.getSchemaForSqlTable
13+
import org.jetbrains.kotlinx.dataframe.io.inferNullability
1314
import org.jetbrains.kotlinx.dataframe.io.readAllSqlTables
1415
import org.jetbrains.kotlinx.dataframe.io.readSqlQuery
1516
import org.jetbrains.kotlinx.dataframe.io.readSqlTable
@@ -417,4 +418,9 @@ class MariadbH2Test {
417418
schema.columns["doublecol"]!!.type shouldBe typeOf<Double>()
418419
schema.columns["decimalcol"]!!.type shouldBe typeOf<BigDecimal>()
419420
}
421+
422+
@Test
423+
fun `infer nullability`() {
424+
inferNullability(connection)
425+
}
420426
}

0 commit comments

Comments
 (0)