Create Cluster TypeDBOptions when connecting to Cluster

Author: alexjpwalker

state/connection/ClientState.kt

@@ -82,6 +82,7 @@ class ClientState constructor(private val notificationMgr: NotificationManager)
     private var runningClosingCommands = AtomicIntegerState(0)
     private var databaseListRefreshedTime = System.currentTimeMillis()
     private val asyncDepth = AtomicInteger(0)
+    internal val isCluster get() = _client is TypeDBClient.Cluster
 
     private val coroutineScope = CoroutineScope(EmptyCoroutineContext)
 

state/connection/QueryRunner.kt

@@ -18,8 +18,6 @@
 
 package com.vaticle.typedb.studio.state.connection
 
-import com.vaticle.typedb.client.api.TypeDBOptions
-import com.vaticle.typedb.client.api.TypeDBTransaction
 import com.vaticle.typedb.client.api.answer.ConceptMap
 import com.vaticle.typedb.client.api.answer.ConceptMapGroup
 import com.vaticle.typedb.client.api.answer.NumericGroup
@@ -52,9 +50,8 @@ import kotlinx.coroutines.launch
 
 @OptIn(ExperimentalTime::class)
 class QueryRunner constructor(
-    val transaction: TypeDBTransaction, // TODO: restrict in the future, TypeDB 3.0 answers return complete info
+    private val transactionState: TransactionState, // TODO: restrict in the future, TypeDB 3.0 answers return complete info
     private val queries: String,
-    private val hasStopSignal: AtomicBoolean,
     private val onComplete: () -> Unit
 ) {
 
@@ -118,6 +115,8 @@ class QueryRunner constructor(
     private val lastResponse = AtomicLong(0)
     private val consumerLatch = CountDownLatch(1)
     private val coroutineScope = CoroutineScope(EmptyCoroutineContext)
+    private val hasStopSignal get() = transactionState.hasStopSignalAtomic.atomic
+    val transaction get() = transactionState.transaction
 
     fun launch() {
         isRunning.set(true)
@@ -218,7 +217,7 @@ class QueryRunner constructor(
             noResultMsg = INSERT_QUERY_NO_RESULT,
             queryStr = query.toString(),
             stream = Response.Stream.ConceptMaps()
-        ) { transaction.query().insert(query, TypeDBOptions.core().prefetch(true)) }
+        ) { transaction.query().insert(query, transactionState.typeDBOptions().prefetch(true)) }
     }
 
     private fun runUpdateQuery(query: TypeQLUpdate) {
@@ -228,7 +227,7 @@ class QueryRunner constructor(
             noResultMsg = UPDATE_QUERY_NO_RESULT,
             queryStr = query.toString(),
             stream = Response.Stream.ConceptMaps()
-        ) { transaction.query().update(query, TypeDBOptions.core().prefetch(true)) }
+        ) { transaction.query().update(query, transactionState.typeDBOptions().prefetch(true)) }
     }
 
     private fun runMatchQuery(query: TypeQLMatch) {

state/connection/SessionState.kt

@@ -36,7 +36,7 @@ import java.util.concurrent.atomic.AtomicReference
 import mu.KotlinLogging
 
 class SessionState constructor(
-    private val client: ClientState,
+    internal val client: ClientState,
     internal val notificationMgr: NotificationManager
 ) {
 

state/connection/TransactionState.kt

@@ -67,10 +67,11 @@ class TransactionState constructor(
     val hasStopSignal get() = hasStopSignalAtomic.state
     val hasRunningQuery get() = hasRunningQueryAtomic.state
     private val onSchemaWrite = LinkedBlockingQueue<() -> Unit>()
-    private val hasStopSignalAtomic = AtomicBooleanState(false)
+    internal val hasStopSignalAtomic = AtomicBooleanState(false)
     private var hasRunningQueryAtomic = AtomicBooleanState(false)
     private val isOpenAtomic = AtomicBooleanState(false)
     private var _transaction: TypeDBTransaction? by mutableStateOf(null)
+    internal val transaction get() = _transaction!!
 
     val snapshot = ConfigState(
         activatedFn = { it || type.isWrite },
@@ -94,7 +95,7 @@ class TransactionState constructor(
     private fun tryOpen() {
         if (isOpen) return
         try {
-            val options = TypeDBOptions.core().infer(infer.activated)
+            val options = typeDBOptions().infer(infer.activated)
                 .explain(infer.activated).transactionTimeoutMillis(ONE_HOUR_IN_MILLS)
             _transaction = session.transaction(type, options)!!.apply {
                 onClose { close(TRANSACTION_CLOSED_ON_SERVER, it?.message ?: "Unknown") }
@@ -112,7 +113,7 @@ class TransactionState constructor(
             try {
                 hasStopSignalAtomic.set(false)
                 tryOpen()
-                return if (isOpen) QueryRunner(_transaction!!, content, hasStopSignalAtomic.atomic) {
+                return if (isOpen) QueryRunner(this, content) {
                     if (!snapshot.activated) close()
                     else if (!isOpen) close(TRANSACTION_CLOSED_IN_QUERY)
                     hasStopSignalAtomic.set(false)
@@ -155,4 +156,8 @@ class TransactionState constructor(
             message?.let { notificationMgr.userError(LOGGER, it, *params) }
         }
     }
-}
+
+    internal fun typeDBOptions(): TypeDBOptions {
+        return if (session.client.isCluster) TypeDBOptions.cluster() else TypeDBOptions.core()
+    }
+}

state/schema/SchemaManager.kt

@@ -183,8 +183,7 @@ class SchemaManager(
         synchronized(this) {
             lastTransactionUse.set(System.currentTimeMillis())
             if (readTx.get() != null) return readTx.get()
-            val options = TypeDBOptions.core().transactionTimeoutMillis(ONE_HOUR_IN_MILLS)
-            readTx.set(session.transaction(options = options).also { it?.onClose { closeReadTx() } })
+            readTx.set(session.transaction().also { it?.onClose { closeReadTx() } })
             scheduleCloseReadTx()
             return readTx.get()
         }