Merge pull request #42 from gfinger/master

New feature in OVertexTransformer: skip duplicate vertices
orientechnologies · Feb 20, 2015 · 43f58e2 · 43f58e2
2 parents 5edffe6 + fa37c4b
commit 43f58e2
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 2 deletions.
diff --git a/src/main/java/com/orientechnologies/orient/etl/transformer/OVertexTransformer.java b/src/main/java/com/orientechnologies/orient/etl/transformer/OVertexTransformer.java
@@ -22,6 +22,7 @@
 import com.orientechnologies.orient.core.exception.OSchemaException;
 import com.orientechnologies.orient.core.metadata.schema.OClass;
 import com.orientechnologies.orient.core.record.impl.ODocument;
+import com.orientechnologies.orient.core.storage.ORecordDuplicatedException;
 import com.orientechnologies.orient.etl.OETLProcessHaltedException;
 import com.orientechnologies.orient.etl.OETLProcessor;
 import com.tinkerpop.blueprints.impls.orient.OrientBaseGraph;
@@ -30,12 +31,14 @@
 public class OVertexTransformer extends OAbstractTransformer {
   protected String        vertexClass;
   private OrientBaseGraph graph;
+  private Boolean skipDuplicates;
 
   @Override
   public ODocument getConfiguration() {
     return new ODocument().fromJSON("{parameters:[" + getCommonConfigurationParameters() + ","
         + "{class:{optional:true,description:'Vertex class name to assign. Default is V'}}]"
-        + ",input:['OrientVertex','ODocument'],output:'OrientVertex'}");
+        + ",input:['OrientVertex','ODocument'],output:'OrientVertex'}"
+        + ",skipDuplicates:{optional:true,description:'Vertices with duplicate keys are skipped', default:false}");
   }
 
   @Override
@@ -44,6 +47,9 @@ public void configure(final OETLProcessor iProcessor, final ODocument iConfigura
 
     if (iConfiguration.containsField("class"))
       vertexClass = (String) resolve(iConfiguration.field("class"));
+    if (iConfiguration.containsField("skipDuplicates")) {
+      skipDuplicates = (Boolean) resolve(iConfiguration.field("skipDuplicates"));
+    }
   }
 
   @Override
@@ -74,7 +80,15 @@ public Object executeTransform(final Object input) {
       return null;
 
     if (vertexClass != null && !vertexClass.equals(v.getRecord().getClassName()))
-      v.setProperty("@class", vertexClass);
+      try {
+        v.setProperty("@class", vertexClass);
+      } catch (ORecordDuplicatedException e) {
+        if (skipDuplicates) {
+          return null;
+        } else {
+          throw e;
+        }
+      }
     return v;
   }
 }
diff --git a/src/test/java/com/orientechnologies/orient/etl/transformer/OVertexTransformerNGTest.java b/src/test/java/com/orientechnologies/orient/etl/transformer/OVertexTransformerNGTest.java
@@ -0,0 +1,89 @@
+/*
+ *
+ *  * Copyright 2010-2014 Orient Technologies LTD (info(at)orientechnologies.com)
+ *  *
+ *  * Licensed under the Apache License, Version 2.0 (the "License");
+ *  * you may not use this file except in compliance with the License.
+ *  * You may obtain a copy of the License at
+ *  *
+ *  *      http://www.apache.org/licenses/LICENSE-2.0
+ *  *
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS,
+ *  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  * See the License for the specific language governing permissions and
+ *  * limitations under the License.
+ *
+ */
+package com.orientechnologies.orient.etl.transformer;
+
+import com.orientechnologies.orient.etl.OETLProcessor;
+import com.orientechnologies.orient.etl.extractor.ETLBaseTest;
+import com.tinkerpop.blueprints.Parameter;
+import com.tinkerpop.blueprints.Vertex;
+import com.tinkerpop.blueprints.impls.orient.OrientGraph;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+/**
+ * Tests ETL Vertex Transformer.
+ *
+ * @author Gregor Frey
+ */
+public class OVertexTransformerNGTest extends ETLBaseTest {
+
+  OrientGraph graph;
+
+  @Test
+  public void testCreateVertex() {
+    OETLProcessor proc = getProcessor(
+            "{source: { content: { value: 'name,\nGregor' } }, extractor : { row: {} },"
+            + " transformers: [{csv: {}}, {vertex: {class:'Person', skipDuplicates:false}},"
+            + "], loader: { orientdb: { dbURL: 'memory:EdgeTransformerTest', dbType:'graph', useLightweightEdges:false } } }")
+            .execute();
+
+    assertEquals(graph.countVertices("Person"), 1);
+  }
+
+  @Test
+  public void testErrorOnDuplicateVertex() {
+    OETLProcessor proc = getProcessor(
+            "{source: { content: { value: 'name,\nGregor\nGregor\nHans' } }, extractor : { row: {} },"
+            + " transformers: [{csv: {}}, {vertex: {class:'Person', skipDuplicates:false}},"
+            + "], loader: { orientdb: { dbURL: 'memory:EdgeTransformerTest', dbType:'graph', useLightweightEdges:false } } }")
+            .execute();
+
+    assertEquals(graph.countVertices("Person"), 1);
+  }
+
+  @Test
+  public void testSkipDuplicateVertex() {
+    OETLProcessor proc = getProcessor(
+            "{source: { content: { value: 'name,\nGregor\nGregor\nHans' } }, extractor : { row: {} },"
+            + " transformers: [{csv: {}}, {vertex: {class:'Person', skipDuplicates:true}},"
+            + "], loader: { orientdb: { dbURL: 'memory:EdgeTransformerTest', dbType:'graph', useLightweightEdges:false } } }")
+            .execute();
+
+    assertEquals(graph.countVertices("Person"), 2);
+  }
+
+  @BeforeMethod
+  @Override
+  public void setUp() {
+    graph = new OrientGraph("memory:EdgeTransformerTest");
+    graph.setUseLightweightEdges(false);
+
+    graph.createVertexType("Person");
+    graph.createKeyIndex("name", Vertex.class,
+            new Parameter<String, String>("type", "UNIQUE"),
+            new Parameter<String, String>("class", "Person"));
+    graph.commit();
+  }
+
+  @Override
+  @AfterMethod
+  public void tearDown() {
+    graph.drop();
+  }
+}