Merge pull request #221 from JetBrains/graph-store-getType-restore

#XD-1142 fixed
JetBrains · Nov 20, 2024 · ec6b9d4 · ec6b9d4
2 parents 06d1a6c + 7d148ed
commit ec6b9d4
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 8 deletions.
diff --git a/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OPersistentEntityStore.kt b/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OPersistentEntityStore.kt
@@ -168,7 +168,8 @@ class OPersistentEntityStore(
     }
 
     override fun getEntityType(entityTypeId: Int): String {
-       TODO()
+        val currentTx = requireActiveTransaction()
+       return currentTx.getType(entityTypeId)
     }
 
     override fun renameEntityType(oldEntityTypeName: String, newEntityTypeName: String) {

diff --git a/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OSchemaBuddy.kt b/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OSchemaBuddy.kt
@@ -22,6 +22,7 @@ import com.orientechnologies.orient.core.metadata.sequence.OSequence.CreateParam
 import com.orientechnologies.orient.core.metadata.sequence.OSequence.SEQUENCE_TYPE
 import com.orientechnologies.orient.core.record.OVertex
 import com.orientechnologies.orient.core.sql.executor.OResultSet
+import jetbrains.exodus.entitystore.EntityRemovedInDatabaseException
 import jetbrains.exodus.entitystore.PersistentEntityId
 import jetbrains.exodus.entitystore.orientdb.OVertexEntity.Companion.CLASS_ID_CUSTOM_PROPERTY_NAME
 import jetbrains.exodus.entitystore.orientdb.OVertexEntity.Companion.CLASS_ID_SEQUENCE_NAME
@@ -39,6 +40,8 @@ interface OSchemaBuddy {
      */
     fun getTypeId(session: ODatabaseSession, entityType: String): Int
 
+    fun getType(session: ODatabaseSession, entityTypeId: Int): String
+
     fun requireTypeExists(session: ODatabaseSession, entityType: String)
 
     fun getOrCreateSequence(session: ODatabaseSession, sequenceName: String, initialValue: Long): OSequence
@@ -62,7 +65,7 @@ class OSchemaBuddyImpl(
         val INTERNAL_CLASS_NAMES = hashSetOf(OClass.VERTEX_CLASS_NAME)
     }
 
-    private val classIdToOClassId = ConcurrentHashMap<Int, Int>()
+    private val classIdToOClassId = ConcurrentHashMap<Int, Pair<Int, String>>()
 
     init {
         if (autoInitialize) {
@@ -76,7 +79,7 @@ class OSchemaBuddyImpl(
         session.createClassIdSequenceIfAbsent()
         for (oClass in session.metadata.schema.classes) {
             if (oClass.isVertexType && !INTERNAL_CLASS_NAMES.contains(oClass.name)) {
-                classIdToOClassId[oClass.requireClassId()] = oClass.defaultClusterId
+                classIdToOClassId[oClass.requireClassId()] = oClass.defaultClusterId to oClass.name
             }
         }
     }
@@ -130,7 +133,7 @@ class OSchemaBuddyImpl(
 
         val classId = entityId.typeId
         val localEntityId = entityId.localId
-        val oClassId = classIdToOClassId[classId] ?: return ORIDEntityId.EMPTY_ID
+        val oClassId = classIdToOClassId[classId]?.first ?: return ORIDEntityId.EMPTY_ID
         val schema = session.metadata.schema
         val oClass = schema.getClassByClusterId(oClassId) ?: return ORIDEntityId.EMPTY_ID
 
@@ -149,6 +152,17 @@ class OSchemaBuddyImpl(
         return session.getClass(entityType)?.requireClassId() ?: -1
     }
 
+    override fun getType(
+        session: ODatabaseSession,
+        entityTypeId: Int
+    ): String {
+        val (_, typeName) = classIdToOClassId.computeIfAbsent(entityTypeId) {
+            val oClass = session.schema.classes.find { oClass -> oClass.getCustom(CLASS_ID_CUSTOM_PROPERTY_NAME)?.toInt() == entityTypeId } ?: throw EntityRemovedInDatabaseException("Invalid type ID $entityTypeId")
+            oClass.requireClassId() to oClass.name
+        }
+        return typeName
+    }
+
     override fun requireTypeExists(session: ODatabaseSession, entityType: String) {
         val oClass = session.getClass(entityType)
         check(oClass != null) { "$entityType has not been found" }
@@ -213,4 +227,4 @@ internal fun ODatabaseSession.getOrCreateVertexClass(className: String): OClass
     if (existingClass != null) return existingClass
 
     return createVertexClassWithClassId(className)
-}
+}
diff --git a/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OStoreTransaction.kt b/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OStoreTransaction.kt
@@ -59,11 +59,18 @@ interface OStoreTransaction : StoreTransaction {
      */
     fun getTypeId(entityType: String): Int
 
+    /**
+     If the class has not been found, will throw EntityRemovedInDatabaseException with invalid type id
+     */
+    fun getType(entityTypeId: Int): String
+
     fun getOSequence(sequenceName: String): OSequence
 
     fun updateOSequence(sequenceName: String, currentValue: Long)
 
     fun renameOClass(oldName: String, newName: String)
 
     fun getOrCreateEdgeClass(linkName: String, outClassName: String, inClassName: String): OClass
-}
+
+    fun bindResultSet(resultSet: OResultSet)
+}
diff --git a/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OStoreTransactionImpl.kt b/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/OStoreTransactionImpl.kt
@@ -30,6 +30,7 @@ import jetbrains.exodus.entitystore.orientdb.iterate.OEntityOfTypeIterable
 import jetbrains.exodus.entitystore.orientdb.iterate.link.*
 import jetbrains.exodus.entitystore.orientdb.iterate.property.*
 import jetbrains.exodus.entitystore.orientdb.query.OQueryCancellingPolicy
+import java.sql.ResultSet
 
 internal typealias TransactionEventHandler = (ODatabaseSession, OStoreTransaction) -> Unit
 
@@ -56,6 +57,8 @@ class OStoreTransactionImpl(
         (session as ODatabaseSessionInternal).transaction.id.toLong()
     }
 
+    private val resultSets: MutableCollection<OResultSet> = arrayListOf()
+
     override fun getTransactionId(): Long {
         return transactionIdImpl
     }
@@ -195,6 +198,9 @@ class OStoreTransactionImpl(
 
     private fun cleanUpTxIfNeeded() {
         if (session.status == ODatabaseSession.STATUS.OPEN && session.activeTxCount() == 0) {
+            println("Closing ${resultSets.size} result sets")
+            resultSets.forEach(OResultSet::close)
+            resultSets.clear()
             onFinished(session, this)
         }
     }
@@ -490,4 +496,14 @@ class OStoreTransactionImpl(
         requireActiveTransaction()
         return schemaBuddy.getTypeId(session, entityType)
     }
+
+    override fun getType(entityTypeId: Int): String {
+        requireActiveTransaction()
+        return schemaBuddy.getType(session, entityTypeId)
+    }
+
+    override fun bindResultSet(resultSet: OResultSet) {
+        resultSets.add(resultSet)
+    }
+
 }
diff --git a/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/query/OQueryExecution.kt b/entity-store/src/main/kotlin/jetbrains/exodus/entitystore/orientdb/query/OQueryExecution.kt
@@ -31,7 +31,7 @@ object OQueryExecution : KLogging() {
             val executionPlan = resultSet.executionPlan.get().prettyPrint(10, 8)
             "Query: $sqlQuery, \n execution plan:\n  $executionPlan, \n stats: ${resultSet.queryStats}"
         }
-
+        tx.bindResultSet(resultSet)
         return resultSet
     }
 }
@@ -46,4 +46,4 @@ internal fun OStoreTransaction.buildSql(query: OQuery): SqlQuery {
     }
 
     return builder.build()
-}
+}
diff --git a/entity-store/src/test/kotlin/jetbrains/exodus/entitystore/orientdb/OSchemaBuddyTest.kt b/entity-store/src/test/kotlin/jetbrains/exodus/entitystore/orientdb/OSchemaBuddyTest.kt
@@ -147,4 +147,16 @@ class OSchemaBuddyTest: OTestMixin {
         }
     }
 
+    @Test
+    fun `require both classId and localEntityId to create an instance`() {
+        val oClass = orientDb.provider.withSession { oSession ->
+            oSession.createVertexClassWithClassId("type1")
+        }
+        val typeID = oClass.requireClassId()
+        orientDb.provider.withSession { oSession ->
+            assertEquals("type1", orientDb.schemaBuddy.getType(oSession, typeID))
+        }
+
+    }
+
 }