Cosmos DB improvements II (#215)

- Rewrite `skip` → `range` - Convert ids to String Signed-off-by: Dwitry [email protected]
opencypher · Nov 7, 2018 · d807862 · d807862
1 parent 73b3967
commit d807862
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 2 deletions.
diff --git a/...slation/src/main/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavor.scala b/...slation/src/main/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavor.scala
@@ -15,7 +15,8 @@
  */
 package org.opencypher.gremlin.translation.ir.rewrite
 
-import org.apache.tinkerpop.gremlin.process.traversal.Order
+import org.apache.tinkerpop.gremlin.process.traversal.{Order, Scope}
+import org.apache.tinkerpop.gremlin.structure.VertexProperty.Cardinality
 import org.opencypher.gremlin.translation.ir.TraversalHelper._
 import org.opencypher.gremlin.translation.ir.model._
 import org.opencypher.gremlin.translation.traversal.DeprecatedOrderAccessor.{decr, incr}
@@ -29,6 +30,8 @@ object CosmosDbFlavor extends GremlinRewriter {
       rewriteValues(_),
       rewriteRange(_),
       rewriteChoose(_),
+      rewriteSkip(_),
+      stringIds(_),
       tinkerPop334Workaround(_)
     ).foldLeft(steps) { (steps, rewriter) =>
       mapTraversals(rewriter)(steps)
@@ -79,4 +82,20 @@ object CosmosDbFlavor extends GremlinRewriter {
         ChooseP3(predicate, trueChoice, Identity :: Nil) :: rest
     })(steps)
   }
+
+  private def rewriteSkip(steps: Seq[GremlinStep]): Seq[GremlinStep] = {
+    replace({
+      case Skip(skip) :: rest =>
+        Range(Scope.global, skip, Int.MaxValue) :: rest
+    })(steps)
+  }
+
+  private def stringIds(steps: Seq[GremlinStep]): Seq[GremlinStep] = {
+    replace({
+      case PropertyVC(Cardinality.single, "id", value) :: rest =>
+        PropertyVC(Cardinality.single, "id", "" + value) :: rest
+      case PropertyTC(Cardinality.single, "id", Constant(value) :: Nil) :: rest =>
+        PropertyTC(Cardinality.single, "id", Constant("" + value) :: Nil) :: rest
+    })(steps)
+  }
 }
diff --git a/...ion/src/test/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavorTest.scala b/...ion/src/test/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavorTest.scala
@@ -15,7 +15,8 @@
  */
 package org.opencypher.gremlin.translation.ir.rewrite
 
-import org.apache.tinkerpop.gremlin.process.traversal.Order
+import org.apache.tinkerpop.gremlin.process.traversal.{Order, Scope}
+import org.apache.tinkerpop.gremlin.structure.VertexProperty.Cardinality
 import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.assertj.core.api.ThrowableAssert
 import org.junit.Test
@@ -144,4 +145,30 @@ class CosmosDbFlavorTest {
         __.select("value").choose(P.neq("  cypher.null"), __.id(), __.identity())
       )
   }
+
+  @Test
+  def skip(): Unit = {
+    assertThat(parse("MATCH (n) RETURN n SKIP 2"))
+      .withFlavor(flavor)
+      .rewritingWith(CosmosDbFlavor)
+      .removes(
+        __.skip(2)
+      )
+      .adds(
+        __.range(Scope.global, 2, Integer.MAX_VALUE)
+      )
+  }
+
+  @Test
+  def stringIds(): Unit = {
+    assertThat(parse("CREATE ({id: 1})"))
+      .withFlavor(flavor)
+      .rewritingWith(CosmosDbFlavor)
+      .removes(
+        __.property(Cardinality.single, "id", __.constant(1L))
+      )
+      .adds(
+        __.property(Cardinality.single, "id", __.constant("1"))
+      )
+  }
 }