Merge pull request #2797 from massooti/feat-timeseries

feat: Add Timeseries Collection Support

Author: bagerard

mongoengine/document.py

@@ -220,6 +220,8 @@ def _get_collection(cls):
             # Get the collection, either capped or regular.
             if cls._meta.get("max_size") or cls._meta.get("max_documents"):
                 cls._collection = cls._get_capped_collection()
+            elif cls._meta.get("timeseries"):
+                cls._collection = cls._get_timeseries_collection()
             else:
                 db = cls._get_db()
                 collection_name = cls._get_collection_name()
@@ -271,6 +273,27 @@ def _get_capped_collection(cls):
 
         return db.create_collection(collection_name, **opts)
 
+    @classmethod
+    def _get_timeseries_collection(cls):
+        """Create a new or get an existing timeseries PyMongo collection."""
+        db = cls._get_db()
+        collection_name = cls._get_collection_name()
+        timeseries_opts = cls._meta.get("timeseries")
+
+        if collection_name in list_collection_names(
+            db, include_system_collections=True
+        ):
+            collection = db[collection_name]
+            collection.options()
+            return collection
+
+        opts = {"expireAfterSeconds": timeseries_opts.pop("expireAfterSeconds", None)}
+        return db.create_collection(
+            name=collection_name,
+            timeseries=timeseries_opts,
+            **opts,
+        )
+
     def to_mongo(self, *args, **kwargs):
         data = super().to_mongo(*args, **kwargs)
 

tests/document/test_timeseries_collection.py

@@ -0,0 +1,202 @@
+import unittest
+from datetime import datetime, timedelta
+
+from mongoengine import (
+    DateTimeField,
+    Document,
+    FloatField,
+    StringField,
+    connect,
+    get_db,
+)
+from mongoengine.connection import disconnect
+from tests.utils import requires_mongodb_gte_50
+
+
+class TestTimeSeriesCollections(unittest.TestCase):
+    def setUp(self):
+        connect(db="mongoenginetest")
+        self.db = get_db()
+
+        class SensorData(Document):
+            timestamp = DateTimeField(required=True)
+            temperature = FloatField()
+
+            meta = {
+                "timeseries": {
+                    "timeField": "timestamp",
+                    "metaField": "temperature",
+                    "granularity": "seconds",
+                    "expireAfterSeconds": 5,
+                },
+                "collection": "sensor_data",
+            }
+
+        self.SensorData = SensorData
+
+    def test_get_db(self):
+        """Ensure that get_db returns the expected db."""
+        db = self.SensorData._get_db()
+        assert self.db == db
+
+    def tearDown(self):
+        for collection_name in self.db.list_collection_names():
+            if not collection_name.startswith("system."):
+                self.db.drop_collection(collection_name)
+        disconnect()
+
+    def test_definition(self):
+        """Ensure that document may be defined using fields."""
+        assert ["id", "temperature", "timestamp"] == sorted(
+            self.SensorData._fields.keys()
+        )
+        assert ["DateTimeField", "FloatField", "ObjectIdField"] == sorted(
+            x.__class__.__name__ for x in self.SensorData._fields.values()
+        )
+
+    @requires_mongodb_gte_50
+    def test_get_collection(self):
+        """Ensure that get_collection returns the expected collection."""
+        collection_name = "sensor_data"
+        collection = self.SensorData._get_collection()
+        assert self.db[collection_name] == collection
+
+    @requires_mongodb_gte_50
+    def test_create_timeseries_collection(self):
+        """Ensure that a time-series collection can be created."""
+        collection_name = self.SensorData._get_collection_name()
+        collection = self.SensorData._get_collection()
+
+        assert collection_name in self.db.list_collection_names()
+
+        options = collection.options()
+
+        assert options.get("timeseries") is not None
+        assert options["timeseries"]["timeField"] == "timestamp"
+        assert options["timeseries"]["granularity"] == "seconds"
+
+    @requires_mongodb_gte_50
+    def test_insert_document_into_timeseries_collection(self):
+        """Ensure that a document can be inserted into a time-series collection."""
+        collection_name = self.SensorData._get_collection_name()
+        collection = self.SensorData._get_collection()
+        assert collection_name in self.db.list_collection_names()
+
+        # Insert a document and ensure it was inserted
+        self.SensorData(timestamp=datetime.utcnow(), temperature=23.4).save()
+        assert collection.count_documents({}) == 1
+
+    @requires_mongodb_gte_50
+    def test_timeseries_expiration(self):
+        """Ensure that documents in a time-series collection expire after the specified time."""
+
+        self.SensorData._meta["timeseries"]["expireAfterSeconds"] = 1
+        self.SensorData._get_collection_name()
+        collection = self.SensorData._get_collection()
+        options = collection.options()
+        assert options.get("timeseries", {}) is not None
+        assert options["expireAfterSeconds"] == 1
+
+        self.SensorData(timestamp=datetime.utcnow(), temperature=23.4).save()
+
+        assert collection.count_documents({}) == 1
+
+        # Wait for more than the expiration time
+        import time
+
+        time.sleep(3)
+        assert collection.count_documents({}) > 0
+
+    @requires_mongodb_gte_50
+    def test_index_creation(self):
+        """Test if the index defined in the meta dictionary is created properly."""
+
+        # Define the Document with indexes
+        class SensorDataWithIndex(Document):
+            timestamp = DateTimeField(required=True)
+            temperature = FloatField()
+            location = StringField()  # Field to be indexed
+
+            meta = {
+                "timeseries": {
+                    "timeField": "timestamp",
+                    "metaField": "temperature",
+                    "granularity": "seconds",
+                    "expireAfterSeconds": 5,
+                },
+                "collection": "sensor_data",
+                "indexes": [
+                    {"fields": ["timestamp"], "name": "timestamp_index"},
+                    {"fields": ["temperature"], "name": "temperature_index"},
+                ],
+            }
+
+        collection = SensorDataWithIndex._get_collection()
+
+        indexes = collection.index_information()
+
+        assert (
+            "timestamp_index" in indexes
+        ), "Index on 'timestamp' field was not created"
+        assert (
+            "temperature_index" in indexes
+        ), "Index on 'temperature' field was not created"
+
+    @requires_mongodb_gte_50
+    def test_timeseries_data_insertion_order(self):
+        """Ensure that data in the time-series collection is inserted and queried in the correct time order."""
+        self.SensorData._get_collection_name()
+        self.SensorData._get_collection()
+
+        # Insert documents out of order
+        now = datetime.utcnow()
+        self.SensorData(timestamp=now, temperature=23.4).save()
+        self.SensorData(timestamp=now - timedelta(seconds=5), temperature=22.0).save()
+        self.SensorData(timestamp=now + timedelta(seconds=5), temperature=24.0).save()
+
+        documents = list(self.SensorData.objects.order_by("timestamp"))
+
+        # Check the insertion order
+        assert len(documents) == 3
+        assert documents[0].temperature == 22.0  # Earliest document
+        assert documents[1].temperature == 23.4  # Middle document
+        assert documents[2].temperature == 24.0  # Latest document
+
+    @requires_mongodb_gte_50
+    def test_timeseries_query_by_time_range(self):
+        """Ensure that data can be queried by a specific time range in the time-series collection."""
+
+        self.SensorData._get_collection_name()
+        self.SensorData._get_collection()
+
+        now = datetime.utcnow()
+        self.SensorData(timestamp=now - timedelta(seconds=10), temperature=22.0).save()
+        self.SensorData(timestamp=now - timedelta(seconds=5), temperature=23.0).save()
+        self.SensorData(timestamp=now, temperature=24.0).save()
+
+        # Query documents within the last 6 seconds
+        start_time = now - timedelta(seconds=6)
+        documents = self.SensorData.objects(timestamp__gte=start_time)
+
+        assert len(documents) == 2
+        assert documents[0].temperature == 23.0
+        assert documents[1].temperature == 24.0
+
+    @requires_mongodb_gte_50
+    def test_timeseries_large_data_volume(self):
+        """Ensure that the time-series collection can handle a large volume of data insertion."""
+
+        self.SensorData._get_collection_name()
+        collection = self.SensorData._get_collection()
+
+        for i in range(10000):
+            self.SensorData(
+                timestamp=datetime.utcnow() - timedelta(seconds=i),
+                temperature=20.0 + i % 5,
+            ).save()
+
+        assert collection.count_documents({}) == 10000
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/utils.py

@@ -45,6 +45,18 @@ def requires_mongodb_gte_44(func):
     return _decorated_with_ver_requirement(func, (4, 4), oper=operator.ge)
 
 
+def requires_mongodb_gte_50(func):
+    return _decorated_with_ver_requirement(func, (5, 0), oper=operator.ge)
+
+
+def requires_mongodb_gte_60(func):
+    return _decorated_with_ver_requirement(func, (6, 0), oper=operator.ge)
+
+
+def requires_mongodb_gte_70(func):
+    return _decorated_with_ver_requirement(func, (7, 0), oper=operator.ge)
+
+
 def _decorated_with_ver_requirement(func, mongo_version_req, oper):
     """Return a MongoDB version requirement decorator.