fix: make sure to return every time the correct document instance in …

…queries, issue #10280
orientechnologies · Aug 12, 2024 · cdc41e0 · cdc41e0
1 parent ff3ec17
commit cdc41e0
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 9 deletions.
diff --git a/core/src/main/java/com/orientechnologies/orient/core/sql/executor/CreateRecordStep.java b/core/src/main/java/com/orientechnologies/orient/core/sql/executor/CreateRecordStep.java
@@ -14,10 +14,13 @@ public class CreateRecordStep extends AbstractExecutionStep {
 
   private int created = 0;
   private int total = 0;
+  private Optional<String> cl;
 
-  public CreateRecordStep(OCommandContext ctx, int total, boolean profilingEnabled) {
+  public CreateRecordStep(
+      OCommandContext ctx, int total, boolean profilingEnabled, Optional<String> cl) {
     super(ctx, profilingEnabled);
     this.total = total;
+    this.cl = cl;
   }
 
   @Override
@@ -39,6 +42,9 @@ public OResult next() {
                 throw new IllegalStateException();
               }
               created++;
+              if (cl.isPresent()) {
+                return new OUpdatableResult((ODocument) ctx.getDatabase().newInstance(cl.get()));
+              }
               return new OUpdatableResult((ODocument) ctx.getDatabase().newInstance());
             } finally {
               if (profilingEnabled) {

diff --git a/...src/main/java/com/orientechnologies/orient/core/sql/executor/OInsertExecutionPlanner.java b/...src/main/java/com/orientechnologies/orient/core/sql/executor/OInsertExecutionPlanner.java
@@ -3,6 +3,7 @@
 import com.orientechnologies.orient.core.command.OCommandContext;
 import com.orientechnologies.orient.core.db.ODatabaseSession;
 import com.orientechnologies.orient.core.index.OIndexAbstract;
+import com.orientechnologies.orient.core.metadata.OMetadataInternal;
 import com.orientechnologies.orient.core.metadata.schema.OClass;
 import com.orientechnologies.orient.core.metadata.schema.OSchema;
 import com.orientechnologies.orient.core.sql.parser.OCluster;
@@ -18,6 +19,7 @@
 import com.orientechnologies.orient.core.sql.parser.OUpdateItem;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 /** Created by luigidellaquila on 08/08/16. */
 public class OInsertExecutionPlanner {
@@ -57,10 +59,10 @@ public OInsertExecutionPlan createExecutionPlan(OCommandContext ctx, boolean ena
     } else {
       if (selectStatement != null) {
         handleInsertSelect(result, this.selectStatement, ctx, enableProfiling);
+        handleTargetClass(result, ctx, enableProfiling);
       } else {
         handleCreateRecord(result, this.insertBody, ctx, enableProfiling);
       }
-      handleTargetClass(result, ctx, enableProfiling);
       handleSetFields(result, insertBody, ctx, enableProfiling);
       ODatabaseSession database = ctx.getDatabase();
       if (targetCluster != null) {
@@ -134,7 +136,14 @@ private void handleSetFields(
   private void handleTargetClass(
       OInsertExecutionPlan result, OCommandContext ctx, boolean profilingEnabled) {
     ODatabaseSession database = ctx.getDatabase();
-    OSchema schema = database.getMetadata().getSchema();
+    Optional<String> tc = resolveTargetClass(database);
+    if (tc.isPresent()) {
+      result.chain(new SetDocumentClassStep(tc.get(), ctx, profilingEnabled));
+    }
+  }
+
+  protected Optional<String> resolveTargetClass(ODatabaseSession database) {
+    OSchema schema = ((OMetadataInternal) database.getMetadata()).getImmutableSchemaSnapshot();
     OIdentifier tc = null;
     if (targetClass != null) {
       tc = targetClass;
@@ -157,7 +166,9 @@ private void handleTargetClass(
       }
     }
     if (tc != null) {
-      result.chain(new SetDocumentClassStep(tc, ctx, profilingEnabled));
+      return Optional.of(tc.getStringValue());
+    } else {
+      return Optional.empty();
     }
   }
 
@@ -185,7 +196,8 @@ private void handleCreateRecord(
         tot = body.getContent().size();
       }
     }
-    result.chain(new CreateRecordStep(ctx, tot, profilingEnabled));
+    Optional<String> cl = resolveTargetClass(ctx.getDatabase());
+    result.chain(new CreateRecordStep(ctx, tot, profilingEnabled, cl));
   }
 
   private void handleInsertSelect(

diff --git a/core/src/main/java/com/orientechnologies/orient/core/sql/executor/SetDocumentClassStep.java b/core/src/main/java/com/orientechnologies/orient/core/sql/executor/SetDocumentClassStep.java
@@ -5,7 +5,6 @@
 import com.orientechnologies.orient.core.db.record.OIdentifiable;
 import com.orientechnologies.orient.core.record.impl.ODocument;
 import com.orientechnologies.orient.core.sql.executor.resultset.OResultSetMapper;
-import com.orientechnologies.orient.core.sql.parser.OIdentifier;
 
 /**
  * Assigns a class to documents coming from upstream
@@ -15,10 +14,9 @@
 public class SetDocumentClassStep extends AbstractExecutionStep {
   private final String targetClass;
 
-  public SetDocumentClassStep(
-      OIdentifier targetClass, OCommandContext ctx, boolean profilingEnabled) {
+  public SetDocumentClassStep(String targetClass, OCommandContext ctx, boolean profilingEnabled) {
     super(ctx, profilingEnabled);
-    this.targetClass = targetClass.getStringValue();
+    this.targetClass = targetClass;
   }
 
   @Override

diff --git a/server/src/test/java/com/orientechnologies/orient/core/db/ODatabasePoolRemoteTest.java b/server/src/test/java/com/orientechnologies/orient/core/db/ODatabasePoolRemoteTest.java
@@ -1,12 +1,15 @@
 package com.orientechnologies.orient.core.db;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import com.orientechnologies.common.io.OFileUtils;
 import com.orientechnologies.orient.core.Orient;
 import com.orientechnologies.orient.core.config.OGlobalConfiguration;
 import com.orientechnologies.orient.core.db.document.ODatabaseDocument;
+import com.orientechnologies.orient.core.record.OVertex;
 import com.orientechnologies.orient.core.record.impl.ODocument;
+import com.orientechnologies.orient.core.sql.executor.OResultSet;
 import com.orientechnologies.orient.server.OServer;
 import java.io.File;
 import org.junit.After;
@@ -76,6 +79,59 @@ public void testPoolDoubleClose() {
     orientDb.close();
   }
 
+  @Test
+  public void testPoolCorrectVertex() {
+    OrientDB orientDb =
+        new OrientDB(
+            "remote:localhost:",
+            "root",
+            "root",
+            OrientDBConfig.builder().addConfig(OGlobalConfiguration.DB_POOL_MAX, 1).build());
+
+    if (!orientDb.exists("test")) {
+      orientDb.execute(
+          "create database test memory users (admin identified by 'admin' role admin)");
+    }
+
+    ODatabasePool pool = new ODatabasePool(orientDb, "test", "admin", "admin");
+    ODatabaseDocument db = pool.acquire();
+    db.command("create vertex v set name='abc'").close();
+    OResultSet res = db.query("select * from V ");
+    assertTrue(res.next().getElement().get() instanceof OVertex);
+    db.close();
+    pool.close();
+    orientDb.close();
+  }
+
+  @Test
+  public void testPoolCorrectVertexCloseOpen() {
+    OrientDB orientDb =
+        new OrientDB(
+            "remote:localhost:",
+            "root",
+            "root",
+            OrientDBConfig.builder().addConfig(OGlobalConfiguration.DB_POOL_MAX, 1).build());
+
+    if (!orientDb.exists("test")) {
+      orientDb.execute(
+          "create database test memory users (admin identified by 'admin' role admin)");
+    }
+
+    ODatabasePool pool = new ODatabasePool(orientDb, "test", "admin", "admin");
+    ODatabaseDocument db = pool.acquire();
+    db.command("create vertex v set name='abc'").close();
+    OResultSet res = db.query("select * from V ");
+    assertTrue(res.next().getElement().get() instanceof OVertex);
+    db.close();
+
+    db = pool.acquire();
+    res = db.query("select * from V ");
+    assertTrue(res.next().getElement().get() instanceof OVertex);
+
+    pool.close();
+    orientDb.close();
+  }
+
   @After
   public void after() {