Add timestamp_objectid function in MongoDB connector

Additionally, add description and documentation for objectid_timestamp function.
trinodb · Aug 10, 2021 · a46e7cd · a46e7cd
1 parent 2f3889e
commit a46e7cd
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 1 deletion.
diff --git a/docs/src/main/sphinx/connector/mongodb.rst b/docs/src/main/sphinx/connector/mongodb.rst
@@ -294,6 +294,41 @@ You can render the ``_id`` field to readable values with a cast to ``VARCHAR``:
      55b151633864d6438c61a9ce  |        1 | bad         |       50.0 | 2015-07-23
     (1 row)
 
+ObjectId timestamp functions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The first four bytes of each `ObjectId <https://docs.mongodb.com/manual/reference/method/ObjectId>`_ represent
+an embedded timestamp of its creation time. Trino provides a couple of functions to take advantage of this MongoDB feature.
+
+.. function:: objectid_timestamp(ObjectId) -> timestamp
+
+    Extracts the timestamp with time zone from a given ObjectId::
+
+        SELECT objectid_timestamp(ObjectId('507f191e810c19729de860ea'));
+        -- 2012-10-17 20:46:22.000 UTC
+
+.. function:: timestamp_objectid(timestamp) -> ObjectId
+
+    Creates an ObjectId from a timestamp with time zone::
+
+        SELECT timestamp_objectid(TIMESTAMP '2021-08-07 17:51:36 +00:00');
+        -- 61 0e c8 28 00 00 00 00 00 00 00 00
+
+In MongoDB, you can filter all the documents created after ``2021-08-07 17:51:36``
+with a query like this:
+
+.. code-block:: text
+
+    db.collection.find({"_id": {"$gt": ObjectId("610ec8280000000000000000")}})
+
+In Trino, the same can be achieved with this query:
+
+.. code-block:: sql
+
+    SELECT *
+    FROM collection
+    WHERE _id > timestamp_objectid(TIMESTAMP '2021-08-07 17:51:36 +00:00');
+
 Limitations
 -----------
 

diff --git a/docs/src/main/sphinx/functions/list.rst b/docs/src/main/sphinx/functions/list.rst
@@ -462,6 +462,7 @@ T
 - :func:`tan`
 - :func:`tanh`
 - :func:`tdigest_agg`
+- :func:`timestamp_objectid`
 - :func:`timezone_hour`
 - :func:`timezone_minute`
 - :func:`to_base`
@@ -546,4 +547,3 @@ Z
 
 - :func:`zip`
 - :func:`zip_with`
-
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ObjectIdFunctions.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ObjectIdFunctions.java
@@ -28,8 +28,10 @@
 import static io.airlift.slice.Slices.utf8Slice;
 import static io.trino.spi.function.OperatorType.CAST;
 import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone;
+import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc;
 import static io.trino.spi.type.TimeZoneKey.UTC_KEY;
 import static java.lang.Math.toIntExact;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 public final class ObjectIdFunctions
@@ -52,6 +54,7 @@ public static Slice objectid(@SqlType(StandardTypes.VARCHAR) Slice value)
         return Slices.wrappedBuffer(new ObjectId(CharMatcher.is(' ').removeFrom(value.toStringUtf8())).toByteArray());
     }
 
+    @Description("Timestamp from the given Mongodb ObjectId")
     @ScalarFunction
     @SqlType("timestamp(3) with time zone") // ObjectId's timestamp is a point in time
     public static long objectidTimestamp(@SqlType("ObjectId") Slice value)
@@ -60,6 +63,15 @@ public static long objectidTimestamp(@SqlType("ObjectId") Slice value)
         return packDateTimeWithZone(SECONDS.toMillis(epochSeconds), UTC_KEY);
     }
 
+    @Description("Mongodb ObjectId from the given timestamp")
+    @ScalarFunction
+    @SqlType("ObjectId")
+    public static Slice timestampObjectid(@SqlType("timestamp(0) with time zone") long timestamp)
+    {
+        long epochSeconds = MILLISECONDS.toSeconds(unpackMillisUtc(timestamp));
+        return Slices.wrappedBuffer(new ObjectId((int) epochSeconds, 0, (short) 0, 0).toByteArray());
+    }
+
     @ScalarOperator(CAST)
     @LiteralParameters("x")
     @SqlType("varchar(x)")

diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestObjectIdFunctions.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestObjectIdFunctions.java
@@ -63,6 +63,24 @@ public void testObjectidTimestamp()
                 toTimestampWithTimeZone(ZonedDateTime.of(1979, 9, 5, 22, 51, 36, 0, UTC)));
     }
 
+    @Test
+    public void testTimestampObjectid()
+    {
+        assertFunction(
+                "timestamp_objectid(TIMESTAMP '1979-09-05 22:51:36 +00:00')",
+                OBJECT_ID,
+                new SqlVarbinary(new ObjectId("123456780000000000000000").toByteArray()));
+    }
+
+    @Test
+    public void testTimestampObjectidNull()
+    {
+        assertFunction(
+                "timestamp_objectid(null)",
+                OBJECT_ID,
+                null);
+    }
+
     private SqlTimestampWithTimeZone toTimestampWithTimeZone(ZonedDateTime zonedDateTime)
     {
         return SqlTimestampWithTimeZone.newInstance(3, zonedDateTime.toInstant().toEpochMilli(), 0, TimeZoneKey.getTimeZoneKey(zonedDateTime.getZone().getId()));