Skip to content

Commit 651e7ec

Browse files
aparajonsvc-squareup-copybara
aparajon
authored and
svc-squareup-copybara
committed
VitessTestDb: expose query executor methods
GitOrigin-RevId: 2c93c1d4999824adfe0d25b8ef101f26e028e2bc
1 parent a1a6e8f commit 651e7ec

File tree

9 files changed

+137
-75
lines changed

9 files changed

+137
-75
lines changed

‎misk-vitess/api/misk-vitess.api‎

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -218,13 +218,43 @@ public final class misk/vitess/testing/VSchemaUpdate {
218218
public fun toString ()Ljava/lang/String;
219219
}
220220

221+
public final class misk/vitess/testing/VitessTable {
222+
public fun <init> (Ljava/lang/String;Lmisk/vitess/testing/VitessTableType;)V
223+
public final fun component1 ()Ljava/lang/String;
224+
public final fun component2 ()Lmisk/vitess/testing/VitessTableType;
225+
public final fun copy (Ljava/lang/String;Lmisk/vitess/testing/VitessTableType;)Lmisk/vitess/testing/VitessTable;
226+
public static synthetic fun copy$default (Lmisk/vitess/testing/VitessTable;Ljava/lang/String;Lmisk/vitess/testing/VitessTableType;ILjava/lang/Object;)Lmisk/vitess/testing/VitessTable;
227+
public fun equals (Ljava/lang/Object;)Z
228+
public final fun getTableName ()Ljava/lang/String;
229+
public final fun getType ()Lmisk/vitess/testing/VitessTableType;
230+
public fun hashCode ()I
231+
public fun toString ()Ljava/lang/String;
232+
}
233+
234+
public final class misk/vitess/testing/VitessTableType : java/lang/Enum {
235+
public static final field REFERENCE Lmisk/vitess/testing/VitessTableType;
236+
public static final field SEQUENCE Lmisk/vitess/testing/VitessTableType;
237+
public static final field STANDARD Lmisk/vitess/testing/VitessTableType;
238+
public static fun getEntries ()Lkotlin/enums/EnumEntries;
239+
public static fun valueOf (Ljava/lang/String;)Lmisk/vitess/testing/VitessTableType;
240+
public static fun values ()[Lmisk/vitess/testing/VitessTableType;
241+
}
242+
221243
public final class misk/vitess/testing/VitessTestDb {
222244
public static final field Companion Lmisk/vitess/testing/VitessTestDb$Companion;
223245
public fun <init> ()V
224246
public fun <init> (ZLjava/lang/String;ZZZZZLjava/lang/String;ILjava/lang/String;Ljava/lang/String;Lmisk/vitess/testing/TransactionIsolationLevel;Ljava/time/Duration;Ljava/lang/String;I)V
225247
public synthetic fun <init> (ZLjava/lang/String;ZZZZZLjava/lang/String;ILjava/lang/String;Ljava/lang/String;Lmisk/vitess/testing/TransactionIsolationLevel;Ljava/time/Duration;Ljava/lang/String;IILkotlin/jvm/internal/DefaultConstructorMarker;)V
226248
public static final fun Builder ()Lmisk/vitess/testing/VitessTestDb$Builder;
227249
public final fun applySchema ()Lmisk/vitess/testing/ApplySchemaResult;
250+
public final fun executeQuery (Ljava/lang/String;Ljava/lang/String;)Ljava/util/List;
251+
public static synthetic fun executeQuery$default (Lmisk/vitess/testing/VitessTestDb;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Ljava/util/List;
252+
public final fun executeTransaction (Ljava/lang/String;Ljava/lang/String;)Z
253+
public static synthetic fun executeTransaction$default (Lmisk/vitess/testing/VitessTestDb;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Z
254+
public final fun executeUpdate (Ljava/lang/String;Ljava/lang/String;)I
255+
public static synthetic fun executeUpdate$default (Lmisk/vitess/testing/VitessTestDb;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)I
256+
public final fun getKeyspaces ()Ljava/util/List;
257+
public final fun getTables (Ljava/lang/String;)Ljava/util/List;
228258
public final fun run ()Lmisk/vitess/testing/VitessTestDbStartupResult;
229259
public final fun shutdown ()Lmisk/vitess/testing/VitessTestDbShutdownResult;
230260
public final fun truncate ()V
@@ -386,15 +416,6 @@ public final class misk/vitess/testing/internal/VitessSchemaManagerException : j
386416
public synthetic fun <init> (Ljava/lang/String;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
387417
}
388418

389-
public final class misk/vitess/testing/internal/VitessTableType : java/lang/Enum {
390-
public static final field REFERENCE Lmisk/vitess/testing/internal/VitessTableType;
391-
public static final field SEQUENCE Lmisk/vitess/testing/internal/VitessTableType;
392-
public static final field STANDARD Lmisk/vitess/testing/internal/VitessTableType;
393-
public static fun getEntries ()Lkotlin/enums/EnumEntries;
394-
public static fun valueOf (Ljava/lang/String;)Lmisk/vitess/testing/internal/VitessTableType;
395-
public static fun values ()[Lmisk/vitess/testing/internal/VitessTableType;
396-
}
397-
398419
public final class misk/vitess/testing/internal/VschemaAdapter {
399420
public fun <init> ()V
400421
public final fun fromJson (Ljava/lang/String;)Ljava/util/Map;

‎misk-vitess/src/test/kotlin/misk/vitess/testing/CustomArgsTest.kt‎

Lines changed: 15 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
package misk.vitess.testing
22

3-
import misk.vitess.testing.internal.VitessClusterConfig
4-
import misk.vitess.testing.internal.VitessQueryExecutor
53
import misk.vitess.testing.internal.VitessQueryExecutorException
64
import org.junit.jupiter.api.AfterAll
75
import org.junit.jupiter.api.Assertions.assertArrayEquals
@@ -24,8 +22,6 @@ class CustomArgsTest {
2422
private lateinit var testDb2: VitessTestDb
2523
private lateinit var testDb1RunResult: VitessTestDbStartupResult
2624
private lateinit var testDb2RunResult: VitessTestDbStartupResult
27-
private lateinit var testDb1QueryExecutor: VitessQueryExecutor
28-
private lateinit var testDb2QueryExecutor: VitessQueryExecutor
2925

3026
// testDb1 args
3127
private const val DB1_CONTAINER_NAME = "custom_args_test_vitess_db"
@@ -69,9 +65,6 @@ class CustomArgsTest {
6965
val future2: Future<VitessTestDbStartupResult> = executorService.submit<VitessTestDbStartupResult> { testDb2.run() }
7066
testDb1RunResult = future1.get()
7167
testDb2RunResult = future2.get()
72-
73-
testDb1QueryExecutor = VitessQueryExecutor(VitessClusterConfig(DB1_PORT))
74-
testDb2QueryExecutor = VitessQueryExecutor(VitessClusterConfig(DB2_PORT))
7568
}
7669

7770
@JvmStatic
@@ -84,40 +77,40 @@ class CustomArgsTest {
8477

8578
@Test
8679
fun `test user defined MySql version`() {
87-
val results = testDb1QueryExecutor.executeQuery("SELECT version();")
80+
val results = testDb1.executeQuery("SELECT version();")
8881
val actualMysqlVersion = results[0]["version()"]
8982
assertEquals("$DB1_MYSQL_VERSION-Vitess", actualMysqlVersion)
9083
}
9184

9285
@Test
9386
fun `test user defined SQL mode`() {
94-
val results = testDb1QueryExecutor.executeQuery("SELECT @@global.sql_mode;")
87+
val results = testDb1.executeQuery("SELECT @@global.sql_mode;")
9588
val actualSqlMode = results[0]["@@global.sql_mode"]
9689
assertEquals(DB1_SQL_MODE, actualSqlMode)
9790
}
9891

9992
@Test
10093
fun `test user defined transaction isolation level`() {
101-
val results = testDb1QueryExecutor.executeQuery("SELECT @@global.transaction_ISOLATION;")
94+
val results = testDb1.executeQuery("SELECT @@global.transaction_ISOLATION;")
10295
val actualTransactionIsolationLevel = results[0]["@@global.transaction_ISOLATION"]
10396
assertEquals(DB1_TXN_ISO_LEVEL.value, actualTransactionIsolationLevel)
10497
}
10598

10699
@Test
107100
fun `test scatter queries fail by default but succeed with query hint`() {
108101
testDb1.truncate()
109-
testDb1QueryExecutor.executeUpdate("INSERT INTO customers (email, token) VALUES ('[email protected]', 'token');")
102+
testDb1.executeUpdate("INSERT INTO customers (email, token) VALUES ('[email protected]', 'token');")
110103

111-
val exception = assertThrows<VitessQueryExecutorException> { testDb1QueryExecutor.executeQuery( "SELECT * FROM customers;") }
104+
val exception = assertThrows<VitessQueryExecutorException> { testDb1.executeQuery( "SELECT * FROM customers;") }
112105
assertTrue(exception.cause?.message!!.contains("plan includes scatter, which is disallowed"))
113106

114-
val results = testDb1QueryExecutor.executeQuery("SELECT /*vt+ ALLOW_SCATTER */ * FROM customers;")
107+
val results = testDb1.executeQuery("SELECT /*vt+ ALLOW_SCATTER */ * FROM customers;")
115108
assertEquals(1, results.size )
116109
}
117110

118111
@Test
119112
fun `test transaction timeout`() {
120-
val exception = assertThrows<VitessQueryExecutorException> { testDb1QueryExecutor.executeTransaction("SELECT SLEEP(6);") }
113+
val exception = assertThrows<VitessQueryExecutorException> { testDb1.executeTransaction("SELECT SLEEP(6);") }
121114
val actualMessage = exception.cause?.message!!
122115
val expectedMessageSubstring = "maximum statement execution time exceeded"
123116
assertTrue(
@@ -137,22 +130,20 @@ class CustomArgsTest {
137130

138131
@Test
139132
fun `test declarative schema changes are auto applied`() {
140-
val vitessQueryExecutor = VitessQueryExecutor(VitessClusterConfig(DB1_PORT))
141-
val keyspaces = vitessQueryExecutor.getKeyspaces()
133+
val keyspaces = testDb1.getKeyspaces()
142134
assertArrayEquals(arrayOf("gameworld", "gameworld_sharded"), keyspaces.toTypedArray())
143135

144-
val unshardedTables = vitessQueryExecutor.getTables("gameworld").map { it.tableName }
136+
val unshardedTables = testDb1.getTables("gameworld").map { it.tableName }
145137
assertArrayEquals(arrayOf("customers_seq", "games_seq"), unshardedTables.toTypedArray())
146138

147-
val shardedTables = vitessQueryExecutor.getTables("gameworld_sharded").map { it.tableName }
139+
val shardedTables = testDb1.getTables("gameworld_sharded").map { it.tableName }
148140
assertArrayEquals(arrayOf("customers", "games"), shardedTables.toTypedArray())
149141
}
150142

151143
@Test
152144
fun `test applySchema after run with keepAlive = false`() {
153-
val vitessQueryExecutor = VitessQueryExecutor(VitessClusterConfig(DB2_PORT))
154145
// Keyspaces are still applied even if autoApplySchemaChanges is false.
155-
var keyspaces = vitessQueryExecutor.getKeyspaces()
146+
var keyspaces = testDb2.getKeyspaces()
156147
assertArrayEquals(arrayOf("gameworld", "gameworld_sharded"), keyspaces.toTypedArray())
157148

158149
// However tables should not yet be applied.
@@ -213,15 +204,15 @@ class CustomArgsTest {
213204
}
214205

215206
private fun assertTablesNotApplied() {
216-
assertEquals(testDb2QueryExecutor.getTables("gameworld").size, 0)
217-
assertEquals(testDb2QueryExecutor.getTables("gameworld_sharded").size, 0)
207+
assertEquals(testDb2.getTables("gameworld").size, 0)
208+
assertEquals(testDb2.getTables("gameworld_sharded").size, 0)
218209
}
219210

220211
private fun assertTablesApplied() {
221-
val unshardedTables = testDb2QueryExecutor.getTables("gameworld").map { it.tableName }
212+
val unshardedTables = testDb2.getTables("gameworld").map { it.tableName }
222213
assertArrayEquals(arrayOf("customers_seq", "games_seq"), unshardedTables.toTypedArray())
223214

224-
val shardedTables = testDb2QueryExecutor.getTables("gameworld_sharded").map { it.tableName }
215+
val shardedTables = testDb2.getTables("gameworld_sharded").map { it.tableName }
225216
assertArrayEquals(arrayOf("customers", "games"), shardedTables.toTypedArray())
226217
}
227218

‎misk-vitess/src/test/kotlin/misk/vitess/testing/VitessTestDbTest.kt‎

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
package misk.vitess.testing
22

3-
import misk.vitess.testing.internal.VitessClusterConfig
4-
import misk.vitess.testing.internal.VitessQueryExecutor
53
import misk.vitess.testing.internal.VitessQueryExecutorException
64
import org.junit.jupiter.api.AfterAll
75
import org.junit.jupiter.api.Assertions.assertArrayEquals
@@ -15,14 +13,12 @@ class VitessTestDbTest {
1513
companion object {
1614
private lateinit var vitessTestDb: VitessTestDb
1715
private lateinit var vitessTestDbRunResult: VitessTestDbStartupResult
18-
private lateinit var vitessQueryExecutor: VitessQueryExecutor
1916

2017
@JvmStatic
2118
@BeforeAll
2219
fun setup() {
2320
vitessTestDb = VitessTestDb()
2421
vitessTestDbRunResult = vitessTestDb.run()
25-
vitessQueryExecutor = VitessQueryExecutor(VitessClusterConfig(DefaultSettings.PORT))
2622
}
2723

2824
@JvmStatic
@@ -36,7 +32,7 @@ class VitessTestDbTest {
3632

3733
@Test
3834
fun `test querying the database`() {
39-
val results = vitessQueryExecutor.executeQuery("SHOW KEYSPACES;")
35+
val results = vitessTestDb.executeQuery("SHOW KEYSPACES;")
4036
assertEquals(2, results.size)
4137
}
4238

@@ -49,7 +45,7 @@ class VitessTestDbTest {
4945
"""
5046
.trimIndent()
5147

52-
val insertCustomersResultCount = vitessQueryExecutor.executeUpdate(insertCustomersSql)
48+
val insertCustomersResultCount = vitessTestDb.executeUpdate(insertCustomersSql)
5349
assertEquals(1, insertCustomersResultCount)
5450

5551
val insertGamesSql =
@@ -59,7 +55,7 @@ class VitessTestDbTest {
5955
"""
6056
.trimIndent()
6157

62-
val insertGamesResultCount = vitessQueryExecutor.executeUpdate(insertGamesSql)
58+
val insertGamesResultCount = vitessTestDb.executeUpdate(insertGamesSql)
6359
assertEquals(1, insertGamesResultCount)
6460

6561
assertDoesNotThrow(vitessTestDb::truncate)
@@ -80,7 +76,7 @@ class VitessTestDbTest {
8076

8177
@Test
8278
fun `test each shard has expected tablet types`() {
83-
val results = vitessQueryExecutor.executeQuery("SHOW VITESS_REPLICATION_STATUS;")
79+
val results = vitessTestDb.executeQuery("SHOW VITESS_REPLICATION_STATUS;")
8480
val shardReplicaMap = mutableMapOf<String, Int>()
8581
val shardRdonlyMap = mutableMapOf<String, Int>()
8682

@@ -116,14 +112,13 @@ class VitessTestDbTest {
116112

117113
@Test
118114
fun `test schema changes are auto-applied`() {
119-
val vitessQueryExecutor = VitessQueryExecutor(VitessClusterConfig(27003))
120-
val keyspaces = vitessQueryExecutor.getKeyspaces()
115+
val keyspaces = vitessTestDb.getKeyspaces()
121116
assertArrayEquals(arrayOf("gameworld", "gameworld_sharded"), keyspaces.toTypedArray())
122117

123-
val unshardedTables = vitessQueryExecutor.getTables("gameworld").map { it.tableName }
118+
val unshardedTables = vitessTestDb.getTables("gameworld").map { it.tableName }
124119
assertArrayEquals(arrayOf("customers_seq", "games_seq"), unshardedTables.toTypedArray())
125120

126-
val shardedTables = vitessQueryExecutor.getTables("gameworld_sharded").map { it.tableName }
121+
val shardedTables = vitessTestDb.getTables("gameworld_sharded").map { it.tableName }
127122
assertArrayEquals(arrayOf("customers", "games"), shardedTables.toTypedArray())
128123
}
129124

@@ -138,29 +133,29 @@ class VitessTestDbTest {
138133

139134
@Test
140135
fun `test default timezone set to UTC`() {
141-
val results = vitessQueryExecutor.executeQuery("SELECT @@global.time_zone;")
136+
val results = vitessTestDb.executeQuery("SELECT @@global.time_zone;")
142137
val actualTimeZone = results[0]["@@global.time_zone"]
143138
assertEquals("+00:00", actualTimeZone)
144139
}
145140

146141
@Test
147142
fun `test default MySql version`() {
148143
// The version behavior was changed in v21+, using SELECT version() is more reliable.
149-
val results = vitessQueryExecutor.executeQuery("SELECT version();")
144+
val results = vitessTestDb.executeQuery("SELECT version();")
150145
val actualMysqlVersion = results[0]["version()"]
151146
assertEquals("${DefaultSettings.MYSQL_VERSION}-Vitess", actualMysqlVersion)
152147
}
153148

154149
@Test
155150
fun `test default SQL mode`() {
156-
val results = vitessQueryExecutor.executeQuery("SELECT @@global.sql_mode;")
151+
val results = vitessTestDb.executeQuery("SELECT @@global.sql_mode;")
157152
val actualSqlMode = results[0]["@@global.sql_mode"]
158153
assertEquals(DefaultSettings.SQL_MODE, actualSqlMode)
159154
}
160155

161156
@Test
162157
fun `test default transaction isolation level`() {
163-
val results = vitessQueryExecutor.executeQuery("SELECT @@global.transaction_ISOLATION;")
158+
val results = vitessTestDb.executeQuery("SELECT @@global.transaction_ISOLATION;")
164159
val actualTransactionIsolationLevel = results[0]["@@global.transaction_ISOLATION"]
165160
assertEquals(DefaultSettings.TRANSACTION_ISOLATION_LEVEL.value, actualTransactionIsolationLevel)
166161
}
@@ -169,23 +164,23 @@ class VitessTestDbTest {
169164
fun `explicit exception thrown on query errors`() {
170165
val executeQueryException =
171166
assertThrows<VitessQueryExecutorException> {
172-
vitessQueryExecutor.executeQuery("SELECT * FROM non_existent_table")
167+
vitessTestDb.executeQuery("SELECT * FROM non_existent_table")
173168
}
174169
assertEquals("Failed to run executeQuery on query: SELECT * FROM non_existent_table", executeQueryException.message)
175170

176171
val executeException =
177172
assertThrows<VitessQueryExecutorException> {
178-
vitessQueryExecutor.execute("UPDATE non_existent_table SET column = value WHERE id = 1")
173+
vitessTestDb.executeUpdate("UPDATE non_existent_table SET column = value WHERE id = 1")
179174
}
180175
assertEquals(
181-
"Failed to run execute on query: UPDATE non_existent_table SET column = value WHERE id = 1",
176+
"Failed to run executeUpdate on query: UPDATE non_existent_table SET column = value WHERE id = 1",
182177
executeException.message,
183178
)
184179
}
185180

186181
private fun getRowCount(table: String): Int {
187182
val query = "SELECT COUNT(*) FROM $table;"
188-
val results = vitessQueryExecutor.executeQuery(query)
183+
val results = vitessTestDb.executeQuery(query)
189184
return (results[0]["count(*)"] as Long).toInt()
190185
}
191186
}

‎misk-vitess/src/testFixtures/kotlin/misk/vitess/testing/VitessTestDb.kt‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,63 @@ class VitessTestDb(
244244
}
245245
}
246246

247+
/**
248+
* Get the list of keyspaces that are present on the vtgate.
249+
*
250+
* @return A list of keyspace names as strings.
251+
*/
252+
fun getKeyspaces(): List<String> {
253+
return VitessQueryExecutor(vitessClusterConfig).getKeyspaces()
254+
}
255+
256+
/**
257+
* Get the list of tables for a specific keyspace.
258+
*
259+
* @param keyspace The name of the keyspace to retrieve tables from.
260+
*
261+
* @return A list of [VitessTable] objects representing the tables in the specified keyspace.
262+
*/
263+
fun getTables(keyspace: String): List<VitessTable> {
264+
return VitessQueryExecutor(vitessClusterConfig).getTables(keyspace)
265+
}
266+
267+
/**
268+
* Execute a SELECT SQL query against the vtgate.
269+
*
270+
* @param query The SQL query to execute.
271+
* @param target The target for the query, which defaults to "@primary". This can be used to specify a specific shard.
272+
*
273+
* @return A list of maps representing the rows returned by the query, where each map corresponds to a row and each
274+
* key-value pair corresponds to a column name and its value.
275+
*/
276+
fun executeQuery(query: String, target: String = "@primary"): List<Map<String, Any>> {
277+
return VitessQueryExecutor(vitessClusterConfig).executeQuery(query, target)
278+
}
279+
280+
/**
281+
* Execute a SQL statement (i.e. INSERT, UPDATE, DELETE, DDL) against the vtgate.
282+
*
283+
* @param query The SQL statement to execute.
284+
* @param target The target for the query, which defaults to "@primary". This can be used to specify a specific shard.
285+
*
286+
* @return `true` if the execution was successful, `false` otherwise.
287+
*/
288+
fun executeUpdate(query: String, target: String = "@primary"): Int {
289+
return VitessQueryExecutor(vitessClusterConfig).executeUpdate(query, target)
290+
}
291+
292+
/**
293+
* Execute a transaction against the vtgate.
294+
*
295+
* @param query The SQL statement to execute as a transaction.
296+
* @param target The target for the query, which defaults to "@primary". This can be used to specify a specific shard.
297+
*
298+
* @return `true` if the transaction was successful, `false` otherwise.
299+
*/
300+
fun executeTransaction(query: String, target: String = "@primary"): Boolean {
301+
return VitessQueryExecutor(vitessClusterConfig).executeTransaction(query, target)
302+
}
303+
247304
private fun getVitessDockerContainer(): VitessDockerContainer {
248305
val container =
249306
VitessDockerContainer(
@@ -326,3 +383,12 @@ data class ApplySchemaResult(
326383
data class VSchemaUpdate(val vschema: String, val keyspace: String)
327384

328385
data class DdlUpdate(val ddl: String, val keyspace: String)
386+
387+
data class VitessTable(val tableName: String, val type: VitessTableType)
388+
389+
enum class VitessTableType {
390+
SEQUENCE,
391+
REFERENCE,
392+
STANDARD,
393+
}
394+

0 commit comments

Comments
 (0)