From d807862512681e65342ee3aa4b95b779303cad68 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?D=D5=A1=C4=B1try?=
 <32328713+dwitry@users.noreply.github.com>
Date: Wed, 7 Nov 2018 17:01:44 +0200
Subject: [PATCH] Cosmos DB improvements II (#215)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Rewrite `skip` → `range`
- Convert ids to String

Signed-off-by: Dwitry dwitry@users.noreply.github.com
---
 .../ir/rewrite/CosmosDbFlavor.scala           | 21 +++++++++++++-
 .../ir/rewrite/CosmosDbFlavorTest.scala       | 29 ++++++++++++++++++-
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/translation/src/main/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavor.scala b/translation/src/main/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavor.scala
index 646c6bcb..a577bb69 100644
--- a/translation/src/main/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavor.scala
+++ b/translation/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/translation/src/test/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavorTest.scala b/translation/src/test/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavorTest.scala
index 832c8f6a..a3784957 100644
--- a/translation/src/test/scala/org/opencypher/gremlin/translation/ir/rewrite/CosmosDbFlavorTest.scala
+++ b/translation/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"))
+      )
+  }
 }