From d40a6bc00654ec214a050adc0aa0a828d2fd496a Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Fri, 25 Oct 2024 12:04:23 +0200
Subject: [PATCH 01/22] Database: Store IP addresses as packed binary.

Modified collection getters to include codec.
The "JSON friendly" is passed to outside API.
---
 dp3/api/internal/entity_response_models.py |  36 ++-
 dp3/api/routers/entity.py                  |   5 +-
 dp3/database/database.py                   | 288 ++++++++++++---------
 tests/test_api/test_get_entity_eid_data.py |   3 +-
 4 files changed, 190 insertions(+), 142 deletions(-)

diff --git a/dp3/api/internal/entity_response_models.py b/dp3/api/internal/entity_response_models.py
index 7845d364..5aac18b8 100644
--- a/dp3/api/internal/entity_response_models.py
+++ b/dp3/api/internal/entity_response_models.py
@@ -1,9 +1,10 @@
 from datetime import datetime
-from typing import Annotated, Any, Optional
+from typing import Annotated, Any, Optional, Union
 
-from pydantic import BaseModel, Field, NonNegativeInt, RootModel
+from pydantic import BaseModel, Field, NonNegativeInt, PlainSerializer
 
 from dp3.common.attrspec import AttrSpecType, AttrType
+from dp3.common.datapoint import to_json_friendly
 
 
 class EntityState(BaseModel):
@@ -19,6 +20,20 @@ class EntityState(BaseModel):
     eid_estimate_count: NonNegativeInt
 
 
+# This is necessary to allow for non-JSON-serializable types in the model
+JsonVal = Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")]
+
+PlainVal = dict[str, JsonVal]
+HistoryVal = list[dict[str, JsonVal]]
+
+Dp3Val = Union[HistoryVal, PlainVal, JsonVal]
+
+EntityEidMasterRecord = dict[str, Dp3Val]
+
+SnapshotType = dict[str, Dp3Val]
+EntityEidSnapshots = list[SnapshotType]
+
+
 class EntityEidList(BaseModel):
     """List of entity eids and their data based on latest snapshot
 
@@ -31,12 +46,7 @@ class EntityEidList(BaseModel):
     time_created: Optional[datetime] = None
     count: int
     total_count: int
-    data: list[dict]
-
-
-EntityEidMasterRecord = RootModel[dict]
-
-EntityEidSnapshots = RootModel[list[dict]]
+    data: EntityEidSnapshots
 
 
 class EntityEidData(BaseModel):
@@ -48,8 +58,8 @@ class EntityEidData(BaseModel):
     """
 
     empty: bool
-    master_record: dict
-    snapshots: list[dict]
+    master_record: EntityEidMasterRecord
+    snapshots: EntityEidSnapshots
 
 
 class EntityEidAttrValueOrHistory(BaseModel):
@@ -62,8 +72,8 @@ class EntityEidAttrValueOrHistory(BaseModel):
     """
 
     attr_type: AttrType
-    current_value: Any = None
-    history: list[dict] = []
+    current_value: JsonVal = None
+    history: HistoryVal = []
 
 
 class EntityEidAttrValue(BaseModel):
@@ -72,4 +82,4 @@ class EntityEidAttrValue(BaseModel):
     The value is fetched from master record.
     """
 
-    value: Any = None
+    value: JsonVal = None
diff --git a/dp3/api/routers/entity.py b/dp3/api/routers/entity.py
index 57cb759c..f8dc3654 100644
--- a/dp3/api/routers/entity.py
+++ b/dp3/api/routers/entity.py
@@ -12,6 +12,7 @@
     EntityEidList,
     EntityEidMasterRecord,
     EntityEidSnapshots,
+    JsonVal,
 )
 from dp3.api.internal.helpers import api_to_dp3_datapoint
 from dp3.api.internal.models import (
@@ -56,7 +57,7 @@ def get_eid_master_record_handler(
 
 def get_eid_snapshots_handler(
     etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
-):
+) -> list[dict[str, Any]]:
     """Handler for getting snapshots of EID"""
     snapshots = list(DB.get_snapshots(etype, eid, t1=date_from, t2=date_to))
 
@@ -248,7 +249,7 @@ async def set_eid_attr_value(
     "/{etype}/_/distinct/{attr}",
     responses={400: {"description": "Query can't be processed", "model": ErrorResponse}},
 )
-async def get_distinct_attribute_values(etype: str, attr: str) -> dict[Any, int]:
+async def get_distinct_attribute_values(etype: str, attr: str) -> dict[JsonVal, int]:
     """Gets distinct attribute values and their counts based on latest snapshots
 
     Useful for displaying `<select>` enumeration fields.
diff --git a/dp3/database/database.py b/dp3/database/database.py
index a5ee50f1..e0a6233f 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -5,15 +5,19 @@
 import time
 import urllib
 from collections import defaultdict
-from collections.abc import Iterable, Iterator
+from collections.abc import Generator, Iterable, Iterator
 from datetime import datetime, timedelta
+from ipaddress import IPv4Address, IPv6Address
 from typing import Any, Callable, Literal, Optional, Union
 
 import bson
 import pymongo
+from bson.binary import USER_DEFINED_SUBTYPE, Binary
+from bson.codec_options import TypeDecoder
 from event_count_logger import DummyEventGroup
 from pydantic import BaseModel, Field, field_validator
 from pymongo import ReplaceOne, UpdateMany, UpdateOne, WriteConcern
+from pymongo.collection import Collection
 from pymongo.command_cursor import CommandCursor
 from pymongo.cursor import Cursor
 from pymongo.errors import BulkWriteError, DocumentTooLarge, OperationFailure, WriteError
@@ -21,7 +25,7 @@
 
 from dp3.common.attrspec import AttrType, timeseries_types
 from dp3.common.config import HierarchicalDict, ModelSpec
-from dp3.common.datapoint import DataPointBase, to_json_friendly
+from dp3.common.datapoint import DataPointBase
 from dp3.common.scheduler import Scheduler
 from dp3.common.task import HASH
 from dp3.common.types import EventGroupType
@@ -29,6 +33,28 @@
 
 BSON_OBJECT_TOO_LARGE = 10334
 
+BSON_IPV4_SUBTYPE = USER_DEFINED_SUBTYPE + 1
+BSON_IPV6_SUBTYPE = BSON_IPV4_SUBTYPE + 1
+
+
+def fallback_encoder(value):
+    if isinstance(value, IPv4Address):
+        return Binary(value.packed, BSON_IPV4_SUBTYPE)
+    if isinstance(value, IPv6Address):
+        return Binary(value.packed, BSON_IPV6_SUBTYPE)
+    return value
+
+
+class DP3BinaryDecoder(TypeDecoder):
+    bson_type = Binary
+
+    def transform_bson(self, value):
+        if value.subtype == BSON_IPV4_SUBTYPE:
+            return IPv4Address(value)
+        if value.subtype == BSON_IPV6_SUBTYPE:
+            return IPv6Address(value)
+        return value
+
 
 class DatabaseError(Exception):
     pass
@@ -148,7 +174,9 @@ def __init__(
 
         # Init and switch to correct database
         self._db = self._db[db_config.db_name]
-        type_registry = bson.codec_options.TypeRegistry(fallback_encoder=to_json_friendly)
+        type_registry = bson.codec_options.TypeRegistry(
+            [DP3BinaryDecoder()], fallback_encoder=fallback_encoder
+        )
         self._codec_opts = bson.codec_options.CodecOptions(type_registry=type_registry)
 
         if process_index == 0:
@@ -260,26 +288,43 @@ def register_on_entity_delete(
         self._on_entity_delete_one.append(f_one)
         self._on_entity_delete_many.append(f_many)
 
-    @staticmethod
-    def _master_col_name(entity: str) -> str:
-        """Returns name of master collection for `entity`."""
-        return f"{entity}#master"
+    def _master_col(self, entity: str, **kwargs) -> Collection:
+        """Returns master collection for `entity`.
 
-    @staticmethod
-    def _snapshots_col_name(entity: str) -> str:
-        """Returns name of snapshots collection for `entity`."""
-        return f"{entity}#snapshots"
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        return self._db.get_collection(f"{entity}#master", codec_options=self._codec_opts, **kwargs)
 
-    @staticmethod
-    def _oversized_snapshots_col_name(entity: str) -> str:
-        """Returns name of oversized snapshots collection for `entity`."""
-        return f"{entity}#snapshots_oversized"
+    def _snapshot_col(self, entity: str, **kwargs) -> Collection:
+        """Returns snapshots collection for `entity`.
+
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        return self._db.get_collection(
+            f"{entity}#snapshots", codec_options=self._codec_opts, **kwargs
+        )
+
+    def _oversized_snapshots_col(self, entity: str, **kwargs) -> Collection:
+        """Returns oversized snapshots collection for `entity`.
+
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        return self._db.get_collection(
+            f"{entity}#snapshots_oversized", codec_options=self._codec_opts, **kwargs
+        )
 
     @staticmethod
     def _raw_col_name(entity: str) -> str:
         """Returns name of raw data collection for `entity`."""
         return f"{entity}#raw"
 
+    def _raw_col(self, entity: str, **kwargs) -> Collection:
+        """Returns raw data collection for `entity`.
+
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        return self._db.get_collection(f"{entity}#raw", codec_options=self._codec_opts, **kwargs)
+
     @staticmethod
     def _get_new_archive_col_name(entity: str) -> str:
         """Returns name of new archive collection for `entity`."""
@@ -289,6 +334,14 @@ def _archive_col_names(self, entity: str) -> list[str]:
         """Returns names of archive collections for `entity`."""
         return self._db.list_collection_names(filter={"name": {"$regex": f"{entity}#archive.*"}})
 
+    def _archive_cols(self, entity: str, **kwargs) -> Generator[Collection, None, None]:
+        """Returns archive collections for `entity`.
+
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        for col_name in self._archive_col_names(entity):
+            yield self._db.get_collection(col_name, codec_options=self._codec_opts, **kwargs)
+
     def _init_database_schema(self, db_name) -> None:
         """Runs full check and update of database schema.
 
@@ -300,44 +353,44 @@ def _init_database_schema(self, db_name) -> None:
         """
         for etype in self._db_schema_config.entities:
             # Snapshots index on `eid` and `_time_created` fields
-            snapshot_col = self._oversized_snapshots_col_name(etype)
-            self._db[snapshot_col].create_index("eid", background=True)
-            self._db[snapshot_col].create_index("_time_created", background=True)
+            snapshot_col = self._oversized_snapshots_col(etype)
+            snapshot_col.create_index("eid", background=True)
+            snapshot_col.create_index("_time_created", background=True)
 
-            snapshot_col = self._snapshots_col_name(etype)
+            snapshot_col = self._snapshot_col(etype)
             # To get the empty bucket for an entity without fetching all its buckets
-            for ix in self._db[snapshot_col].list_indexes():
+            for ix in snapshot_col.list_indexes():
                 if set(ix["key"].keys()) == {"_id", "count"}:
                     if ix.get("partialFilterExpression") is None:
-                        self._db[snapshot_col].drop_index(ix["name"])
-                        self._db[snapshot_col].create_index(
+                        snapshot_col.drop_index(ix["name"])
+                        snapshot_col.create_index(
                             [("_id", pymongo.DESCENDING), ("count", pymongo.ASCENDING)],
                             partialFilterExpression={"count": {"$lt": self._snapshot_bucket_size}},
                             background=True,
                         )
                     break
             else:
-                self._db[snapshot_col].create_index(
+                snapshot_col.create_index(
                     [("_id", pymongo.DESCENDING), ("count", pymongo.ASCENDING)],
                     partialFilterExpression={"count": {"$lt": self._snapshot_bucket_size}},
                     background=True,
                 )
 
             # To fetch all the latest snapshots without relying on date
-            self._db[snapshot_col].create_index(
+            snapshot_col.create_index(
                 [("latest", pymongo.DESCENDING), ("_id", pymongo.ASCENDING)],
                 partialFilterExpression={"latest": {"$eq": True}},
                 background=True,
             )
 
             # To fetch the oversized entities only
-            self._db[snapshot_col].create_index(
+            snapshot_col.create_index(
                 [("oversized", pymongo.DESCENDING)],
                 partialFilterExpression={"oversized": {"$eq": True}},
                 background=True,
             )
             # For deleting old snapshots
-            self._db[snapshot_col].create_index("_time_created", background=True)
+            snapshot_col.create_index("_time_created", background=True)
 
         # Create a TTL index for metadata collection
         self._db["#metadata"].create_index(
@@ -346,7 +399,7 @@ def _init_database_schema(self, db_name) -> None:
 
         # Master indexes
         for etype in self._db_schema_config.entities:
-            master_col = self._db[self._master_col_name(etype)]
+            master_col = self._master_col(etype)
             history_attrs = set()
             plain_attrs = set()
 
@@ -503,11 +556,9 @@ def _push_raw(self):
                 if not dps:
                     continue
             try:
-                self._db.get_collection(
-                    self._raw_col_name(etype),
-                    write_concern=WriteConcern(w=0),
-                    codec_options=self._codec_opts,
-                ).insert_many(dps, ordered=False)
+                self._raw_col(etype, write_concern=WriteConcern(w=0)).insert_many(
+                    dps, ordered=False
+                )
                 end = time.time()
                 self.log.debug(
                     "Inserted %s raw datapoints to %s in %.3fs, %.3fus lock, %.3fs insert",
@@ -523,11 +574,7 @@ def _push_raw(self):
     def _push_master(self):
         """Push master changes to database."""
         for etype, lock in self._master_buffer_locks.items():
-            master_col = self._db.get_collection(
-                self._master_col_name(etype),
-                write_concern=WriteConcern(w=1),
-                codec_options=self._codec_opts,
-            )
+            master_col = self._master_col(etype, write_concern=WriteConcern(w=1))
             begin = time.time()
             with lock:
                 locked = time.time()
@@ -574,9 +621,9 @@ def update_master_records(self, etype: str, eids: list[str], records: list[dict]
 
         Raises DatabaseError when update fails.
         """
-        master_col = self._master_col_name(etype)
+        master_col = self._master_col(etype)
         try:
-            res = self._db[master_col].bulk_write(
+            res = master_col.bulk_write(
                 [
                     ReplaceOne({"_id": eid}, record, upsert=True)
                     for eid, record in zip(eids, records)
@@ -591,9 +638,9 @@ def update_master_records(self, etype: str, eids: list[str], records: list[dict]
 
     def extend_ttl(self, etype: str, eid: str, ttl_tokens: dict[str, datetime]):
         """Extends TTL of given `etype`:`eid` by `ttl_tokens`."""
-        master_col = self._master_col_name(etype)
+        master_col = self._master_col(etype)
         try:
-            self._db[master_col].update_one(
+            master_col.update_one(
                 {"_id": eid},
                 {
                     "$max": {
@@ -608,9 +655,9 @@ def extend_ttl(self, etype: str, eid: str, ttl_tokens: dict[str, datetime]):
 
     def remove_expired_ttls(self, etype: str, expired_eid_ttls: dict[str, list[str]]):
         """Removes expired TTL of given `etype`:`eid`."""
-        master_col = self._master_col_name(etype)
+        master_col = self._master_col(etype)
         try:
-            res = self._db[master_col].bulk_write(
+            res = master_col.bulk_write(
                 [
                     UpdateOne(
                         {"_id": eid},
@@ -630,11 +677,11 @@ def remove_expired_ttls(self, etype: str, expired_eid_ttls: dict[str, list[str]]
 
     def delete_eids(self, etype: str, eids: list[str]):
         """Delete master record and all snapshots of `etype`:`eids`."""
-        master_col = self._master_col_name(etype)
-        snapshot_col = self._snapshots_col_name(etype)
-        os_snapshot_col = self._oversized_snapshots_col_name(etype)
+        master_col = self._master_col(etype)
+        snapshot_col = self._snapshot_col(etype)
+        os_snapshot_col = self._oversized_snapshots_col(etype)
         try:
-            res = self._db[master_col].delete_many({"_id": {"$in": eids}})
+            res = master_col.delete_many({"_id": {"$in": eids}})
             self.log.debug(
                 "Deleted %s master records of %s (%s).", res.deleted_count, etype, len(eids)
             )
@@ -642,10 +689,10 @@ def delete_eids(self, etype: str, eids: list[str]):
         except Exception as e:
             raise DatabaseError(f"Delete of master record failed: {e}\n{eids}") from e
         try:
-            res = self._db[snapshot_col].delete_many(self._snapshot_bucket_eids_filter(eids))
+            res = snapshot_col.delete_many(self._snapshot_bucket_eids_filter(eids))
             del_cnt = res.deleted_count * self._snapshot_bucket_size
             self.log.debug("Deleted %s snapshots of %s (%s).", del_cnt, etype, len(eids))
-            res = self._db[os_snapshot_col].delete_many({"eid": {"$in": eids}})
+            res = os_snapshot_col.delete_many({"eid": {"$in": eids}})
             self.log.debug(
                 "Deleted %s oversized snapshots of %s (%s).", res.deleted_count, etype, len(eids)
             )
@@ -659,20 +706,20 @@ def delete_eids(self, etype: str, eids: list[str]):
 
     def delete_eid(self, etype: str, eid: str):
         """Delete master record and all snapshots of `etype`:`eid`."""
-        master_col = self._master_col_name(etype)
-        snapshot_col = self._snapshots_col_name(etype)
-        os_snapshot_col = self._oversized_snapshots_col_name(etype)
+        master_col = self._master_col(etype)
+        snapshot_col = self._snapshot_col(etype)
+        os_snapshot_col = self._oversized_snapshots_col(etype)
         try:
-            self._db[master_col].delete_one({"_id": eid})
+            master_col.delete_one({"_id": eid})
             self.log.debug("Deleted master record of %s/%s.", etype, eid)
             self.elog.log("record_removed")
         except Exception as e:
             raise DatabaseError(f"Delete of master record failed: {e}\n{eid}") from e
         try:
-            res = self._db[snapshot_col].delete_many(self._snapshot_bucket_eid_filter(eid))
+            res = snapshot_col.delete_many(self._snapshot_bucket_eid_filter(eid))
             del_cnt = res.deleted_count * self._snapshot_bucket_size
             self.log.debug("deleted %s snapshots of %s/%s.", del_cnt, etype, eid)
-            res = self._db[os_snapshot_col].delete_many({"eid": eid})
+            res = os_snapshot_col.delete_many({"eid": eid})
             self.log.debug(
                 "Deleted %s oversized snapshots of %s/%s.", res.deleted_count, etype, eid
             )
@@ -690,9 +737,9 @@ def mark_all_entity_dps_t2(self, etype: str, attrs: list[str]) -> UpdateResult:
 
         Periodically called for all `etype`s from HistoryManager.
         """
-        master_col = self._master_col_name(etype)
+        master_col = self._master_col(etype)
         try:
-            return self._db[master_col].update_many(
+            return master_col.update_many(
                 {},
                 [
                     {
@@ -737,9 +784,7 @@ def delete_old_dps(self, etype: str, attr_name: str, t_old: datetime) -> UpdateR
 
         Periodically called for all `etype`s from HistoryManager.
         """
-        master_col = self._db.get_collection(
-            self._master_col_name(etype), write_concern=WriteConcern(w=1)
-        )
+        master_col = self._master_col(etype, write_concern=WriteConcern(w=1))
         try:
             return master_col.update_many(
                 {f"#min_t2s.{attr_name}": {"$lt": t_old}},
@@ -782,16 +827,16 @@ def delete_link_dps(
 
         Called from LinkManager for deleted entities.
         """
-        master_col = self._master_col_name(etype)
+        master_col = self._master_col(etype)
         attr_type = self._db_schema_config.attr(etype, attr_name).t
         filter_cond = {"_id": {"$in": affected_eids}}
         try:
             if attr_type == AttrType.OBSERVATIONS:
                 update_pull = {"$pull": {attr_name: {"v.eid": eid_to}}}
-                self._db[master_col].update_many(filter_cond, update_pull)
+                master_col.update_many(filter_cond, update_pull)
             elif attr_type == AttrType.PLAIN:
                 update_unset = {"$unset": {attr_name: ""}}
-                self._db[master_col].update_many(filter_cond, update_unset)
+                master_col.update_many(filter_cond, update_unset)
             else:
                 raise ValueError(f"Unsupported attribute type: {attr_type}")
         except Exception as e:
@@ -813,7 +858,7 @@ def delete_many_link_dps(
             for etype, affected_eid_list, attr_name, eid_to_list in zip(
                 etypes, affected_eids, attr_names, eids_to
             ):
-                master_col = self._master_col_name(etype)
+                master_col = self._master_col(etype)
                 attr_type = self._db_schema_config.attr(etype, attr_name).t
                 filter_cond = {"_id": {"$in": affected_eid_list}}
                 if attr_type == AttrType.OBSERVATIONS:
@@ -824,7 +869,7 @@ def delete_many_link_dps(
                     updates.append(UpdateMany(filter_cond, update_unset))
                 else:
                     raise ValueError(f"Unsupported attribute type: {attr_type}")
-                self._db[master_col].bulk_write(updates)
+                master_col.bulk_write(updates)
         except Exception as e:
             raise DatabaseError(f"Delete of link datapoints failed: {e}") from e
 
@@ -836,8 +881,8 @@ def get_master_record(self, etype: str, eid: str, **kwargs) -> dict:
         # Check `etype`
         self._assert_etype_exists(etype)
 
-        master_col = self._master_col_name(etype)
-        return self._db[master_col].find_one({"_id": eid}, **kwargs) or {}
+        master_col = self._master_col(etype)
+        return master_col.find_one({"_id": eid}, **kwargs) or {}
 
     def ekey_exists(self, etype: str, eid: str) -> bool:
         """Checks whether master record for etype/eid exists"""
@@ -848,8 +893,8 @@ def get_master_records(self, etype: str, **kwargs) -> Cursor:
         # Check `etype`
         self._assert_etype_exists(etype)
 
-        master_col = self._master_col_name(etype)
-        return self._db[master_col].find({}, **kwargs)
+        master_col = self._master_col(etype)
+        return master_col.find({}, **kwargs)
 
     def get_worker_master_records(
         self, worker_index: int, worker_cnt: int, etype: str, query_filter: dict = None, **kwargs
@@ -859,8 +904,8 @@ def get_worker_master_records(
             raise DatabaseError(f"Entity '{etype}' does not exist")
 
         query_filter = {} if query_filter is None else query_filter
-        master_col = self._master_col_name(etype)
-        return self._db[master_col].find(
+        master_col = self._master_col(etype)
+        return master_col.find(
             {"#hash": {"$mod": [worker_cnt, worker_index]}, **query_filter}, **kwargs
         )
 
@@ -872,9 +917,9 @@ def get_latest_snapshot(self, etype: str, eid: str) -> dict:
         # Check `etype`
         self._assert_etype_exists(etype)
 
-        snapshot_col = self._snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
         return (
-            self._db[snapshot_col].find_one(
+            snapshot_col.find_one(
                 self._snapshot_bucket_eid_filter(eid), {"last": 1}, sort=[("_id", -1)]
             )
             or {}
@@ -938,7 +983,7 @@ def get_latest_snapshots(
         # Check `etype`
         self._assert_etype_exists(etype)
 
-        snapshot_col = self._snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
 
         if not fulltext_filters:
             fulltext_filters = {}
@@ -972,9 +1017,9 @@ def get_latest_snapshots(
                 query["last." + attr] = fulltext_filter
 
         try:
-            return self._db[snapshot_col].find(query, {"last": 1}).sort(
+            return snapshot_col.find(query, {"last": 1}).sort(
                 [("_id", pymongo.ASCENDING)]
-            ), self._db[snapshot_col].count_documents(query)
+            ), snapshot_col.count_documents(query)
         except OperationFailure as e:
             raise DatabaseError("Invalid query") from e
 
@@ -1006,12 +1051,11 @@ def get_snapshots(
         # Check `etype`
         self._assert_etype_exists(etype)
 
-        snapshot_col = self._snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
 
         # Find out if the snapshot is oversized
         doc = (
-            self._db[snapshot_col]
-            .find(self._snapshot_bucket_eid_filter(eid), {"oversized": 1})
+            snapshot_col.find(self._snapshot_bucket_eid_filter(eid), {"oversized": 1})
             .sort([("_id", -1)])
             .limit(1)
         )
@@ -1036,13 +1080,13 @@ def get_snapshots(
         if query["_time_created"]:
             pipeline.append({"$match": query})
         pipeline.append({"$sort": {"_time_created": pymongo.ASCENDING}})
-        return self._db[snapshot_col].aggregate(pipeline)
+        return snapshot_col.aggregate(pipeline)
 
     def _get_snapshots_oversized(
         self, etype: str, eid: str, t1: Optional[datetime] = None, t2: Optional[datetime] = None
     ) -> Cursor:
         """Get all (or filtered) snapshots of given `eid` from oversized snapshots collection."""
-        snapshot_col = self._oversized_snapshots_col_name(etype)
+        snapshot_col = self._oversized_snapshots_col(etype)
         query = {"eid": eid, "_time_created": {}}
 
         # Filter by date
@@ -1055,10 +1099,8 @@ def _get_snapshots_oversized(
         if not query["_time_created"]:
             del query["_time_created"]
 
-        return (
-            self._db[snapshot_col]
-            .find(query, projection={"_id": False})
-            .sort([("_time_created", pymongo.ASCENDING)])
+        return snapshot_col.find(query, projection={"_id": False}).sort(
+            [("_time_created", pymongo.ASCENDING)]
         )
 
     def get_value_or_history(
@@ -1106,34 +1148,30 @@ def estimate_count_eids(self, etype: str) -> int:
         # Check `etype`
         self._assert_etype_exists(etype)
 
-        master_col = self._master_col_name(etype)
-        return self._db[master_col].estimated_document_count({})
+        master_col = self._master_col(etype)
+        return master_col.estimated_document_count({})
 
     def _migrate_to_oversized_snapshot(self, etype: str, eid: str, snapshot: dict):
-        snapshot_col = self._snapshots_col_name(etype)
-        os_col = self._oversized_snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
+        os_col = self._oversized_snapshots_col(etype)
 
         try:
             move_to_oversized = []
             last_id = None
-            for doc in (
-                self._db[snapshot_col].find(self._snapshot_bucket_eid_filter(eid)).sort({"_id": 1})
-            ):
+            for doc in snapshot_col.find(self._snapshot_bucket_eid_filter(eid)).sort({"_id": 1}):
                 move_to_oversized.extend(doc.get("history", []))
                 last_id = doc["_id"]
             move_to_oversized.insert(0, snapshot)
 
             self._db[os_col].insert_many(move_to_oversized)
-            self._db[snapshot_col].update_one(
+            snapshot_col.update_one(
                 {"_id": last_id},
                 {
                     "$set": {"oversized": True, "last": snapshot, "count": 0},
                     "$unset": {"history": ""},
                 },
             )
-            self._db[snapshot_col].delete_many(
-                self._snapshot_bucket_eid_filter(eid) | {"oversized": False}
-            )
+            snapshot_col.delete_many(self._snapshot_bucket_eid_filter(eid) | {"oversized": False})
         except Exception as e:
             raise DatabaseError(f"Update of snapshot {eid} failed: {e}, {snapshot}") from e
 
@@ -1151,14 +1189,14 @@ def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
         eid = snapshot["eid"]
         snapshot["_time_created"] = ctime
 
-        snapshot_col = self._snapshots_col_name(etype)
-        os_col = self._oversized_snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
+        os_col = self._oversized_snapshots_col(etype)
 
         # Find out if the snapshot is oversized
         normal, oversized = self._get_snapshot_state(etype, {eid})
         if normal:
             try:
-                res = self._db[snapshot_col].update_one(
+                res = snapshot_col.update_one(
                     self._snapshot_bucket_eid_filter(eid)
                     | {"count": {"$lt": self._snapshot_bucket_size}},
                     {
@@ -1176,7 +1214,7 @@ def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
                 )
 
                 if res.upserted_id is not None:
-                    self._db[snapshot_col].update_many(
+                    snapshot_col.update_many(
                         self._snapshot_bucket_eid_filter(eid)
                         | {"latest": True, "count": self._snapshot_bucket_size},
                         {"$unset": {"latest": 1}},
@@ -1193,10 +1231,10 @@ def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
             return
         elif oversized:
             # Snapshot is already marked as oversized
-            self._db[snapshot_col].update_one(
+            snapshot_col.update_one(
                 self._snapshot_bucket_eid_filter(eid), {"$set": {"last": snapshot}}
             )
-            self._db[os_col].insert_one(snapshot)
+            os_col.insert_one(snapshot)
             return
 
     def _get_snapshot_state(self, etype: str, eids: set[str]) -> tuple[set, set]:
@@ -1209,9 +1247,9 @@ def _get_snapshot_state(self, etype: str, eids: set[str]) -> tuple[set, set]:
         if not unknown:
             return normal, oversized
 
-        snapshot_col = self._snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
         new_oversized = set()
-        for doc in self._db[snapshot_col].find(
+        for doc in snapshot_col.find(
             self._snapshot_bucket_eids_filter(unknown) | {"oversized": True},
             {"oversized": 1},
         ):
@@ -1246,8 +1284,8 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
         for snapshot in snapshots:
             snapshot["_time_created"] = ctime
 
-        snapshot_col = self._snapshots_col_name(etype)
-        os_col = self._oversized_snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
+        os_col = self._oversized_snapshots_col(etype)
 
         snapshots_by_eid = defaultdict(list)
         for snapshot in snapshots:
@@ -1301,7 +1339,7 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
 
         if upserts:
             try:
-                res = self._db[snapshot_col].bulk_write(upserts, ordered=False)
+                res = snapshot_col.bulk_write(upserts, ordered=False)
 
                 # Unset latest snapshots if new snapshots were inserted
                 if res.upserted_count > 0:
@@ -1315,7 +1353,7 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
                                 {"$unset": {"latest": 1}},
                             )
                         )
-                    up_res = self._db[snapshot_col].bulk_write(unset_latest_updates)
+                    up_res = snapshot_col.bulk_write(unset_latest_updates)
                     if up_res.modified_count != res.upserted_count:
                         self.log.info(
                             "Upserted the first snapshot for %d entities.",
@@ -1355,8 +1393,8 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
         if oversized_inserts:
             try:
                 if oversized_updates:
-                    self._db[snapshot_col].bulk_write(oversized_updates)
-                self._db[os_col].insert_many(oversized_inserts)
+                    snapshot_col.bulk_write(oversized_updates)
+                os_col.insert_many(oversized_inserts)
             except Exception as e:
                 raise DatabaseError(f"Insert of snapshots failed: {str(e)[:2048]}") from e
 
@@ -1559,7 +1597,7 @@ def get_distinct_val_count(self, etype: str, attr: str) -> dict[Any, int]:
         """
         self._assert_etype_exists(etype)
 
-        snapshot_col = self._snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
 
         # Get attribute specification
         try:
@@ -1599,7 +1637,7 @@ def get_distinct_val_count(self, etype: str, attr: str) -> dict[Any, int]:
             {"$sort": {"_id": 1, "count": -1}},
         ]
         # Run aggregation
-        distinct_counts_cur = self._db[snapshot_col].aggregate(agg_query)
+        distinct_counts_cur = snapshot_col.aggregate(agg_query)
 
         distinct_counts = {x["_id"]: x["count"] for x in distinct_counts_cur}
 
@@ -1626,8 +1664,8 @@ def move_raw_to_archive(self, etype: str):
 
     def get_archive_summary(self, etype: str, before: datetime) -> Optional[dict]:
         collection_summaries = []
-        for archive_col in self._archive_col_names(etype):
-            result_cursor = self._get_archive_summary(archive_col, before=before)
+        for archive_col_name in self._archive_col_names(etype):
+            result_cursor = self._get_archive_summary(archive_col_name, before=before)
             for summary in result_cursor:
                 collection_summaries.append(summary)
         if not collection_summaries:
@@ -1661,9 +1699,9 @@ def get_archive(self, etype: str, after: datetime, before: datetime) -> Iterator
         All plain datapoints will be returned (default).
         """
         query_filter = {"$or": [{"t1": {"$gte": after, "$lt": before}}, {"t1": None}]}
-        for archive_col in self._archive_col_names(etype):
+        for archive_col in self._archive_cols(etype):
             try:
-                yield self._db[archive_col].find(query_filter)
+                yield archive_col.find(query_filter)
             except Exception as e:
                 raise DatabaseError(f"Archive collection {archive_col} fetch failed: {e}") from e
 
@@ -1673,18 +1711,17 @@ def delete_old_archived_dps(self, etype: str, before: datetime) -> Iterator[Dele
         Deletes all plain datapoints if `plain` is `True` (default).
         """
         query_filter = {"$or": [{"t1": {"$lt": before}}, {"t1": None}]}
-        for archive_col in self._archive_col_names(etype):
+        for archive_col in self._archive_cols(etype):
             try:
-                yield self._db[archive_col].delete_many(query_filter)
+                yield archive_col.delete_many(query_filter)
             except Exception as e:
                 raise DatabaseError(f"Delete of old datapoints failed: {e}") from e
 
     def drop_empty_archives(self, etype: str) -> int:
         """Drop empty archive collections."""
         dropped_count = 0
-        for archive_col in self._archive_col_names(etype):
+        for col in self._archive_cols(etype):
             try:
-                col = self._db[archive_col]
                 if col.estimated_document_count() < 1000 and col.count_documents({}) == 0:
                     col.drop()
                     dropped_count += 1
@@ -1697,15 +1734,15 @@ def delete_old_snapshots(self, etype: str, t_old: datetime) -> int:
 
         Periodically called for all `etype`s from HistoryManager.
         """
-        snapshot_col_name = self._snapshots_col_name(etype)
-        os_snapshot_col_name = self._oversized_snapshots_col_name(etype)
+        snapshot_col = self._snapshot_col(etype)
+        os_snapshot_col = self._oversized_snapshots_col(etype)
         deleted = 0
         try:
-            res = self._db[snapshot_col_name].delete_many(
+            res = snapshot_col.delete_many(
                 {"_time_created": {"$lte": t_old - self._bucket_delta}},
             )
             deleted += res.deleted_count * self._snapshot_bucket_size
-            res = self._db[os_snapshot_col_name].delete_many({"_time_created": {"$lt": t_old}})
+            res = os_snapshot_col.delete_many({"_time_created": {"$lt": t_old}})
             deleted += res.deleted_count
         except Exception as e:
             raise DatabaseError(f"Delete of olds snapshots failed: {e}") from e
@@ -1722,5 +1759,4 @@ def get_module_cache(self, override_called_id: Optional[str] = None):
 
     def get_estimated_entity_count(self, entity_type: str) -> int:
         """Get count of entities of given type."""
-        master_col = self._master_col_name(entity_type)
-        return self._db[master_col].estimated_document_count()
+        return self._master_col(entity_type).estimated_document_count()
diff --git a/tests/test_api/test_get_entity_eid_data.py b/tests/test_api/test_get_entity_eid_data.py
index 378597e9..024ea9b4 100644
--- a/tests/test_api/test_get_entity_eid_data.py
+++ b/tests/test_api/test_get_entity_eid_data.py
@@ -2,6 +2,7 @@
 from datetime import datetime, timedelta
 
 import common
+from pydantic import RootModel
 
 from dp3.api.internal.entity_response_models import EntityEidData, EntityEidMasterRecord
 
@@ -48,7 +49,7 @@ def test_get_entity_data(self):
     def test_get_entity_master_record(self):
         data = self.query_expected_value(
             lambda: self.get_entity_data(
-                f"entity/test_entity_type/{self.eid}/master", EntityEidMasterRecord
+                f"entity/test_entity_type/{self.eid}/master", RootModel[EntityEidMasterRecord]
             ),
             lambda data: "test_attr_history" in data.root,
         )

From 5c25b90cccbe603a1f803fa8eddda35c98d2ea1e Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 5 Nov 2024 11:05:28 +0100
Subject: [PATCH 02/22] Config: Can specify EID data type.

---
 dp3/common/datatype.py   | 27 +++++++++++++++++++++++++--
 dp3/common/entityspec.py |  2 ++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/dp3/common/datatype.py b/dp3/common/datatype.py
index 0370f571..9fb36f5f 100644
--- a/dp3/common/datatype.py
+++ b/dp3/common/datatype.py
@@ -40,6 +40,8 @@
     "json": Union[Json[Any], dict, list],
 }
 
+eid_data_types = ["string", "int", "ipv4", "ipv6"]
+
 
 class ReadOnly(BaseModel):
     """The ReadOnly data_type is used to avoid datapoint insertion for an attribute."""
@@ -95,8 +97,7 @@ class DataType(RootModel):
     _mirror_as = PrivateAttr()
     _link_data = PrivateAttr()
 
-    @model_validator(mode="after")
-    def determine_value_validator(self):
+    def _determine_value_validator(self):
         """Determines value validator (inner `data_type`)."""
         str_type = self.root
 
@@ -201,6 +202,11 @@ def determine_value_validator(self):
             self._type_info = str(data_type)
         return self
 
+    @model_validator(mode="after")
+    def determine_value_validator(self):
+        """Determines value validator (inner `data_type`)."""
+        return self._determine_value_validator()
+
     @property
     def data_type(self) -> Union[type, BaseModel]:
         """Type for incoming value validation"""
@@ -267,5 +273,22 @@ def __repr__(self):
         return f"'{str(self)}'"
 
 
+class EidDataType(DataType):
+    """Data type container for entity id
+
+    Represents one of primitive data types:
+    - string
+    - int
+    - ipv4
+    - ipv6
+    """
+
+    @model_validator(mode="after")
+    def determine_value_validator(self):
+        if self.root not in eid_data_types:
+            raise ValueError(f"Data type of entity ID must be one of {eid_data_types}")
+        return self._determine_value_validator()
+
+
 def is_primitive_element_type(data_type: DataType) -> bool:
     return data_type.root in primitive_data_types or data_type.is_link
diff --git a/dp3/common/entityspec.py b/dp3/common/entityspec.py
index d08b6d49..a7c9a6cb 100644
--- a/dp3/common/entityspec.py
+++ b/dp3/common/entityspec.py
@@ -2,6 +2,7 @@
 
 from pydantic import BaseModel, Extra, Field
 
+from dp3.common.datatype import EidDataType
 from dp3.common.types import ParsedTimedelta
 
 
@@ -51,6 +52,7 @@ class EntitySpec(SpecModel):
 
     id: str
     name: str
+    data_type: EidDataType = EidDataType("string")
     snapshot: bool
     lifetime: Union[ImmortalLifetime, TimeToLiveLifetime, WeakLifetime] = Field(
         default_factory=lambda: ImmortalLifetime(type="immortal"), discriminator="type"

From ac4957b79abbea3a993ae444af526fff01148630 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Wed, 6 Nov 2024 14:09:25 +0100
Subject: [PATCH 03/22] Support validation and enforce eid data_type.

---
 dp3/api/internal/entity_response_models.py |  32 ++++-
 dp3/api/routers/root.py                    |   1 +
 dp3/common/attrspec.py                     |  25 +++-
 dp3/common/config.py                       |  31 +++--
 dp3/common/datapoint.py                    |   2 +-
 dp3/common/entityspec.py                   |  41 +++++-
 dp3/common/task.py                         |  17 ++-
 dp3/common/types.py                        |  15 ++-
 dp3/common/utils.py                        |  12 ++
 dp3/database/database.py                   | 142 +++++++++++++++------
 dp3/history_management/history_manager.py  |  18 +--
 dp3/task_processing/task_executor.py       |   2 +-
 12 files changed, 256 insertions(+), 82 deletions(-)

diff --git a/dp3/api/internal/entity_response_models.py b/dp3/api/internal/entity_response_models.py
index 5aac18b8..e1b0c5d8 100644
--- a/dp3/api/internal/entity_response_models.py
+++ b/dp3/api/internal/entity_response_models.py
@@ -1,10 +1,11 @@
 from datetime import datetime
 from typing import Annotated, Any, Optional, Union
 
-from pydantic import BaseModel, Field, NonNegativeInt, PlainSerializer
+from pydantic import BaseModel, Field, NonNegativeInt, PlainSerializer, model_validator
 
-from dp3.common.attrspec import AttrSpecType, AttrType
+from dp3.common.attrspec import AttrSpec, AttrSpecGeneric, AttrSpecType, AttrType
 from dp3.common.datapoint import to_json_friendly
+from dp3.common.entityspec import EntitySpec, entity_context
 
 
 class EntityState(BaseModel):
@@ -15,10 +16,37 @@ class EntityState(BaseModel):
     """
 
     id: str
+    id_data_type: str
     name: str
     attribs: dict[str, Annotated[AttrSpecType, Field(discriminator="type")]]
     eid_estimate_count: NonNegativeInt
 
+    @model_validator(mode="before")
+    def _validate_attribs(cls, data: Any) -> Any:
+        if not isinstance(data, dict):
+            return data
+        if any(x not in data for x in ["id", "name"]):
+            return data
+
+        entity_spec = EntitySpec.model_validate(
+            {
+                "id": data["id"],
+                "name": data["name"],
+                "id_data_type": data.get("id_data_type", "string"),
+                "snapshot": True,
+            }
+        )
+        print(entity_spec)
+
+        assert isinstance(data["attribs"], dict), "'attribs' must be a dictionary"
+        with entity_context(entity_spec):
+            data["attribs"] = {
+                attr_id: AttrSpec(attr_id, spec) if not isinstance(spec, AttrSpecGeneric) else spec
+                for attr_id, spec in data["attribs"].items()
+            }
+
+        return data
+
 
 # This is necessary to allow for non-JSON-serializable types in the model
 JsonVal = Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")]
diff --git a/dp3/api/routers/root.py b/dp3/api/routers/root.py
index 9a6cf21b..df47fb1c 100644
--- a/dp3/api/routers/root.py
+++ b/dp3/api/routers/root.py
@@ -75,6 +75,7 @@ async def list_entity_types() -> dict[str, EntityState]:
         entity_spec = MODEL_SPEC.entity(etype)
         entities[etype] = {
             "id": etype,
+            "id_data_type": entity_spec.id_data_type.root,
             "name": entity_spec.name,
             "attribs": MODEL_SPEC.attribs(etype),
             "eid_estimate_count": DB.estimate_count_eids(etype),
diff --git a/dp3/common/attrspec.py b/dp3/common/attrspec.py
index 2b991974..1ac9849b 100644
--- a/dp3/common/attrspec.py
+++ b/dp3/common/attrspec.py
@@ -24,7 +24,7 @@
     dp_ts_v_validator,
 )
 from dp3.common.datatype import DataType, ReadOnly
-from dp3.common.entityspec import SpecModel
+from dp3.common.entityspec import SpecModel, _init_attr_spec_context_var
 from dp3.common.types import ParsedTimedelta
 
 # Regex of attribute and series id's
@@ -186,9 +186,14 @@ class AttrSpecPlain(AttrSpecClassic):
     def __init__(self, **data):
         super().__init__(**data)
 
+        entity_spec = _init_attr_spec_context_var.get()
+        if entity_spec is None:
+            raise ValueError("Entity specification not found in context")
+
         self._dp_model = create_model(
             f"DataPointPlain_{self.id}",
             __base__=DataPointPlainBase,
+            eid=(entity_spec.id_data_type.data_type, ...),
             v=(self.data_type.data_type, ...),
         )
 
@@ -199,9 +204,14 @@ class AttrSpecReadOnly(AttrSpecPlain):
     def __init__(self, **data):
         super().__init__(**data)
 
+        entity_spec = _init_attr_spec_context_var.get()
+        if entity_spec is None:
+            raise ValueError("Entity specification not found in context")
+
         self._dp_model = create_model(
             f"DataPointReadOnly_{self.id}",
             __base__=DataPointPlainBase,
+            eid=(entity_spec.id_data_type.data_type, ...),
             v=(ReadOnly, ...),
         )
 
@@ -221,10 +231,14 @@ def __init__(self, **data):
         super().__init__(**data)
 
         value_validator = self.data_type.data_type
+        entity_spec = _init_attr_spec_context_var.get()
+        if entity_spec is None:
+            raise ValueError("Entity specification not found in context")
 
         self._dp_model = create_model(
             f"DataPointObservations_{self.id}",
             __base__=DataPointObservationsBase,
+            eid=(entity_spec.id_data_type.data_type, ...),
             v=(value_validator, ...),
         )
 
@@ -242,6 +256,10 @@ class AttrSpecTimeseries(AttrSpecGeneric):
     def __init__(self, **data):
         super().__init__(**data)
 
+        entity_spec = _init_attr_spec_context_var.get()
+        if entity_spec is None:
+            raise ValueError("Entity specification not found in context")
+
         # Typing of `v` field
         dp_value_typing = {}
         for s in self.series:
@@ -268,6 +286,7 @@ def __init__(self, **data):
             f"DataPointTimeseries_{self.id}",
             __base__=DataPointTimeseriesBase,
             __validators__=dp_validators,
+            eid=(entity_spec.id_data_type.data_type, ...),
             v=(create_model(f"DataPointTimeseriesValue_{self.id}", **dp_value_typing), ...),
         )
 
@@ -304,4 +323,6 @@ def AttrSpec(id: str, spec: dict[str, Any]) -> AttrSpecType:
         AttrType.OBSERVATIONS: AttrSpecObservations,
         AttrType.TIMESERIES: AttrSpecTimeseries,
     }
-    return subclasses[attr_type](id=id, **spec)
+
+    spec["id"] = id  # Fill in the id if not present, else overwrite
+    return subclasses[attr_type](**spec)
diff --git a/dp3/common/config.py b/dp3/common/config.py
index 9754beeb..23351e0d 100644
--- a/dp3/common/config.py
+++ b/dp3/common/config.py
@@ -3,7 +3,7 @@
 """
 
 import os
-from typing import Annotated, Optional, Union
+from typing import Annotated, Any, Optional, Union
 
 import yaml
 from pydantic import (
@@ -25,7 +25,7 @@
     AttrSpecReadOnly,
     AttrSpecType,
 )
-from dp3.common.entityspec import EntitySpec
+from dp3.common.entityspec import EntitySpec, entity_context
 
 
 class NoDefault:
@@ -200,19 +200,22 @@ class EntitySpecDict(BaseModel):
     def __getitem__(self, item):
         return self.__getattribute__(item)
 
-    @field_validator("entity", mode="before")
-    def _parse_entity_spec(cls, v):
-        if isinstance(v, EntitySpec):
-            return v
-        return EntitySpec.model_validate(v)
+    @model_validator(mode="before")
+    def _validate_attr_spec(cls, data: Any) -> Any:
+        if not isinstance(data, dict) or "entity" not in data or "attribs" not in data:
+            return data
 
-    @field_validator("attribs", mode="before")
-    def _parse_attr_spec(cls, v):
-        assert isinstance(v, dict), "'attribs' must be a dictionary"
-        return {
-            attr_id: AttrSpec(attr_id, spec) if not isinstance(spec, AttrSpecGeneric) else spec
-            for attr_id, spec in v.items()
-        }
+        if not isinstance(data["entity"], EntitySpec):
+            data["entity"] = EntitySpec.model_validate(data["entity"])
+
+        assert isinstance(data["attribs"], dict), "'attribs' must be a dictionary"
+        with entity_context(data["entity"]):
+            data["attribs"] = {
+                attr_id: AttrSpec(attr_id, spec) if not isinstance(spec, AttrSpecGeneric) else spec
+                for attr_id, spec in data["attribs"].items()
+            }
+
+        return data
 
 
 class ModelSpec(BaseModel):
diff --git a/dp3/common/datapoint.py b/dp3/common/datapoint.py
index cbbd6575..4e650098 100644
--- a/dp3/common/datapoint.py
+++ b/dp3/common/datapoint.py
@@ -31,7 +31,7 @@ class DataPointBase(BaseModel, use_enum_values=True):
     """
 
     etype: str
-    eid: str
+    eid: Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")] = None
     attr: str
     src: Optional[str] = None
     v: Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")] = None
diff --git a/dp3/common/entityspec.py b/dp3/common/entityspec.py
index a7c9a6cb..077dc4dc 100644
--- a/dp3/common/entityspec.py
+++ b/dp3/common/entityspec.py
@@ -1,12 +1,15 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
 from typing import Literal, Union
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, Field, PrivateAttr, model_validator
 
 from dp3.common.datatype import EidDataType
 from dp3.common.types import ParsedTimedelta
 
 
-class SpecModel(BaseModel, extra=Extra.forbid): ...
+class SpecModel(BaseModel, extra="forbid"): ...
 
 
 class ImmortalLifetime(SpecModel):
@@ -48,14 +51,46 @@ class EntitySpec(SpecModel):
     """Entity specification
 
     This class represents specification of an entity type (e.g. ip, asn, ...)
+
+    Attributes:
+        id: Entity type identifier
+        id_data_type: Entity type identifier data type
+        name: User-friendly entity type name for display
+        snapshot: If `True`, the entity type supports snapshots
+        lifetime: Entity lifetime specification
+        description: Entity type description
     """
 
     id: str
+    id_data_type: EidDataType = EidDataType("string")
     name: str
-    data_type: EidDataType = EidDataType("string")
     snapshot: bool
     lifetime: Union[ImmortalLifetime, TimeToLiveLifetime, WeakLifetime] = Field(
         default_factory=lambda: ImmortalLifetime(type="immortal"), discriminator="type"
     )
 
     description: str = ""
+
+    _eid_validator: PrivateAttr()
+
+    @model_validator(mode="after")
+    def fill_validator(self):
+        """Fill the `eid_validator` attribute."""
+        self._eid_validator = self.id_data_type.data_type
+        return self
+
+    def validate_eid(self, v):
+        return self._eid_validator(v)
+
+
+_init_attr_spec_context_var = ContextVar("_init_attr_spec_context_var", default=None)
+
+
+@contextmanager
+def entity_context(entity_spec: EntitySpec) -> Iterator[None]:
+    """Context manager for setting the `model_spec` context variable."""
+    token = _init_attr_spec_context_var.set(entity_spec)
+    try:
+        yield
+    finally:
+        _init_attr_spec_context_var.reset(token)
diff --git a/dp3/common/task.py b/dp3/common/task.py
index 78007e87..22601bf6 100644
--- a/dp3/common/task.py
+++ b/dp3/common/task.py
@@ -11,13 +11,14 @@
     AfterValidator,
     BaseModel,
     BeforeValidator,
+    PlainSerializer,
     ValidationError,
     field_validator,
 )
 from pydantic_core.core_schema import FieldValidationInfo
 
 from dp3.common.config import ModelSpec
-from dp3.common.datapoint import DataPointBase
+from dp3.common.datapoint import DataPointBase, to_json_friendly
 
 _init_context_var = ContextVar("_init_context_var", default=None)
 
@@ -128,7 +129,7 @@ class DataPointTask(Task):
     """
 
     etype: str
-    eid: str
+    eid: Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")]
     data_points: list[ValidatedDataPoint] = []
     tags: list[Any] = []
     ttl_tokens: Optional[dict[str, datetime]] = None
@@ -160,6 +161,18 @@ def validate_etype(cls, v, info: FieldValidationInfo):
             raise AssertionError("Missing `model_spec` in context")
         return v
 
+    @field_validator("eid")
+    def validate_eid(cls, v, info: FieldValidationInfo):
+        if "etype" not in info.data:
+            return v
+
+        context = info.context
+        if context and "model_spec" in context:
+            ms: ModelSpec = context["model_spec"]
+            return ms.entities[info.data["etype"]].validate_eid(v)
+        else:
+            raise AssertionError("Missing `model_spec` in context")
+
 
 def parse_data_point_task(task: str, model_spec: ModelSpec) -> DataPointTask:
     with task_context(model_spec):
diff --git a/dp3/common/types.py b/dp3/common/types.py
index f8178767..296d0a7b 100644
--- a/dp3/common/types.py
+++ b/dp3/common/types.py
@@ -1,5 +1,7 @@
 from datetime import datetime, timedelta
-from typing import Annotated, Union
+from ipaddress import IPv4Address, IPv6Address
+from json import JSONEncoder
+from typing import Annotated, Any, Union
 
 from event_count_logger import DummyEventGroup, EventGroup
 from pydantic import AfterValidator, BeforeValidator
@@ -36,3 +38,14 @@ def t2_after_t1(v, info: FieldValidationInfo):
 T2Datetime = Annotated[datetime, BeforeValidator(t2_implicity_t1), AfterValidator(t2_after_t1)]
 
 EventGroupType = Union[EventGroup, DummyEventGroup]
+
+
+class DP3Encoder(JSONEncoder):
+    """JSONEncoder to encode python types using DP3 conventions."""
+
+    def default(self, o: Any) -> Any:
+        if isinstance(o, datetime):
+            return o.strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-4]
+        if isinstance(o, (IPv4Address, IPv6Address)):
+            return str(o)
+        return super().default(o)
diff --git a/dp3/common/utils.py b/dp3/common/utils.py
index ccaa5098..574c9afc 100644
--- a/dp3/common/utils.py
+++ b/dp3/common/utils.py
@@ -33,6 +33,18 @@ def int2ipstr(i):
     return ".".join((str(i >> 24), str((i >> 16) & 0xFF), str((i >> 8) & 0xFF), str(i & 0xFF)))
 
 
+def int2bytes(num: int) -> bytes:
+    """Convert signed int to however many bytes necessary in big-endian order"""
+    return num.to_bytes(
+        length=(8 + (num + (num < 0)).bit_length()) // 8, byteorder="big", signed=True
+    )
+
+
+def bytes2int(b: bytes) -> int:
+    """Convert bytes to signed int in big-endian order"""
+    return int.from_bytes(b, "big", signed=True)
+
+
 # *** Time conversion ***
 # Regex for RFC 3339 time format
 timestamp_re = re.compile(
diff --git a/dp3/database/database.py b/dp3/database/database.py
index e0a6233f..48451beb 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -29,6 +29,7 @@
 from dp3.common.scheduler import Scheduler
 from dp3.common.task import HASH
 from dp3.common.types import EventGroupType
+from dp3.common.utils import int2bytes
 from dp3.database.schema_cleaner import SchemaCleaner
 
 BSON_OBJECT_TOO_LARGE = 10334
@@ -487,7 +488,7 @@ def await_updated_schema(self):
         self.schema_cleaner.await_updated()
 
     def insert_datapoints(
-        self, etype: str, eid: str, dps: list[DataPointBase], new_entity: bool = False
+        self, eid: Any, dps: list[DataPointBase], new_entity: bool = False
     ) -> None:
         """Inserts datapoint to raw data collection and updates master record.
 
@@ -636,7 +637,7 @@ def update_master_records(self, etype: str, eids: list[str], records: list[dict]
         except Exception as e:
             raise DatabaseError(f"Update of master records failed: {e}\n{records}") from e
 
-    def extend_ttl(self, etype: str, eid: str, ttl_tokens: dict[str, datetime]):
+    def extend_ttl(self, etype: str, eid: Any, ttl_tokens: dict[str, datetime]):
         """Extends TTL of given `etype`:`eid` by `ttl_tokens`."""
         master_col = self._master_col(etype)
         try:
@@ -689,7 +690,7 @@ def delete_eids(self, etype: str, eids: list[str]):
         except Exception as e:
             raise DatabaseError(f"Delete of master record failed: {e}\n{eids}") from e
         try:
-            res = snapshot_col.delete_many(self._snapshot_bucket_eids_filter(eids))
+            res = snapshot_col.delete_many(self._snapshot_bucket_eids_filter(etype, eids))
             del_cnt = res.deleted_count * self._snapshot_bucket_size
             self.log.debug("Deleted %s snapshots of %s (%s).", del_cnt, etype, len(eids))
             res = os_snapshot_col.delete_many({"eid": {"$in": eids}})
@@ -704,7 +705,7 @@ def delete_eids(self, etype: str, eids: list[str]):
             except Exception as e:
                 self.log.exception("Error in on_entity_delete_many callback %s: %s", f, e)
 
-    def delete_eid(self, etype: str, eid: str):
+    def delete_eid(self, etype: str, eid: Any):
         """Delete master record and all snapshots of `etype`:`eid`."""
         master_col = self._master_col(etype)
         snapshot_col = self._snapshot_col(etype)
@@ -716,7 +717,7 @@ def delete_eid(self, etype: str, eid: str):
         except Exception as e:
             raise DatabaseError(f"Delete of master record failed: {e}\n{eid}") from e
         try:
-            res = snapshot_col.delete_many(self._snapshot_bucket_eid_filter(eid))
+            res = snapshot_col.delete_many(self._snapshot_bucket_eid_filter(etype, eid))
             del_cnt = res.deleted_count * self._snapshot_bucket_size
             self.log.debug("deleted %s snapshots of %s/%s.", del_cnt, etype, eid)
             res = os_snapshot_col.delete_many({"eid": eid})
@@ -873,18 +874,19 @@ def delete_many_link_dps(
         except Exception as e:
             raise DatabaseError(f"Delete of link datapoints failed: {e}") from e
 
-    def get_master_record(self, etype: str, eid: str, **kwargs) -> dict:
+    def get_master_record(self, etype: str, eid: Any, **kwargs) -> dict:
         """Get current master record for etype/eid.
 
         If doesn't exist, returns {}.
         """
         # Check `etype`
         self._assert_etype_exists(etype)
-
+        entity_spec = self._db_schema_config.entities[etype]
         master_col = self._master_col(etype)
-        return master_col.find_one({"_id": eid}, **kwargs) or {}
 
-    def ekey_exists(self, etype: str, eid: str) -> bool:
+        return master_col.find_one({"_id": entity_spec.validate_eid(eid)}, **kwargs) or {}
+
+    def ekey_exists(self, etype: str, eid: Any) -> bool:
         """Checks whether master record for etype/eid exists"""
         return bool(self.get_master_record(etype, eid))
 
@@ -909,7 +911,7 @@ def get_worker_master_records(
             {"#hash": {"$mod": [worker_cnt, worker_index]}, **query_filter}, **kwargs
         )
 
-    def get_latest_snapshot(self, etype: str, eid: str) -> dict:
+    def get_latest_snapshot(self, etype: str, eid: Any) -> dict:
         """Get latest snapshot of given etype/eid.
 
         If doesn't exist, returns {}.
@@ -920,7 +922,7 @@ def get_latest_snapshot(self, etype: str, eid: str) -> dict:
         snapshot_col = self._snapshot_col(etype)
         return (
             snapshot_col.find_one(
-                self._snapshot_bucket_eid_filter(eid), {"last": 1}, sort=[("_id", -1)]
+                self._snapshot_bucket_eid_filter(etype, eid), {"last": 1}, sort=[("_id", -1)]
             )
             or {}
         ).get("last", {})
@@ -1023,20 +1025,75 @@ def get_latest_snapshots(
         except OperationFailure as e:
             raise DatabaseError("Invalid query") from e
 
-    @staticmethod
-    def _snapshot_bucket_id(eid: str, ctime: datetime) -> str:
-        return f"{eid}_#{int(ctime.timestamp())}"
+    def _snapshot_bucket_id(self, etype: str, eid: Any, ctime: datetime) -> Union[str, Binary]:
+        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
+        ts_int = int(ctime.timestamp())
 
-    @staticmethod
-    def _snapshot_bucket_eid_filter(eid: str) -> dict:
-        return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
+        if entity_data_type == "string":
+            return f"{eid}_#{ts_int}"
+        if entity_data_type == "int" and isinstance(eid, int):
+            return self._pack_binary_snapshot_bucket_id(int2bytes(eid), ts_int)
+        if entity_data_type == "ipv4" and isinstance(eid, IPv4Address):
+            return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
+        if entity_data_type == "ipv6" and isinstance(eid, IPv6Address):
+            return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
+        raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}', {eid}")
 
     @staticmethod
-    def _snapshot_bucket_eids_filter(eids: Iterable[str]) -> dict:
-        return {"_id": {"$regex": "|".join([f"^{re.escape(eid)}_#" for eid in eids])}}
+    def _pack_binary_snapshot_bucket_id(eid: bytes, ts_int: int) -> Binary:
+        return Binary(b"".join([eid, ts_int.to_bytes(8, "big", signed=True)]))
+
+    def _snapshot_bucket_eid_filter(self, etype: str, eid: Any) -> dict:
+        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
+
+        if entity_data_type == "string":
+            eid = str(eid)
+            return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
+        if entity_data_type == "int":
+            eid = int(eid)
+            return {"_id": self._binary_snapshot_bucket_range(int2bytes(eid))}
+        if entity_data_type == "ipv4":
+            eid = IPv4Address(eid)
+            return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
+        if entity_data_type == "ipv6":
+            eid = IPv6Address(eid)
+            return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
+        raise ValueError(
+            f"Unsupported data type '{entity_data_type}' for entity {etype} '{eid}' ({type(eid)})"
+        )
+
+    def _snapshot_eid_filter_from_bid(self, etype: str, b_id: Union[bytes, str]) -> dict:
+        """Returns filter for snapshots with same eid as given bucket document _id.
+        Args:
+            etype: entity type
+            b_id: the _id of the snapshot bucket, type depends on etype's data type
+        """
+        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
+
+        if entity_data_type == "string":
+            eid = b_id.rsplit("_#", maxsplit=1)[0]
+            return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
+        if entity_data_type in ["int", "ipv4", "ipv6"]:
+            eid_bytes = b_id[:-8]
+            return {"_id": self._binary_snapshot_bucket_range(eid_bytes)}
+        raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}'")
+
+    def _binary_snapshot_bucket_range(self, eid: bytes) -> dict:
+        return {
+            "$gte": self._pack_binary_snapshot_bucket_id(eid, 0),
+            "$lt": self._pack_binary_snapshot_bucket_id(eid, -1),
+        }
+
+    def _snapshot_bucket_eids_filter(self, etype: str, eids: Iterable[Any]) -> dict:
+        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
+        if entity_data_type == "string":
+            return {"_id": {"$regex": "|".join([f"^{re.escape(eid)}_#" for eid in eids])}}
+        if entity_data_type in ["int", "ipv4", "ipv6"]:
+            return {"$or": [self._snapshot_bucket_eid_filter(etype, eid) for eid in eids]}
+        raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}'")
 
     def get_snapshots(
-        self, etype: str, eid: str, t1: Optional[datetime] = None, t2: Optional[datetime] = None
+        self, etype: str, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
     ) -> Union[Cursor, CommandCursor]:
         """Get all (or filtered) snapshots of given `eid`.
 
@@ -1055,7 +1112,7 @@ def get_snapshots(
 
         # Find out if the snapshot is oversized
         doc = (
-            snapshot_col.find(self._snapshot_bucket_eid_filter(eid), {"oversized": 1})
+            snapshot_col.find(self._snapshot_bucket_eid_filter(etype, eid), {"oversized": 1})
             .sort([("_id", -1)])
             .limit(1)
         )
@@ -1065,7 +1122,7 @@ def get_snapshots(
 
         query = {"_time_created": {}}
         pipeline = [
-            {"$match": self._snapshot_bucket_eid_filter(eid)},
+            {"$match": self._snapshot_bucket_eid_filter(etype, eid)},
             {"$unwind": "$history"},
             {"$replaceRoot": {"newRoot": "$history"}},
         ]
@@ -1083,7 +1140,7 @@ def get_snapshots(
         return snapshot_col.aggregate(pipeline)
 
     def _get_snapshots_oversized(
-        self, etype: str, eid: str, t1: Optional[datetime] = None, t2: Optional[datetime] = None
+        self, etype: str, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
     ) -> Cursor:
         """Get all (or filtered) snapshots of given `eid` from oversized snapshots collection."""
         snapshot_col = self._oversized_snapshots_col(etype)
@@ -1107,7 +1164,7 @@ def get_value_or_history(
         self,
         etype: str,
         attr_name: str,
-        eid: str,
+        eid: Any,
         t1: Optional[datetime] = None,
         t2: Optional[datetime] = None,
     ) -> dict:
@@ -1151,19 +1208,21 @@ def estimate_count_eids(self, etype: str) -> int:
         master_col = self._master_col(etype)
         return master_col.estimated_document_count({})
 
-    def _migrate_to_oversized_snapshot(self, etype: str, eid: str, snapshot: dict):
+    def _migrate_to_oversized_snapshot(self, etype: str, eid: Any, snapshot: dict):
         snapshot_col = self._snapshot_col(etype)
         os_col = self._oversized_snapshots_col(etype)
 
         try:
             move_to_oversized = []
             last_id = None
-            for doc in snapshot_col.find(self._snapshot_bucket_eid_filter(eid)).sort({"_id": 1}):
+            for doc in snapshot_col.find(self._snapshot_bucket_eid_filter(etype, eid)).sort(
+                {"_id": 1}
+            ):
                 move_to_oversized.extend(doc.get("history", []))
                 last_id = doc["_id"]
             move_to_oversized.insert(0, snapshot)
 
-            self._db[os_col].insert_many(move_to_oversized)
+            os_col.insert_many(move_to_oversized)
             snapshot_col.update_one(
                 {"_id": last_id},
                 {
@@ -1171,7 +1230,9 @@ def _migrate_to_oversized_snapshot(self, etype: str, eid: str, snapshot: dict):
                     "$unset": {"history": ""},
                 },
             )
-            snapshot_col.delete_many(self._snapshot_bucket_eid_filter(eid) | {"oversized": False})
+            snapshot_col.delete_many(
+                self._snapshot_bucket_eid_filter(etype, eid) | {"oversized": False}
+            )
         except Exception as e:
             raise DatabaseError(f"Update of snapshot {eid} failed: {e}, {snapshot}") from e
 
@@ -1197,14 +1258,14 @@ def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
         if normal:
             try:
                 res = snapshot_col.update_one(
-                    self._snapshot_bucket_eid_filter(eid)
+                    self._snapshot_bucket_eid_filter(etype, eid)
                     | {"count": {"$lt": self._snapshot_bucket_size}},
                     {
                         "$set": {"last": snapshot},
                         "$push": {"history": {"$each": [snapshot], "$position": 0}},
                         "$inc": {"count": 1},
                         "$setOnInsert": {
-                            "_id": self._snapshot_bucket_id(eid, ctime),
+                            "_id": self._snapshot_bucket_id(etype, eid, ctime),
                             "_time_created": ctime,
                             "oversized": False,
                             "latest": True,
@@ -1215,7 +1276,7 @@ def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
 
                 if res.upserted_id is not None:
                     snapshot_col.update_many(
-                        self._snapshot_bucket_eid_filter(eid)
+                        self._snapshot_bucket_eid_filter(etype, eid)
                         | {"latest": True, "count": self._snapshot_bucket_size},
                         {"$unset": {"latest": 1}},
                     )
@@ -1232,7 +1293,7 @@ def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
         elif oversized:
             # Snapshot is already marked as oversized
             snapshot_col.update_one(
-                self._snapshot_bucket_eid_filter(eid), {"$set": {"last": snapshot}}
+                self._snapshot_bucket_eid_filter(etype, eid), {"$set": {"last": snapshot}}
             )
             os_col.insert_one(snapshot)
             return
@@ -1250,7 +1311,7 @@ def _get_snapshot_state(self, etype: str, eids: set[str]) -> tuple[set, set]:
         snapshot_col = self._snapshot_col(etype)
         new_oversized = set()
         for doc in snapshot_col.find(
-            self._snapshot_bucket_eids_filter(unknown) | {"oversized": True},
+            self._snapshot_bucket_eids_filter(etype, unknown) | {"oversized": True},
             {"oversized": 1},
         ):
             eid = doc["_id"].rsplit("_#", maxsplit=1)[0]
@@ -1299,7 +1360,7 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
         normal, oversized = self._get_snapshot_state(etype, set(snapshots_by_eid.keys()))
 
         upserts = []
-        update_originals = []
+        update_originals: list[list[dict]] = []
         oversized_inserts = []
         oversized_updates = []
 
@@ -1307,14 +1368,14 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
         for eid in normal:
             upserts.append(
                 UpdateOne(
-                    self._snapshot_bucket_eid_filter(eid)
+                    self._snapshot_bucket_eid_filter(etype, eid)
                     | {"count": {"$lt": self._snapshot_bucket_size}},
                     {
                         "$set": {"last": snapshots_by_eid[eid][-1]},
                         "$push": {"history": {"$each": snapshots_by_eid[eid], "$position": 0}},
                         "$inc": {"count": len(snapshots_by_eid[eid])},
                         "$setOnInsert": {
-                            "_id": self._snapshot_bucket_id(eid, ctime),
+                            "_id": self._snapshot_bucket_id(etype, eid, ctime),
                             "_time_created": ctime,
                             "oversized": False,
                             "latest": True,
@@ -1330,7 +1391,7 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
             oversized_inserts.extend(snapshots_by_eid[eid])
             oversized_updates.append(
                 UpdateOne(
-                    self._snapshot_bucket_eid_filter(eid),
+                    self._snapshot_bucket_eid_filter(etype, eid),
                     {"$set": {"last": snapshots_by_eid[eid][-1]}},
                 )
             )
@@ -1345,10 +1406,9 @@ def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
                 if res.upserted_count > 0:
                     unset_latest_updates = []
                     for upsert_id in res.upserted_ids.values():
-                        eid = upsert_id.rsplit("_#", maxsplit=1)[0]
                         unset_latest_updates.append(
                             UpdateMany(
-                                self._snapshot_bucket_eid_filter(eid)
+                                self._snapshot_eid_filter_from_bid(etype, upsert_id)
                                 | {"latest": True, "count": self._snapshot_bucket_size},
                                 {"$unset": {"latest": 1}},
                             )
@@ -1441,7 +1501,7 @@ def get_observation_history(
         self,
         etype: str,
         attr_name: str,
-        eid: str,
+        eid: Any,
         t1: datetime = None,
         t2: datetime = None,
         sort: int = None,
@@ -1484,7 +1544,7 @@ def get_timeseries_history(
         self,
         etype: str,
         attr_name: str,
-        eid: str,
+        eid: Any,
         t1: datetime = None,
         t2: datetime = None,
         sort: int = None,
diff --git a/dp3/history_management/history_manager.py b/dp3/history_management/history_manager.py
index 26c0fd0b..7619d89d 100644
--- a/dp3/history_management/history_manager.py
+++ b/dp3/history_management/history_manager.py
@@ -3,9 +3,8 @@
 import logging
 import os
 from datetime import datetime
-from json import JSONEncoder
 from pathlib import Path
-from typing import Any, Optional
+from typing import Optional
 
 from pydantic import BaseModel, Extra
 
@@ -17,22 +16,13 @@
 )
 from dp3.common.callback_registrar import CallbackRegistrar
 from dp3.common.config import CronExpression, PlatformConfig
-from dp3.common.types import ParsedTimedelta
+from dp3.common.types import DP3Encoder, ParsedTimedelta
 from dp3.common.utils import entity_expired
 from dp3.database.database import DatabaseError, EntityDatabase
 
 DB_SEND_CHUNK = 100
 
 
-class DatetimeEncoder(JSONEncoder):
-    """JSONEncoder to encode datetime using the standard ADiCT format string."""
-
-    def default(self, o: Any) -> Any:
-        if isinstance(o, datetime):
-            return o.strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-4]
-        return super().default(o)
-
-
 class SnapshotCleaningConfig(BaseModel):
     """Configuration for snapshot cleaning.
 
@@ -218,9 +208,7 @@ def archive_old_dps(self):
                 for etype in self.model_spec.entities:
                     for result_cursor in self.db.get_archive(etype, after=min_date, before=t_old):
                         for dp in result_cursor:
-                            logfile.write(
-                                f"{json.dumps(self._reformat_dp(dp), cls=DatetimeEncoder)}\n"
-                            )
+                            logfile.write(f"{json.dumps(self._reformat_dp(dp), cls=DP3Encoder)}\n")
                             datapoints += 1
 
             self.log.info("Archived %s datapoints to %s", datapoints, date_logfile)
diff --git a/dp3/task_processing/task_executor.py b/dp3/task_processing/task_executor.py
index 5ab97c62..5703820d 100644
--- a/dp3/task_processing/task_executor.py
+++ b/dp3/task_processing/task_executor.py
@@ -152,7 +152,7 @@ def process_task(self, task: DataPointTask) -> tuple[bool, list[DataPointTask]]:
 
         # Insert into database
         try:
-            self.db.insert_datapoints(task.etype, task.eid, task.data_points, new_entity=new_entity)
+            self.db.insert_datapoints(task.eid, task.data_points, new_entity=new_entity)
             self.log.debug(f"Task {task.etype}/{task.eid}: All changes written to DB")
         except DatabaseError as e:
             self.log.error(f"Task {task.etype}/{task.eid}: DB error: {e}")

From f81080cc483379bc89b78160da3c7eb10a1a76e6 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Fri, 8 Nov 2024 13:34:00 +0100
Subject: [PATCH 04/22] Added MAC address type as eid type and binary subtype.

---
 dp3/common/datapoint.py               |  3 +-
 dp3/common/datatype.py                | 10 +--
 dp3/common/mac_address.py             | 91 +++++++++++++++++++++++++++
 dp3/database/database.py              | 15 ++++-
 tests/test_common/test_mac_address.py | 71 +++++++++++++++++++++
 5 files changed, 183 insertions(+), 7 deletions(-)
 create mode 100644 dp3/common/mac_address.py
 create mode 100644 tests/test_common/test_mac_address.py

diff --git a/dp3/common/datapoint.py b/dp3/common/datapoint.py
index 4e650098..1753e427 100644
--- a/dp3/common/datapoint.py
+++ b/dp3/common/datapoint.py
@@ -4,11 +4,12 @@
 
 from pydantic import BaseModel, BeforeValidator, Field, PlainSerializer
 
+from dp3.common.mac_address import MACAddress
 from dp3.common.types import T2Datetime
 
 
 def to_json_friendly(v):
-    if isinstance(v, (IPv4Address, IPv6Address)):
+    if isinstance(v, (IPv4Address, IPv6Address, MACAddress)):
         return str(v)
     return v
 
diff --git a/dp3/common/datatype.py b/dp3/common/datatype.py
index 9fb36f5f..bef7cb56 100644
--- a/dp3/common/datatype.py
+++ b/dp3/common/datatype.py
@@ -2,12 +2,11 @@
 import re
 from datetime import datetime
 from enum import Enum
-from typing import Annotated, Any, Optional, Union
+from typing import Any, Optional, Union
 
 from pydantic import (
     BaseModel,
     Extra,
-    Field,
     Json,
     PrivateAttr,
     RootModel,
@@ -15,6 +14,8 @@
     model_validator,
 )
 
+from dp3.common.mac_address import MACAddress
+
 # Regular expressions for parsing various data types
 re_array = re.compile(r"^array<(.+)>$")
 re_set = re.compile(r"^set<(.+)>$")
@@ -34,13 +35,13 @@
     "float": float,
     "ipv4": ipaddress.IPv4Address,
     "ipv6": ipaddress.IPv6Address,
-    "mac": Annotated[str, Field(pattern=r"^([0-9a-fA-F]{2}[:-]){5}([0-9a-fA-F]{2})$")],
+    "mac": MACAddress,
     "time": datetime,
     "special": Any,
     "json": Union[Json[Any], dict, list],
 }
 
-eid_data_types = ["string", "int", "ipv4", "ipv6"]
+eid_data_types = ["string", "int", "ipv4", "ipv6", "mac"]
 
 
 class ReadOnly(BaseModel):
@@ -281,6 +282,7 @@ class EidDataType(DataType):
     - int
     - ipv4
     - ipv6
+    - mac
     """
 
     @model_validator(mode="after")
diff --git a/dp3/common/mac_address.py b/dp3/common/mac_address.py
new file mode 100644
index 00000000..6d9d85f2
--- /dev/null
+++ b/dp3/common/mac_address.py
@@ -0,0 +1,91 @@
+from typing import Any, Union
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
+
+
+class MACAddress:
+    def __init__(self, mac: Union[bytes, str, "MACAddress"]):
+        if isinstance(mac, self.__class__):
+            mac = mac.mac  # type: ignore
+        if not isinstance(mac, bytes) or len(mac) != 6:
+            mac = self._parse_mac(mac)
+
+        self.mac: bytes = mac
+
+    @classmethod
+    def _validate(cls, value: Any) -> "MACAddress":
+        return cls(value)
+
+    @classmethod
+    def _serialize(cls, value: "MACAddress", info: Any) -> Any:
+        if info.mode == "json":
+            return str(value)
+        return value
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        base_schema = core_schema.no_info_after_validator_function(
+            cls._validate, handler(Union[str, bytes])
+        )
+
+        python_schema = core_schema.union_schema(
+            [
+                core_schema.is_instance_schema(cls=cls),
+                base_schema,
+            ],
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=base_schema,
+            python_schema=python_schema,
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                cls._serialize, info_arg=True
+            ),
+        )
+
+    @staticmethod
+    def _parse_mac(mac: Union[bytes, str]) -> bytes:
+        if isinstance(mac, str):
+            mac = mac.encode()
+        if not isinstance(mac, bytes):
+            raise TypeError(f"Expected str or bytes, got {type(mac)}")
+
+        if len(mac) != 17:
+            raise ValueError(f"Invalid MAC address length, {len(mac)} != 17")
+        if mac[2] not in [ord(":"), ord("-")]:
+            print(mac[2])
+            raise ValueError("Invalid MAC address format, expected separator ':' or '-'")
+
+        mac_bytes = bytearray(6)
+        for i in range(0, 18, 3):
+            mac_bytes[i // 3] = int(mac[i : i + 2], 16)
+            if i == 15:
+                break
+            if mac[i + 2] != mac[2]:
+                raise ValueError(
+                    f"Invalid MAC address format at pos {i + 2}, "
+                    f"expected separator '{chr(mac[2])}' but got '{chr(mac[i + 2])}'"
+                )
+        return bytes(mac_bytes)
+
+    def __str__(self):
+        return ":".join(f"{b:02x}" for b in self.mac)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self})"
+
+    def __eq__(self, other):
+        return isinstance(other, MACAddress) and self.mac == other.mac
+
+    def __hash__(self):
+        return hash(self.mac)
+
+    def __bytes__(self):
+        return self.mac
+
+    @property
+    def packed(self):
+        return self.mac
diff --git a/dp3/database/database.py b/dp3/database/database.py
index 48451beb..48d9e8d2 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -26,6 +26,7 @@
 from dp3.common.attrspec import AttrType, timeseries_types
 from dp3.common.config import HierarchicalDict, ModelSpec
 from dp3.common.datapoint import DataPointBase
+from dp3.common.mac_address import MACAddress
 from dp3.common.scheduler import Scheduler
 from dp3.common.task import HASH
 from dp3.common.types import EventGroupType
@@ -36,6 +37,7 @@
 
 BSON_IPV4_SUBTYPE = USER_DEFINED_SUBTYPE + 1
 BSON_IPV6_SUBTYPE = BSON_IPV4_SUBTYPE + 1
+BSON_MAC_SUBTYPE = BSON_IPV6_SUBTYPE + 1
 
 
 def fallback_encoder(value):
@@ -43,6 +45,8 @@ def fallback_encoder(value):
         return Binary(value.packed, BSON_IPV4_SUBTYPE)
     if isinstance(value, IPv6Address):
         return Binary(value.packed, BSON_IPV6_SUBTYPE)
+    if isinstance(value, MACAddress):
+        return Binary(value.mac, BSON_MAC_SUBTYPE)
     return value
 
 
@@ -54,6 +58,8 @@ def transform_bson(self, value):
             return IPv4Address(value)
         if value.subtype == BSON_IPV6_SUBTYPE:
             return IPv6Address(value)
+        if value.subtype == BSON_MAC_SUBTYPE:
+            return MACAddress(value)
         return value
 
 
@@ -1037,6 +1043,8 @@ def _snapshot_bucket_id(self, etype: str, eid: Any, ctime: datetime) -> Union[st
             return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
         if entity_data_type == "ipv6" and isinstance(eid, IPv6Address):
             return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
+        if entity_data_type == "mac" and isinstance(eid, MACAddress):
+            return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
         raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}', {eid}")
 
     @staticmethod
@@ -1058,6 +1066,9 @@ def _snapshot_bucket_eid_filter(self, etype: str, eid: Any) -> dict:
         if entity_data_type == "ipv6":
             eid = IPv6Address(eid)
             return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
+        if entity_data_type == "mac":
+            eid = MACAddress(eid)
+            return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
         raise ValueError(
             f"Unsupported data type '{entity_data_type}' for entity {etype} '{eid}' ({type(eid)})"
         )
@@ -1073,7 +1084,7 @@ def _snapshot_eid_filter_from_bid(self, etype: str, b_id: Union[bytes, str]) ->
         if entity_data_type == "string":
             eid = b_id.rsplit("_#", maxsplit=1)[0]
             return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
-        if entity_data_type in ["int", "ipv4", "ipv6"]:
+        if entity_data_type in ["int", "ipv4", "ipv6", "mac"]:
             eid_bytes = b_id[:-8]
             return {"_id": self._binary_snapshot_bucket_range(eid_bytes)}
         raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}'")
@@ -1088,7 +1099,7 @@ def _snapshot_bucket_eids_filter(self, etype: str, eids: Iterable[Any]) -> dict:
         entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
         if entity_data_type == "string":
             return {"_id": {"$regex": "|".join([f"^{re.escape(eid)}_#" for eid in eids])}}
-        if entity_data_type in ["int", "ipv4", "ipv6"]:
+        if entity_data_type in ["int", "ipv4", "ipv6", "mac"]:
             return {"$or": [self._snapshot_bucket_eid_filter(etype, eid) for eid in eids]}
         raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}'")
 
diff --git a/tests/test_common/test_mac_address.py b/tests/test_common/test_mac_address.py
new file mode 100644
index 00000000..30408fca
--- /dev/null
+++ b/tests/test_common/test_mac_address.py
@@ -0,0 +1,71 @@
+"""Test the MAC address parsing of MACAddress class"""
+
+import json
+import unittest
+
+from pydantic import BaseModel
+
+from dp3.common.mac_address import MACAddress
+
+
+class TestModel(BaseModel):
+    mac: MACAddress
+
+
+class TestMACAddress(unittest.TestCase):
+    def test_parse_mac_valid(self):
+        self.assertEqual(MACAddress._parse_mac("01:23:45:67:89:ab"), b"\x01\x23\x45\x67\x89\xab")
+        self.assertEqual(MACAddress._parse_mac("01-23-45-67-89-ab"), b"\x01\x23\x45\x67\x89\xab")
+        self.assertEqual(MACAddress._parse_mac("01:23:45:67:89:AB"), b"\x01\x23\x45\x67\x89\xab")
+        self.assertEqual(MACAddress._parse_mac("01-23-45-67-89-Ab"), b"\x01\x23\x45\x67\x89\xab")
+
+    def test_parse_mac_invalid_length(self):
+        with self.assertRaises(ValueError):
+            MACAddress._parse_mac("01:23:45:67:89")
+        with self.assertRaises(ValueError):
+            MACAddress._parse_mac("01:23:45:67:89:ab:cd")
+
+    def test_parse_mac_invalid_format(self):
+        invalid_macs = [
+            "01:23:45:67:89:ab:",
+            "01:23-45:67-89:ab",
+            "01:23:45:67:89:ag",
+            "01:23:45:67:89:ab:cd",
+            "",
+        ]
+        for mac in invalid_macs:
+            with self.assertRaises(ValueError):
+                MACAddress._parse_mac(mac)
+
+    def test_model_valid(self):
+        model = TestModel(mac="01:23:45:67:89:ab")
+        self.assertEqual(model.mac.packed, b"\x01\x23\x45\x67\x89\xab")
+
+        model = TestModel(mac=MACAddress("01:23:45:67:89:ab"))
+        self.assertEqual(model.mac.packed, b"\x01\x23\x45\x67\x89\xab")
+
+        model = TestModel.model_validate({"mac": b"\x01\x23\x45\x67\x89\xab"})
+        self.assertEqual(model.mac.packed, b"\x01\x23\x45\x67\x89\xab")
+
+        model = TestModel.model_validate({"mac": b"01:23:45:67:89:ab"})
+        self.assertEqual(model.mac.packed, b"\x01\x23\x45\x67\x89\xab")
+
+        model = TestModel.model_validate_json('{"mac": "01:23:45:67:89:ab"}')
+        self.assertEqual(model.mac.packed, b"\x01\x23\x45\x67\x89\xab")
+
+    def test_model_invalid(self):
+        with self.assertRaises(ValueError):
+            TestModel.model_validate({"mac": "01:23:45:67:89"})
+
+        with self.assertRaises(ValueError):
+            TestModel.model_validate({"mac": 420})
+
+        with self.assertRaises(ValueError):
+            TestModel.model_validate_json('{"mac": "01:23:45:67:89"}')
+
+    def test_model_serialization(self):
+        model = TestModel(mac="01:23:45:67:89:ab")
+        self.assertEqual(model.model_dump(), {"mac": MACAddress("01:23:45:67:89:ab")})
+
+        model = TestModel(mac=b"\x01\x23\x45\x67\x89\xab")
+        self.assertEqual(json.loads(model.model_dump_json()), {"mac": "01:23:45:67:89:ab"})

From 464724ee2c15acc0993bf5dca06989ae021f613c Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Wed, 6 Nov 2024 17:22:06 +0100
Subject: [PATCH 05/22] API & DB: Search and replace for generic filter.

---
 dp3/api/routers/entity.py       |  57 ++++++++++-
 dp3/database/database.py        |  17 +++-
 dp3/database/magic.py           | 167 ++++++++++++++++++++++++++++++++
 tests/test_common/test_magic.py | 132 +++++++++++++++++++++++++
 4 files changed, 367 insertions(+), 6 deletions(-)
 create mode 100644 dp3/database/magic.py
 create mode 100644 tests/test_common/test_magic.py

diff --git a/dp3/api/routers/entity.py b/dp3/api/routers/entity.py
index f8dc3654..be39ffce 100644
--- a/dp3/api/routers/entity.py
+++ b/dp3/api/routers/entity.py
@@ -101,6 +101,61 @@ async def list_entity_type_eids(
 
     Generic filter allows filtering using generic MongoDB query (including `$and`, `$or`,
     `$lt`, etc.).
+    For querying non-JSON-native types, you can use the following magic strings:
+
+    - `"$$IPv4{<ip address>}"` - converts to IPv4Address object
+    - `"$$IPv6{<ip address>}"` - converts to IPv6Address object
+    - `"$$int{<value>}"` - may be necessary for filtering when `eid` data type is int
+    - `"$$Date{<YYYY-mm-ddTHH:MM:ssZ>}"` - converts specified UTC date to UTC datetime object
+    - `"$$DateTs{<POSIX timestamp>}"` - converts POSIX timestamp (int/float) to UTC datetime object
+    - `"$$MAC{<mac address>}"` - converts to MACAddress object
+
+    To query an IP prefix, use the following magic strings:
+
+    - `"$$IPv4Prefix{<ip address>/<prefix length>}"` - matches prefix
+    - `"$$IPv6Prefix{<ip address>/<prefix length>}"` - matches prefix
+
+    To query a binary `_id`s of non-string snapshot buckets,
+    use the following magic string:
+
+    - `"$$Binary_ID{<EID object | Valid magic string>}"`
+
+        - converts to filter the exact EID object snapshots, only EID valid types are supported
+
+    There are no attribute name checks (may be added in the future).
+
+    Generic filter examples:
+
+    - `{"attr1": {"$gt": 10}, "attr2": {"$lt": 20}}`
+    - `{"ip_attr": "$$IPv4{127.0.0.1}"}` - converts to IPv4Address object, exact match
+    - `{"ip_attr": "$$IPv4Prefix{127.0.0.0/24}"}`
+        - converts to `{"ip_attr": {"$gte": "$$IPv4{127.0.0.0}",
+          "$lte": "$$IPv4{127.0.0.255}"}}`
+
+    - `{"ip_attr": "$$IPv6{::1}"}` - converts to IPv6Address object, exact match
+    - `{"ip_attr": "$$IPv6Prefix{::1/64}"}`
+        - converts to `{"ip_attr": {"$gte": "$$IPv6{::1}",
+          "$lte": "$$IPv6{::1:ffff:ffff:ffff:ffff}"}}`
+
+    - `{"_id": "$$Binary_ID{$$IPv4{127.0.0.1}}"}`
+        - converts to `{"_id": {"$gte": Binary(<IP bytes + min timestamp>),
+          "$lt": Binary(<IP bytes + max timestamp>)}}`
+
+    - `{"id": "$$Binary_ID{$$IPv4Prefix{127.0.0.0/24}}"}`
+        - converts to `{"_id": {"$gte": Binary(<IP 127.0.0.0 bytes + min timestamp>),
+          "$lte": Binary(<IP 127.0.0.255 bytes + max timestamp>)}}`
+
+    - `{"_time_created": {"$gte": "$$Date{2024-11-07T00:00:00Z}"}}`
+        - converts to `{"_time_created": datetime(2024, 11, 7, 0, 0, 0, tzinfo=timezone.utc)}`
+
+    - `{"_time_created": {"$gte": "$$DateTs{1609459200}"}}`
+        - converts to `{"_time_created": datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc)}`
+
+    - `{"attr": "$$MAC{00:11:22:33:44:55}"}` - converts to MACAddress object, exact match
+    - `{"_id": "$$Binary_ID{$$MAC{Ab-cD-Ef-12-34-56}}"}`
+        - converts to `{"_id": {"$gte": Binary(<MAC bytes + min timestamp>),
+            "$lt": Binary(<MAC bytes + max timestamp>)}}`
+
     There are no attribute name checks (may be added in the future).
 
     Generic and fulltext filters are merged - fulltext overrides conflicting keys.
@@ -128,7 +183,7 @@ async def list_entity_type_eids(
         cursor, total_count = DB.get_latest_snapshots(etype, fulltext_filters, generic_filter)
         cursor_page = cursor.skip(skip).limit(limit)
     except DatabaseError as e:
-        raise HTTPException(status_code=400, detail="Query is invalid") from e
+        raise HTTPException(status_code=400, detail=str(e)) from e
 
     time_created = None
 
diff --git a/dp3/database/database.py b/dp3/database/database.py
index 48d9e8d2..69df79bf 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -31,6 +31,7 @@
 from dp3.common.task import HASH
 from dp3.common.types import EventGroupType
 from dp3.common.utils import int2bytes
+from dp3.database.magic import search_and_replace
 from dp3.database.schema_cleaner import SchemaCleaner
 
 BSON_OBJECT_TOO_LARGE = 10334
@@ -962,7 +963,7 @@ def get_latest_snapshots(
         self,
         etype: str,
         fulltext_filters: Optional[dict[str, str]] = None,
-        generic_filter: Optional[dict[str, any]] = None,
+        generic_filter: Optional[dict[str, Any]] = None,
     ) -> tuple[Cursor, int]:
         """Get latest snapshots of given `etype`.
 
@@ -976,11 +977,14 @@ def get_latest_snapshots(
         Only plain and observation attributes with string-based data types can be queried.
         Array and set data types are supported as well as long as they are not multi value
         at the same time.
-        If you need to filter EIDs, use attribute `eid`.
+        If you need to filter EIDs, ensure the EID is string, then use attribute `eid`.
+        Otherwise, use generic filter.
 
         Generic filter allows filtering using generic MongoDB query (including `$and`, `$or`,
         `$lt`, etc.).
-        There are no attribute name checks (may be added in the future).
+        For querying non-JSON-native types, you can use magic strings, such as
+        `"$$IPv4{<ip address>}"` for IPv4 addresses. The full spec with examples is in the
+        [magic strings module][dp3.database.magic].
 
         Generic and fulltext filters are merged - fulltext overrides conflicting keys.
 
@@ -1000,7 +1004,10 @@ def get_latest_snapshots(
             generic_filter = {}
 
         # Create base of query
-        query = generic_filter
+        try:
+            query = search_and_replace(generic_filter)
+        except ValueError as e:
+            raise DatabaseError(f"Invalid generic filter: {str(e)}") from e
         query["latest"] = True
 
         # Process fulltext filters
@@ -1029,7 +1036,7 @@ def get_latest_snapshots(
                 [("_id", pymongo.ASCENDING)]
             ), snapshot_col.count_documents(query)
         except OperationFailure as e:
-            raise DatabaseError("Invalid query") from e
+            raise DatabaseError(f"Query is invalid: {e}") from e
 
     def _snapshot_bucket_id(self, etype: str, eid: Any, ctime: datetime) -> Union[str, Binary]:
         entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
diff --git a/dp3/database/magic.py b/dp3/database/magic.py
new file mode 100644
index 00000000..e0ba99b1
--- /dev/null
+++ b/dp3/database/magic.py
@@ -0,0 +1,167 @@
+"""
+Magic strings are used to convert non-JSON-native types to their respective objects.
+
+For querying non-JSON-native types, you can use the following magic strings:
+
+- `"$$IPv4{<ip address>}"` - converts to IPv4Address object
+- `"$$IPv6{<ip address>}"` - converts to IPv6Address object
+- `"$$int{<value>}"` - may be necessary for filtering when `eid` data type is int
+- `"$$Date{<YYYY-mm-ddTHH:MM:ssZ>}"` - converts specified UTC date to UTC datetime object
+- `"$$DateTs{<POSIX timestamp>}"` - converts POSIX timestamp (int/float) to UTC datetime object
+- `"$$MAC{<mac address>}"` - converts to MACAddress object
+
+To query an IP prefix, use the following magic strings:
+
+- `"$$IPv4Prefix{<ip address>/<prefix length>}"` - matches prefix
+- `"$$IPv6Prefix{<ip address>/<prefix length>}"` - matches prefix
+
+To query a binary `_id`s of non-string snapshot buckets,
+use the following magic string:
+
+- `"$$Binary_ID{<EID object | Valid magic string>}"`
+
+    - converts to filter the exact EID object snapshots, only EID valid types are supported
+
+There are no attribute name checks (may be added in the future).
+
+Generic filter examples:
+
+- `{"attr1": {"$gt": 10}, "attr2": {"$lt": 20}}`
+- `{"ip_attr": "$$IPv4{127.0.0.1}"}` - converts to IPv4Address object, exact match
+- `{"ip_attr": "$$IPv4Prefix{127.0.0.0/24}"}`
+    - converts to `{"ip_attr": {"$gte": "$$IPv4{127.0.0.0}",
+      "$lte": "$$IPv4{127.0.0.255}"}}`
+
+- `{"ip_attr": "$$IPv6{::1}"}` - converts to IPv6Address object, exact match
+- `{"ip_attr": "$$IPv6Prefix{::1/64}"}`
+    - converts to `{"ip_attr": {"$gte": "$$IPv6{::1}",
+      "$lte": "$$IPv6{::1:ffff:ffff:ffff:ffff}"}}`
+
+- `{"_id": "$$Binary_ID{$$IPv4{127.0.0.1}}"}`
+    - converts to `{"_id": {"$gte": Binary(<IP bytes + min timestamp>),
+      "$lt": Binary(<IP bytes + max timestamp>)}}`
+
+- `{"id": "$$Binary_ID{$$IPv4Prefix{127.0.0.0/24}}"}`
+    - converts to `{"_id": {"$gte": Binary(<IP 127.0.0.0 bytes + min timestamp>),
+      "$lte": Binary(<IP 127.0.0.255 bytes + max timestamp>)}}`
+
+- `{"_time_created": {"$gte": "$$Date{2024-11-07T00:00:00Z}"}}`
+    - converts to `{"_time_created": datetime(2024, 11, 7, 0, 0, 0, tzinfo=timezone.utc)}`
+
+- `{"_time_created": {"$gte": "$$DateTs{1609459200}"}}`
+    - converts to `{"_time_created": datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc)}`
+
+- `{"attr": "$$MAC{00:11:22:33:44:55}"}` - converts to MACAddress object, exact match
+- `{"_id": "$$Binary_ID{$$MAC{Ab-cD-Ef-12-34-56}}"}`
+    - converts to `{"_id": {"$gte": Binary(<MAC bytes + min timestamp>),
+        "$lt": Binary(<MAC bytes + max timestamp>)}}`
+"""
+
+import re
+from datetime import datetime, timezone
+from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
+from typing import Any, Union
+
+from bson import Binary
+
+from dp3.common.mac_address import MACAddress
+from dp3.common.utils import int2bytes
+
+
+def _pack_binary_snapshot_bucket_id(eid: bytes, ts_int: int) -> Binary:
+    return Binary(b"".join([eid, ts_int.to_bytes(8, "big", signed=True)]))
+
+
+def _binary_snapshot_bucket_range(eid: bytes) -> dict[str, Binary]:
+    return {
+        "$gte": _pack_binary_snapshot_bucket_id(eid, 0),
+        "$lt": _pack_binary_snapshot_bucket_id(eid, -1),
+    }
+
+
+def _binary_id_filter(value: Any, _) -> dict[str, Binary]:
+    """Converts the value to a binary ID filter."""
+    if isinstance(value, str):
+        match = magic_regex.match(value)
+        if match:
+            magic_type, value = match.groups()
+            value = magic_string_replacements[magic_type](value, True)
+
+    if isinstance(value, (IPv4Address, IPv6Address, MACAddress)):
+        return _binary_snapshot_bucket_range(value.packed)
+    if isinstance(value, (IPv4Network, IPv6Network)):
+        return {
+            "$gte": _pack_binary_snapshot_bucket_id(value[0].packed, 0),
+            "$lte": _pack_binary_snapshot_bucket_id(value[-1].packed, -1),
+        }
+    if isinstance(value, int):
+        return _binary_snapshot_bucket_range(int2bytes(value))
+
+    raise ValueError(f"Unsupported value type {type(value)}: {value}")
+
+
+def _parse_ipv4_network(
+    value: str, in_id_filter: bool
+) -> Union[IPv4Network, dict[str, IPv4Address]]:
+    ip = IPv4Network(value)
+    if in_id_filter:
+        return ip
+    return {"$gte": ip[0], "$lte": ip[-1]}
+
+
+def _parse_ipv6_network(
+    value: str, in_id_filter: bool
+) -> Union[IPv6Network, dict[str, IPv6Address]]:
+    ip = IPv6Network(value)
+    if in_id_filter:
+        return ip
+    return {"$gte": ip[0], "$lte": ip[-1]}
+
+
+def _parse_mac_address(value: str, _) -> MACAddress:
+    return MACAddress(value)
+
+
+def _parse_date_ts(value: Union[int, float], _) -> datetime:
+    return datetime.fromtimestamp(float(value), timezone.utc)
+
+
+def _parse_date_string(value: str, _) -> datetime:
+    return datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
+
+
+magic_string_replacements = {
+    "int": lambda x, _: int(x),
+    "IPv4": lambda x, _: IPv4Address(x),
+    "IPv6": lambda x, _: IPv6Address(x),
+    "Date": _parse_date_string,
+    "DateTs": _parse_date_ts,
+    "IPv4Prefix": _parse_ipv4_network,
+    "IPv6Prefix": _parse_ipv6_network,
+    "MAC": _parse_mac_address,
+    "Binary_ID": _binary_id_filter,
+}
+
+magic_regex = re.compile(r"\$\$(" + "|".join(magic_string_replacements) + r")\{(.+?)}$")
+
+
+def search_and_replace(query: dict[str, Any]) -> dict[str, Any]:
+    """Search and replace all occurrences of magic strings in the query.
+
+    A helper function to
+    [`get_latest_snapshots`][dp3.database.database.EntityDatabase.get_latest_snapshots]
+    """
+    if isinstance(query, dict):
+        for key, value in query.items():
+            if isinstance(value, (dict, list)):
+                search_and_replace(value)
+            elif isinstance(value, str):
+                match = magic_regex.match(value)
+                if match:
+                    magic_type, magic_value = match.groups()
+                    if magic_type in magic_string_replacements:
+                        query[key] = magic_string_replacements[magic_type](magic_value, False)
+    if isinstance(query, list):
+        for item in query:
+            search_and_replace(item)
+    return query
diff --git a/tests/test_common/test_magic.py b/tests/test_common/test_magic.py
new file mode 100644
index 00000000..0c1edb3d
--- /dev/null
+++ b/tests/test_common/test_magic.py
@@ -0,0 +1,132 @@
+"""Test the search & replace functionality for snapshot generic filter endpoint"""
+
+import unittest
+from datetime import datetime, timezone
+from ipaddress import IPv4Address, IPv6Address
+
+from bson import Binary
+from bson.binary import BINARY_SUBTYPE
+
+from dp3.common.mac_address import MACAddress
+from dp3.database.magic import search_and_replace
+
+
+class TestSearchAndReplace(unittest.TestCase):
+    def test_no_match(self):
+        query = {"ip_attr": "normal_value"}
+        expected = {"ip_attr": "normal_value"}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_ipv4(self):
+        query = {"ip_attr": "$$IPv4{127.0.0.1}"}
+        expected = {"ip_attr": IPv4Address("127.0.0.1")}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_ipv6(self):
+        query = {"ip_attr": "$$IPv6{::1}"}
+        expected = {"ip_attr": IPv6Address("::1")}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_int(self):
+        query = {"int_attr": "$$int{42}"}
+        expected = {"int_attr": 42}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_date(self):
+        query = {"date_attr": "$$Date{2021-01-01T00:00:00Z}"}
+        expected = {"date_attr": datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc)}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_date_ts(self):
+        query = {"date_attr": "$$DateTs{1609459200}"}
+        expected = {"date_attr": datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc)}
+        self.assertEqual(search_and_replace(query), expected)
+
+        # Test with float value
+        query = {"date_attr": "$$DateTs{1609459200.5}"}
+        expected = {"date_attr": datetime(2021, 1, 1, 0, 0, 0, 500000, tzinfo=timezone.utc)}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_ipv4_prefix(self):
+        query = {"ip_attr": "$$IPv4Prefix{127.0.0.0/24}"}
+        expected = {
+            "ip_attr": {"$gte": IPv4Address("127.0.0.0"), "$lte": IPv4Address("127.0.0.255")}
+        }
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_ipv6_prefix(self):
+        query = {"ip_attr": "$$IPv6Prefix{fe80::/64}"}
+        expected = {
+            "ip_attr": {
+                "$gte": IPv6Address("fe80::"),
+                "$lte": IPv6Address("fe80::ffff:ffff:ffff:ffff"),
+            }
+        }
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_binary_id(self):
+        query = {"_id": "$$Binary_ID{$$IPv4{127.0.0.1}}"}
+        expected = {
+            "_id": {
+                "$gte": Binary(b"\x7F\x00\x00\x01" + (b"\x00" * 8), subtype=BINARY_SUBTYPE),
+                "$lt": Binary(b"\x7F\x00\x00\x01" + (b"\xFF" * 8), subtype=BINARY_SUBTYPE),
+            }
+        }
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_mac(self):
+        query = {"mac_attr": "$$MAC{00:11:22:33:44:55}"}
+        expected = {"mac_attr": MACAddress("00:11:22:33:44:55")}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_replace_binary_id_mac(self):
+        query = {"_id": "$$Binary_ID{$$MAC{00:11:22:33:44:55}}"}
+        expected = {
+            "_id": {
+                "$gte": Binary(b"\x00\x11\x22\x33\x44\x55" + (b"\x00" * 8), subtype=BINARY_SUBTYPE),
+                "$lt": Binary(b"\x00\x11\x22\x33\x44\x55" + (b"\xFF" * 8), subtype=BINARY_SUBTYPE),
+            }
+        }
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_nested_query(self):
+        query = {"outer": {"inner": "$$IPv4{127.0.0.1}"}}
+        expected = {"outer": {"inner": IPv4Address("127.0.0.1")}}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_nested_query_list(self):
+        query = {"outer": [{"inner": "$$IPv4{192.168.0.1}"}]}
+        expected = {"outer": [{"inner": IPv4Address("192.168.0.1")}]}
+        self.assertEqual(search_and_replace(query), expected)
+
+    def test_invalid_type(self):
+        query = {"invalid": "$$invalid{42}"}
+        query_copy = query.copy()
+        self.assertEqual(search_and_replace(query), query_copy)
+
+    def test_invalid_sub_type(self):
+        invalid_queries = [
+            {"invalid": "$$IPv4{invalid}"},
+            {"invalid": "$$IPv4Prefix{invalid}"},
+            {"invalid": "$$IPv6{invalid}"},
+            {"invalid": "$$IPv6Prefix{invalid}"},
+            {"invalid": "$$int{invalid}"},
+            {"invalid": "$$Date{24-12-2021}"},
+            {"invalid": "$$DateTs{invalid}"},
+            {"invalid": "$$MAC{invalid}"},
+            {"invalid": "$$Binary_ID{invalid}"},
+            {"invalid": "$$Binary_ID{$$IPv4{invalid}}"},
+        ]
+        for query in invalid_queries:
+            with self.assertRaises(ValueError):
+                search_and_replace(query)
+
+    def test_empty_binary_id(self):
+        query = {"_id": "$$Binary_ID{}"}
+        query_copy = query.copy()
+        # Search and replace should not match
+        self.assertEqual(search_and_replace(query), query_copy)
+
+
+if __name__ == "__main__":
+    unittest.main()

From cb8befbbc521c2152f81ee497484a5aa9dd7cc5a Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Thu, 7 Nov 2024 16:29:46 +0100
Subject: [PATCH 06/22] DB: Separate modules for handling snapshot collection
 logic.

---
 dp3/database/config.py     |  47 +++
 dp3/database/encodings.py  |  41 ++
 dp3/database/exceptions.py |   6 +
 dp3/database/snapshots.py  | 829 +++++++++++++++++++++++++++++++++++++
 4 files changed, 923 insertions(+)
 create mode 100644 dp3/database/config.py
 create mode 100644 dp3/database/encodings.py
 create mode 100644 dp3/database/exceptions.py
 create mode 100644 dp3/database/snapshots.py

diff --git a/dp3/database/config.py b/dp3/database/config.py
new file mode 100644
index 00000000..3137736e
--- /dev/null
+++ b/dp3/database/config.py
@@ -0,0 +1,47 @@
+import urllib
+from typing import Literal, Union
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class MongoHostConfig(BaseModel, extra="forbid"):
+    """MongoDB host."""
+
+    address: str = "localhost"
+    port: int = 27017
+
+
+class MongoStandaloneConfig(BaseModel, extra="forbid"):
+    """MongoDB standalone configuration."""
+
+    mode: Literal["standalone"]
+    host: MongoHostConfig = MongoHostConfig()
+
+
+class MongoReplicaConfig(BaseModel, extra="forbid"):
+    """MongoDB replica set configuration."""
+
+    mode: Literal["replica"]
+    replica_set: str = "dp3"
+    hosts: list[MongoHostConfig]
+
+
+class StorageConfig(BaseModel, extra="forbid"):
+    """Storage configuration."""
+
+    snapshot_bucket_size: int = 32
+
+
+class MongoConfig(BaseModel, extra="forbid"):
+    """Database configuration."""
+
+    db_name: str = "dp3"
+    username: str = "dp3"
+    password: str = "dp3"
+    connection: Union[MongoStandaloneConfig, MongoReplicaConfig] = Field(..., discriminator="mode")
+    storage: StorageConfig = StorageConfig()
+
+    @field_validator("username", "password")
+    @classmethod
+    def url_safety(cls, v):
+        return urllib.parse.quote_plus(v)
diff --git a/dp3/database/encodings.py b/dp3/database/encodings.py
new file mode 100644
index 00000000..36e79bbd
--- /dev/null
+++ b/dp3/database/encodings.py
@@ -0,0 +1,41 @@
+from ipaddress import IPv4Address, IPv6Address
+
+from bson import Binary
+from bson.binary import USER_DEFINED_SUBTYPE
+from bson.codec_options import CodecOptions, TypeDecoder, TypeRegistry
+
+from dp3.common.mac_address import MACAddress
+
+BSON_OBJECT_TOO_LARGE = 10334
+
+BSON_IPV4_SUBTYPE = USER_DEFINED_SUBTYPE + 1
+BSON_IPV6_SUBTYPE = BSON_IPV4_SUBTYPE + 1
+BSON_MAC_SUBTYPE = BSON_IPV6_SUBTYPE + 1
+
+
+def fallback_encoder(value):
+    if isinstance(value, IPv4Address):
+        return Binary(value.packed, BSON_IPV4_SUBTYPE)
+    if isinstance(value, IPv6Address):
+        return Binary(value.packed, BSON_IPV6_SUBTYPE)
+    if isinstance(value, MACAddress):
+        return Binary(value.mac, BSON_MAC_SUBTYPE)
+    return value
+
+
+class DP3BinaryDecoder(TypeDecoder):
+    bson_type = Binary
+
+    def transform_bson(self, value):
+        if value.subtype == BSON_IPV4_SUBTYPE:
+            return IPv4Address(value)
+        if value.subtype == BSON_IPV6_SUBTYPE:
+            return IPv6Address(value)
+        if value.subtype == BSON_MAC_SUBTYPE:
+            return MACAddress(value)
+        return value
+
+
+def get_codec_options():
+    tr = TypeRegistry([DP3BinaryDecoder()], fallback_encoder=fallback_encoder)
+    return CodecOptions(type_registry=tr)
diff --git a/dp3/database/exceptions.py b/dp3/database/exceptions.py
new file mode 100644
index 00000000..7d4d8694
--- /dev/null
+++ b/dp3/database/exceptions.py
@@ -0,0 +1,6 @@
+class DatabaseError(Exception):
+    pass
+
+
+class SnapshotCollectionError(DatabaseError):
+    pass
diff --git a/dp3/database/snapshots.py b/dp3/database/snapshots.py
new file mode 100644
index 00000000..eabb9118
--- /dev/null
+++ b/dp3/database/snapshots.py
@@ -0,0 +1,829 @@
+import abc
+import logging
+import re
+from abc import ABC
+from collections import defaultdict
+from collections.abc import Iterable
+from datetime import datetime, timedelta
+from ipaddress import IPv4Address, IPv6Address
+from typing import Any, Optional, Union
+
+import pymongo
+from bson import Binary
+from pymongo import UpdateMany, UpdateOne
+from pymongo.collection import Collection
+from pymongo.command_cursor import CommandCursor
+from pymongo.cursor import Cursor
+from pymongo.database import Database
+from pymongo.errors import BulkWriteError, DocumentTooLarge, OperationFailure, WriteError
+
+from dp3.common.attrspec import AttrSpecType, AttrType
+from dp3.common.config import ModelSpec
+from dp3.common.mac_address import MACAddress
+from dp3.common.utils import int2bytes
+from dp3.database.config import MongoConfig
+from dp3.database.encodings import BSON_OBJECT_TOO_LARGE, get_codec_options
+from dp3.database.exceptions import SnapshotCollectionError
+from dp3.database.magic import search_and_replace
+
+
+class TypedSnapshotCollection(abc.ABC):
+    """Snapshot collection handler with eid type awareness."""
+
+    def __init__(
+        self,
+        db: Database,
+        entity_type: str,
+        db_config: MongoConfig,
+        model_spec: ModelSpec,
+        snapshots_config: dict,
+    ):
+        self._db = db.with_options(codec_options=get_codec_options())
+
+        if entity_type not in model_spec.entities:
+            raise ValueError(f"Entity type '{entity_type}' not found in model spec")
+        self.entity_type = entity_type
+        self._col_name = f"{entity_type}#snapshots"
+        self._os_col_name = f"{entity_type}#oversized_snapshots"
+
+        self.attr_specs: dict[str, AttrSpecType] = model_spec.entity_attributes[entity_type]
+
+        self.log = logging.getLogger(f"EntityDatabase.SnapshotCollection[{entity_type}]")
+
+        self._normal_snapshot_eids = set()
+        self._oversized_snapshot_eids = set()
+        self._snapshot_bucket_size = db_config.storage.snapshot_bucket_size
+        self._bucket_delta = self._get_snapshot_bucket_delta(snapshots_config)
+
+    def _get_snapshot_bucket_delta(self, config) -> timedelta:
+        """Returns how long it takes to fill a snapshot bucket.
+
+        This depends on the frequency of snapshot creation and the bucket size.
+        """
+        creation_rate = config.get("creation_rate", {"minute": "*/30"})
+        if len(creation_rate) > 1:
+            raise ValueError("Only one snapshot creation rate is supported.")
+
+        time_to_sec = {
+            "second": 1,
+            "minute": 60,
+            "hour": 60 * 60,
+            "day": 60 * 60 * 24,
+        }
+
+        time_to_sec_shifted = {
+            "second": 60,
+            "minute": 60 * 60,
+            "hour": 60 * 60 * 24,
+            "day": 60 * 60 * 24 * 30,
+        }
+
+        for key, value in creation_rate.items():
+            if value.startswith("*/"):
+                seconds = time_to_sec.get(key)
+                if seconds is None:
+                    raise ValueError(f"Unsupported snapshot creation rate: {creation_rate}")
+
+                snapshot_interval = int(value[2:])
+                bucket_delta = seconds * snapshot_interval
+            else:
+                seconds = time_to_sec_shifted.get(key)
+                if seconds is None:
+                    raise ValueError(f"Unsupported snapshot creation rate: {creation_rate}")
+
+                count = len(value.split(","))
+                bucket_delta = seconds // count
+            break
+        else:
+            raise ValueError(f"Unsupported snapshot creation rate: {creation_rate}")
+
+        return timedelta(seconds=bucket_delta * self._snapshot_bucket_size)
+
+    def _col(self, **kwargs) -> Collection:
+        """Returns entity snapshots collection.
+
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        return self._db.get_collection(self._col_name, **kwargs)
+
+    def _os_col(self, **kwargs) -> Collection:
+        """Returns entity snapshots collection.
+
+        **kwargs: additional arguments for `pymongo.Database.get_collection`.
+        """
+        return self._db.get_collection(self._os_col_name, **kwargs)
+
+    @staticmethod
+    def _pack_binary_bucket_id(eid: bytes, ts_int: int) -> Binary:
+        return Binary(b"".join([eid, ts_int.to_bytes(8, "big", signed=True)]))
+
+    def _binary_bucket_range(self, eid: bytes) -> dict:
+        return {
+            "$gte": self._pack_binary_bucket_id(eid, 0),
+            "$lt": self._pack_binary_bucket_id(eid, -1),
+        }
+
+    @abc.abstractmethod
+    def _bucket_id(self, eid: Any, ctime: datetime) -> Union[str, Binary]:
+        """Returns `_id` for snapshot bucket document."""
+
+    @abc.abstractmethod
+    def _filter_from_bid(self, b_id: Union[bytes, str]) -> dict:
+        """Returns filter for snapshots with same eid as given bucket document _id.
+        Args:
+            b_id: the _id of the snapshot bucket, type depends on etype's data type
+        """
+
+    @abc.abstractmethod
+    def _filter_from_eid(self, eid: Any) -> dict:
+        """Returns filter for snapshots with given eid."""
+
+    @abc.abstractmethod
+    def _filter_from_eids(self, eids: Iterable[Any]) -> dict:
+        """Returns filter for snapshots with given eids."""
+
+    def get_latest_one(self, eid: Any) -> dict:
+        """Get latest snapshot of given eid.
+
+        If doesn't exist, returns {}.
+        """
+        return (
+            self._col().find_one(self._filter_from_eid(eid), {"last": 1}, sort=[("_id", -1)]) or {}
+        ).get("last", {})
+
+    def get_latest(
+        self,
+        fulltext_filters: Optional[dict[str, str]] = None,
+        generic_filter: Optional[dict[str, Any]] = None,
+    ) -> tuple[Cursor, int]:
+        """Get latest snapshots of given `etype`.
+
+        This method is useful for displaying data on web.
+
+        Returns only documents matching `generic_filter` and `fulltext_filters`
+        (dictionary attribute - fulltext filter).
+        Fulltext filters are interpreted as regular expressions.
+        Only string values may be filtered this way. There's no validation that queried attribute
+        can be fulltext filtered.
+        Only plain and observation attributes with string-based data types can be queried.
+        Array and set data types are supported as well as long as they are not multi value
+        at the same time.
+        If you need to filter EIDs, ensure the EID is string, then use attribute `eid`.
+        Otherwise, use generic filter.
+
+        Generic filter allows filtering using generic MongoDB query (including `$and`, `$or`,
+        `$lt`, etc.).
+        For querying non-JSON-native types, you can use magic strings, such as
+        `"$$IPv4{<ip address>}"` for IPv4 addresses. The full spec with examples is in the
+        [magic strings module][dp3.database.magic].
+
+        Generic and fulltext filters are merged - fulltext overrides conflicting keys.
+
+        Also returns total document count (after filtering).
+
+        May raise `SnapshotCollectionError` if query is invalid.
+        """
+        snapshot_col = self._col()
+
+        if not fulltext_filters:
+            fulltext_filters = {}
+
+        if not generic_filter:
+            generic_filter = {}
+
+        # Create base of query
+        try:
+            query = search_and_replace(generic_filter)
+        except ValueError as e:
+            raise SnapshotCollectionError(f"Invalid generic filter: {str(e)}") from e
+        query["latest"] = True
+
+        # Process fulltext filters
+        for attr in fulltext_filters:
+            fulltext_filter = {"$regex": fulltext_filters[attr], "$options": "i"}
+
+            # EID filter
+            if attr == "eid":
+                query["_id"] = fulltext_filter
+                continue
+
+            # Check if attribute exists
+            try:
+                attr_spec = self.attr_specs[attr]
+            except KeyError as e:
+                raise SnapshotCollectionError(
+                    f"Attribute '{attr}' in fulltext filter doesn't exist"
+                ) from e
+
+            # Correctly handle link<...> data type
+            if attr_spec.t in AttrType.PLAIN | AttrType.OBSERVATIONS and attr_spec.is_relation:
+                query["last." + attr + ".eid"] = fulltext_filter
+            else:
+                query["last." + attr] = fulltext_filter
+
+        try:
+            return snapshot_col.find(query, {"last": 1}).sort(
+                [("_id", pymongo.ASCENDING)]
+            ), snapshot_col.count_documents(query)
+        except OperationFailure as e:
+            raise SnapshotCollectionError(f"Query is invalid: {e}") from e
+
+    def get_by_eid(
+        self, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
+    ) -> Union[Cursor, CommandCursor]:
+        """Get all (or filtered) snapshots of given `eid`.
+
+        This method is useful for displaying `eid`'s history on web.
+
+        Args:
+            eid: id of entity, to which data-points correspond
+            t1: left value of time interval (inclusive)
+            t2: right value of time interval (inclusive)
+        """
+        snapshot_col = self._col()
+
+        # Find out if the snapshot is oversized
+        doc = (
+            snapshot_col.find(self._filter_from_eid(eid), {"oversized": 1})
+            .sort([("_id", -1)])
+            .limit(1)
+        )
+        doc = next(doc, None)
+        if doc and doc.get("oversized", False):
+            return self._get_oversized(eid, t1, t2)
+
+        query = {"_time_created": {}}
+        pipeline = [
+            {"$match": self._filter_from_eid(eid)},
+            {"$unwind": "$history"},
+            {"$replaceRoot": {"newRoot": "$history"}},
+        ]
+
+        # Filter by date
+        if t1:
+            query["_time_created"]["$gte"] = t1
+        if t2:
+            query["_time_created"]["$lte"] = t2
+
+        # Unset if empty
+        if query["_time_created"]:
+            pipeline.append({"$match": query})
+        pipeline.append({"$sort": {"_time_created": pymongo.ASCENDING}})
+        return snapshot_col.aggregate(pipeline)
+
+    def get_distinct_val_count(self, attr: str) -> dict[Any, int]:
+        """Counts occurrences of distinct values of given attribute in snapshots.
+
+        Returns dictionary mapping value -> count.
+
+        Works for all plain and observation data types except `dict` and `json`.
+        """
+        # Get attribute specification
+        try:
+            attr_spec = self.attr_specs[attr]
+        except KeyError as e:
+            raise SnapshotCollectionError(f"Attribute '{attr}' does not exist") from e
+
+        if attr_spec.t not in AttrType.PLAIN | AttrType.OBSERVATIONS:
+            raise SnapshotCollectionError(f"Attribute '{attr}' isn't plain or observations")
+
+        # Attribute data type must be primitive, array<T> or set<T>
+        if any(needle in attr_spec.data_type.root for needle in ("dict", "json")):
+            raise SnapshotCollectionError(
+                f"Data type '{attr_spec.data_type}' of attribute '{attr}' is not processable"
+            )
+
+        # Build aggregation query
+        attr_path = "$last." + attr
+        unwinding = []
+
+        # Unwind array-like and multi value attributes
+        # If attribute is multi value array, unwind twice
+        if "array" in attr_spec.data_type.root or "set" in attr_spec.data_type.root:
+            unwinding.append({"$unwind": attr_path})
+        if attr_spec.t == AttrType.OBSERVATIONS and attr_spec.multi_value:
+            unwinding.append({"$unwind": attr_path})
+
+        # Group
+        agg_query_group_id = attr_path
+        if "link" in attr_spec.data_type.root:
+            agg_query_group_id += ".eid"
+
+        agg_query = [
+            {"$match": {"latest": True}},
+            *unwinding,
+            {"$group": {"_id": agg_query_group_id, "count": {"$sum": 1}}},
+            {"$sort": {"_id": 1, "count": -1}},
+        ]
+        # Run aggregation
+        distinct_counts_cur = self._col().aggregate(agg_query)
+
+        distinct_counts = {x["_id"]: x["count"] for x in distinct_counts_cur}
+
+        if None in distinct_counts:
+            del distinct_counts[None]
+
+        return distinct_counts
+
+    def _get_oversized(
+        self, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
+    ) -> Cursor:
+        """Get all (or filtered) snapshots of given `eid` from oversized snapshots collection."""
+        snapshot_col = self._os_col()
+        query = {"eid": eid, "_time_created": {}}
+
+        # Filter by date
+        if t1:
+            query["_time_created"]["$gte"] = t1
+        if t2:
+            query["_time_created"]["$lte"] = t2
+
+        # Unset if empty
+        if not query["_time_created"]:
+            del query["_time_created"]
+
+        return snapshot_col.find(query, projection={"_id": False}).sort(
+            [("_time_created", pymongo.ASCENDING)]
+        )
+
+    def _migrate_to_oversized(self, eid: Any, snapshot: dict):
+        snapshot_col = self._col()
+        os_col = self._os_col()
+
+        try:
+            move_to_oversized = []
+            last_id = None
+            for doc in snapshot_col.find(self._filter_from_eid(eid)).sort({"_id": 1}):
+                move_to_oversized.extend(doc.get("history", []))
+                last_id = doc["_id"]
+            move_to_oversized.insert(0, snapshot)
+
+            os_col.insert_many(move_to_oversized)
+            snapshot_col.update_one(
+                {"_id": last_id},
+                {
+                    "$set": {"oversized": True, "last": snapshot, "count": 0},
+                    "$unset": {"history": ""},
+                },
+            )
+            snapshot_col.delete_many(self._filter_from_eid(eid) | {"oversized": False})
+        except Exception as e:
+            raise SnapshotCollectionError(
+                f"Update of snapshot {eid} failed: {e}, {snapshot}"
+            ) from e
+
+    def save_one(self, snapshot: dict, ctime: datetime):
+        """Saves snapshot to specified entity of current master document.
+
+        Will move snapshot to oversized snapshots if the maintained bucket is too large.
+        """
+        if "eid" not in snapshot:
+            self.log.error("Snapshot is missing 'eid' field: %s", snapshot)
+            return
+        eid = snapshot["eid"]
+        snapshot["_time_created"] = ctime
+
+        snapshot_col = self._col()
+        os_col = self._os_col()
+
+        # Find out if the snapshot is oversized
+        normal, oversized = self._get_state({eid})
+        if normal:
+            try:
+                res = snapshot_col.update_one(
+                    self._filter_from_eid(eid) | {"count": {"$lt": self._snapshot_bucket_size}},
+                    {
+                        "$set": {"last": snapshot},
+                        "$push": {"history": {"$each": [snapshot], "$position": 0}},
+                        "$inc": {"count": 1},
+                        "$setOnInsert": {
+                            "_id": self._bucket_id(eid, ctime),
+                            "_time_created": ctime,
+                            "oversized": False,
+                            "latest": True,
+                        },
+                    },
+                    upsert=True,
+                )
+
+                if res.upserted_id is not None:
+                    snapshot_col.update_many(
+                        self._filter_from_eid(eid)
+                        | {"latest": True, "count": self._snapshot_bucket_size},
+                        {"$unset": {"latest": 1}},
+                    )
+            except (WriteError, OperationFailure, DocumentTooLarge) as e:
+                if e.code != BSON_OBJECT_TOO_LARGE:
+                    raise e
+                # The snapshot is too large, move it to oversized snapshots
+                self.log.info(f"Snapshot of {eid} is too large: {e}, marking as oversized.")
+                self._migrate_to_oversized(eid, snapshot)
+                self._cache_snapshot_state(set(), normal)
+            except Exception as e:
+                raise SnapshotCollectionError(
+                    f"Insert of snapshot {eid} failed: {e}, {snapshot}"
+                ) from e
+            return
+        elif oversized:
+            # Snapshot is already marked as oversized
+            snapshot_col.update_one(self._filter_from_eid(eid), {"$set": {"last": snapshot}})
+            os_col.insert_one(snapshot)
+            return
+
+    def _get_state(self, eids: set[str]) -> tuple[set, set]:
+        """Get current state of snapshot of given `eid`."""
+        unknown = eids
+        normal = self._normal_snapshot_eids & unknown
+        oversized = self._oversized_snapshot_eids & unknown
+        unknown = unknown - normal - oversized
+
+        if not unknown:
+            return normal, oversized
+
+        snapshot_col = self._col()
+        new_oversized = set()
+        for doc in snapshot_col.find(
+            self._filter_from_eids(unknown) | {"oversized": True},
+            {"oversized": 1},
+        ):
+            eid = doc["_id"].rsplit("_#", maxsplit=1)[0]
+            new_oversized.add(eid)
+
+        unknown = unknown - new_oversized
+        self._normal_snapshot_eids |= unknown
+        self._oversized_snapshot_eids |= new_oversized
+
+        return normal | unknown, oversized | new_oversized
+
+    def _cache_snapshot_state(self, normal: set, oversized: set):
+        """Cache snapshot state for given `etype`."""
+        self._normal_snapshot_eids |= normal
+
+        self._normal_snapshot_eids -= oversized
+        self._oversized_snapshot_eids |= oversized
+
+    def save_many(self, snapshots: list[dict], ctime: datetime):
+        """
+        Saves a list of snapshots of current master documents.
+
+        All snapshots must belong to same entity type.
+
+        Will move snapshots to oversized snapshots if the maintained bucket is too large.
+        For better understanding, see `save()`.
+        """
+
+        for snapshot in snapshots:
+            snapshot["_time_created"] = ctime
+
+        snapshot_col = self._col()
+        os_col = self._os_col()
+
+        snapshots_by_eid = defaultdict(list)
+        for snapshot in snapshots:
+            if "eid" not in snapshot:
+                continue
+            if "_id" in snapshot:
+                del snapshot["_id"]
+            snapshots_by_eid[snapshot["eid"]].append(snapshot)
+
+        # Find out if any of the snapshots are oversized
+        normal, oversized = self._get_state(set(snapshots_by_eid.keys()))
+
+        upserts = []
+        update_originals: list[list[dict]] = []
+        oversized_inserts = []
+        oversized_updates = []
+
+        # A normal snapshot, shift the last snapshot to history and update last
+        for eid in normal:
+            upserts.append(
+                UpdateOne(
+                    self._filter_from_eid(eid) | {"count": {"$lt": self._snapshot_bucket_size}},
+                    {
+                        "$set": {"last": snapshots_by_eid[eid][-1]},
+                        "$push": {"history": {"$each": snapshots_by_eid[eid], "$position": 0}},
+                        "$inc": {"count": len(snapshots_by_eid[eid])},
+                        "$setOnInsert": {
+                            "_id": self._bucket_id(eid, ctime),
+                            "_time_created": ctime,
+                            "oversized": False,
+                            "latest": True,
+                        },
+                    },
+                    upsert=True,
+                )
+            )
+            update_originals.append(snapshots_by_eid[eid])
+
+        # Snapshot is already marked as oversized
+        for eid in oversized:
+            oversized_inserts.extend(snapshots_by_eid[eid])
+            oversized_updates.append(
+                UpdateOne(
+                    self._filter_from_eid(eid),
+                    {"$set": {"last": snapshots_by_eid[eid][-1]}},
+                )
+            )
+
+        new_oversized = set()
+
+        if upserts:
+            try:
+                res = snapshot_col.bulk_write(upserts, ordered=False)
+
+                # Unset latest snapshots if new snapshots were inserted
+                if res.upserted_count > 0:
+                    unset_latest_updates = []
+                    for upsert_id in res.upserted_ids.values():
+                        unset_latest_updates.append(
+                            UpdateMany(
+                                self._filter_from_bid(upsert_id)
+                                | {"latest": True, "count": self._snapshot_bucket_size},
+                                {"$unset": {"latest": 1}},
+                            )
+                        )
+                    up_res = snapshot_col.bulk_write(unset_latest_updates)
+                    if up_res.modified_count != res.upserted_count:
+                        self.log.info(
+                            "Upserted the first snapshot for %d entities.",
+                            res.upserted_count - up_res.modified_count,
+                        )
+
+                if res.modified_count + res.upserted_count != len(upserts):
+                    self.log.error(
+                        "Some snapshots were not updated, %s != %s",
+                        res.modified_count + res.upserted_count,
+                        len(upserts),
+                    )
+            except (BulkWriteError, OperationFailure) as e:
+                self.log.info("Update of snapshots failed, will retry with oversize.")
+                failed_indexes = [
+                    err["index"]
+                    for err in e.details["writeErrors"]
+                    if err["code"] == BSON_OBJECT_TOO_LARGE
+                ]
+                failed_snapshots = (update_originals[i] for i in failed_indexes)
+                for eid_snapshots in failed_snapshots:
+                    eid = eid_snapshots[0]["eid"]
+                    failed_snapshots = sorted(
+                        eid_snapshots, key=lambda s: s["_time_created"], reverse=True
+                    )
+                    self._migrate_to_oversized(eid, failed_snapshots[0])
+                    oversized_inserts.extend(failed_snapshots[1:])
+                    new_oversized.add(eid)
+
+                if any(err["code"] != BSON_OBJECT_TOO_LARGE for err in e.details["writeErrors"]):
+                    # Some other error occurred
+                    raise e
+            except Exception as e:
+                raise SnapshotCollectionError(f"Upsert of snapshots failed: {str(e)[:2048]}") from e
+
+        # Update the oversized snapshots
+        if oversized_inserts:
+            try:
+                if oversized_updates:
+                    snapshot_col.bulk_write(oversized_updates)
+                os_col.insert_many(oversized_inserts)
+            except Exception as e:
+                raise SnapshotCollectionError(f"Insert of snapshots failed: {str(e)[:2048]}") from e
+
+        # Cache the new state
+        self._cache_snapshot_state(set(), new_oversized)
+
+    def delete_old(self, t_old: datetime) -> int:
+        """Delete old snapshots.
+
+        Periodically called from HistoryManager.
+        """
+        deleted = 0
+        try:
+            res = self._col().delete_many(
+                {"_time_created": {"$lte": t_old - self._bucket_delta}},
+            )
+            deleted += res.deleted_count * self._snapshot_bucket_size
+            res = self._os_col().delete_many({"_time_created": {"$lt": t_old}})
+            deleted += res.deleted_count
+        except Exception as e:
+            raise SnapshotCollectionError(f"Delete of olds snapshots failed: {e}") from e
+        return deleted
+
+    def delete_eid(self, eid: Any):
+        """Delete all snapshots of `eid`."""
+        try:
+            res = self._col().delete_many(self._filter_from_eid(eid))
+            del_cnt = res.deleted_count * self._snapshot_bucket_size
+            self.log.debug("deleted %s snapshots of %s/%s.", del_cnt, self.entity_type, eid)
+
+            res = self._os_col().delete_many({"eid": eid})
+            self.log.debug(
+                "Deleted %s oversized snapshots of %s/%s.", res.deleted_count, self.entity_type, eid
+            )
+        except Exception as e:
+            raise SnapshotCollectionError(f"Delete of failed: {e}\n{eid}") from e
+
+    def delete_eids(self, eids: list[Any]):
+        """Delete all snapshots of `eids`."""
+        try:
+            res = self._col().delete_many(self._filter_from_eids(eids))
+            del_cnt = res.deleted_count * self._snapshot_bucket_size
+            self.log.debug("Deleted %s snapshots of %s (%s).", del_cnt, self.entity_type, len(eids))
+            res = self._os_col().delete_many({"eid": {"$in": eids}})
+            self.log.debug(
+                "Deleted %s oversized snapshots of %s (%s).",
+                res.deleted_count,
+                self.entity_type,
+                len(eids),
+            )
+        except Exception as e:
+            raise SnapshotCollectionError(f"Delete of snapshots failed: {e}\n{eids}") from e
+
+
+class StringEidSnapshots(TypedSnapshotCollection):
+    def _bucket_id(self, eid: str, ctime: datetime) -> str:
+        return f"{eid}_#{int(ctime.timestamp())}"
+
+    def _filter_from_eid(self, eid: str) -> dict:
+        return {"_id": {"$regex": f"^{re.escape(str(eid))}_#"}}
+
+    def _filter_from_eids(self, eids: Iterable[str]) -> dict:
+        return {"_id": {"$regex": "|".join([f"^{re.escape(eid)}_#" for eid in eids])}}
+
+    def _filter_from_bid(self, b_id: str) -> dict:
+        eid = b_id.rsplit("_#", maxsplit=1)[0]
+        return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
+
+
+class BinaryEidSnapshots(TypedSnapshotCollection, ABC):
+    @abc.abstractmethod
+    def _get_packed(self, eid: Any) -> bytes:
+        """Return the packed bytes representation of the given eid."""
+
+    def _filter_from_bid(self, b_id: bytes) -> dict:
+        eid_bytes = b_id[:-8]
+        return {"_id": self._binary_bucket_range(eid_bytes)}
+
+    def _bucket_id(self, eid: Any, ctime: datetime) -> Binary:
+        ts = int(ctime.timestamp())
+        return self._pack_binary_bucket_id(self._get_packed(eid), ts)
+
+    def _filter_from_eid(self, eid: Any) -> dict:
+        return {"_id": self._binary_bucket_range(self._get_packed(eid))}
+
+    def _filter_from_eids(self, eids: Iterable[Any]) -> dict:
+        return {"$or": [self._filter_from_eid(eid) for eid in eids]}
+
+
+class IntEidSnapshots(BinaryEidSnapshots):
+    def _get_packed(self, eid: int) -> bytes:
+        return int2bytes(eid)
+
+
+class IPv4EidSnapshots(BinaryEidSnapshots):
+    def _get_packed(self, eid: IPv4Address) -> bytes:
+        return eid.packed
+
+
+class IPv6EidSnapshots(BinaryEidSnapshots):
+    def _get_packed(self, eid: IPv6Address) -> bytes:
+        return eid.packed
+
+
+class MACAddressEidSnapshots(BinaryEidSnapshots):
+    def _get_packed(self, eid: MACAddress) -> bytes:
+        return eid.packed
+
+
+entity_type2collection = {
+    "string": StringEidSnapshots,
+    "int": IntEidSnapshots,
+    "ipv4": IPv4EidSnapshots,
+    "ipv6": IPv6EidSnapshots,
+    "mac": MACAddressEidSnapshots,
+}
+
+
+class SnapshotCollectionContainer:
+    """Container for all required snapshot collections, exposing the public interface."""
+
+    def __init__(
+        self, db: Database, db_config: MongoConfig, model_spec: ModelSpec, snapshots_config: dict
+    ):
+        self._snapshot_collections = {}
+        for entity_type, entity_spec in model_spec.entities.items():
+            eid_type_name = entity_spec.id_data_type.root
+            typed_collection: type[TypedSnapshotCollection] = entity_type2collection[eid_type_name]
+            self._snapshot_collections[entity_type] = typed_collection(
+                db, entity_type, db_config, model_spec, snapshots_config
+            )
+
+        self.log = logging.getLogger("EntityDatabase.SnapshotCollections")
+
+    def __getitem__(self, entity_type: str) -> TypedSnapshotCollection:
+        if entity_type not in self._snapshot_collections:
+            raise KeyError(f"Entity type '{entity_type}' not found in model spec")
+        return self._snapshot_collections[entity_type]
+
+    def save_one(self, entity_type: str, snapshot: dict, ctime: datetime) -> None:
+        """Save snapshot to specified entity of current master document."""
+        return self[entity_type].save_one(snapshot, ctime)
+
+    def save_many(self, entity_type: str, snapshots: list[dict], ctime: datetime) -> None:
+        """
+        Saves a list of snapshots of current master documents.
+
+        All snapshots must belong to same entity type.
+        """
+        return self[entity_type].save_many(snapshots, ctime)
+
+    def get_latest_one(self, entity_type: str, eid: Any) -> dict:
+        """Get latest snapshot of given eid.
+
+        If doesn't exist, returns {}.
+        """
+        return self[entity_type].get_latest_one(eid)
+
+    def get_latest(
+        self,
+        entity_type: str,
+        fulltext_filters: Optional[dict[str, str]] = None,
+        generic_filter: Optional[dict[str, Any]] = None,
+    ) -> tuple[Cursor, int]:
+        """Get latest snapshots of given `etype`.
+
+        This method is useful for displaying data on web.
+
+        Returns only documents matching `generic_filter` and `fulltext_filters`
+        (dictionary attribute - fulltext filter).
+        Fulltext filters are interpreted as regular expressions.
+        Only string values may be filtered this way. There's no validation that queried attribute
+        can be fulltext filtered.
+        Only plain and observation attributes with string-based data types can be queried.
+        Array and set data types are supported as well as long as they are not multi value
+        at the same time.
+        If you need to filter EIDs, ensure the EID is string, then use attribute `eid`.
+        Otherwise, use generic filter.
+
+        Generic filter allows filtering using generic MongoDB query (including `$and`, `$or`,
+        `$lt`, etc.).
+        For querying non-JSON-native types, you can use magic strings, such as
+        `"$$IPv4{<ip address>}"` for IPv4 addresses. The full spec with examples is in the
+        [magic strings module][dp3.database.magic].
+
+        Generic and fulltext filters are merged - fulltext overrides conflicting keys.
+
+        Also returns total document count (after filtering).
+
+        May raise `SnapshotCollectionError` if query is invalid.
+        """
+        return self[entity_type].get_latest(fulltext_filters, generic_filter)
+
+    def get_by_eid(
+        self,
+        entity_type: str,
+        eid: Any,
+        t1: Optional[datetime] = None,
+        t2: Optional[datetime] = None,
+    ) -> Union[Cursor, CommandCursor]:
+        """Get all (or filtered) snapshots of given `eid`.
+
+        This method is useful for displaying `eid`'s history on web.
+
+        Args:
+            entity_type: name of entity type
+            eid: id of entity, to which data-points correspond
+            t1: left value of time interval (inclusive)
+            t2: right value of time interval (inclusive)
+        """
+        return self[entity_type].get_by_eid(eid, t1, t2)
+
+    def get_distinct_val_count(self, entity_type: str, attr: str) -> dict[Any, int]:
+        """Counts occurrences of distinct values of given attribute in snapshots.
+
+        Returns dictionary mapping value -> count.
+
+        Works for all plain and observation data types except `dict` and `json`.
+        """
+        return self[entity_type].get_distinct_val_count(attr)
+
+    def delete_old(self, t_old: datetime) -> int:
+        """Delete old snapshots, may raise `SnapshotCollectionError`.
+
+        Periodically called from HistoryManager.
+        Returns:
+             number of deleted snapshots.
+        """
+        deleted_total = 0
+        for collection in self._snapshot_collections.values():
+            try:
+                deleted_total += collection.delete_old(t_old)
+            except Exception as e:
+                raise SnapshotCollectionError(f"Delete of old snapshots failed: {e}") from e
+        return deleted_total
+
+    def delete_eid(self, entity_type: str, eid: Any) -> int:
+        """Delete snapshots of given `eids`."""
+        return self[entity_type].delete_eid(eid)
+
+    def delete_eids(self, entity_type: str, eids: list[Any]) -> int:
+        """Delete snapshots of given `eids`."""
+        return self[entity_type].delete_eids(eids)

From 97e17e87f961a7450b0362c12a2750bfcf8bd1f4 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Fri, 8 Nov 2024 15:49:55 +0100
Subject: [PATCH 07/22] DB: Use modularized snapshots.

---
 dp3/api/routers/entity.py                 |   6 +-
 dp3/database/database.py                  | 803 +---------------------
 dp3/history_management/history_manager.py |  10 +-
 dp3/snapshots/snapshooter.py              |   6 +-
 4 files changed, 39 insertions(+), 786 deletions(-)

diff --git a/dp3/api/routers/entity.py b/dp3/api/routers/entity.py
index be39ffce..237bb3d5 100644
--- a/dp3/api/routers/entity.py
+++ b/dp3/api/routers/entity.py
@@ -59,7 +59,7 @@ def get_eid_snapshots_handler(
     etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
 ) -> list[dict[str, Any]]:
     """Handler for getting snapshots of EID"""
-    snapshots = list(DB.get_snapshots(etype, eid, t1=date_from, t2=date_to))
+    snapshots = list(DB.snapshots.get_by_eid(etype, eid, t1=date_from, t2=date_to))
 
     return snapshots
 
@@ -180,7 +180,7 @@ async def list_entity_type_eids(
         fulltext_filters["eid"] = eid_filter
 
     try:
-        cursor, total_count = DB.get_latest_snapshots(etype, fulltext_filters, generic_filter)
+        cursor, total_count = DB.snapshots.get_latest(etype, fulltext_filters, generic_filter)
         cursor_page = cursor.skip(skip).limit(limit)
     except DatabaseError as e:
         raise HTTPException(status_code=400, detail=str(e)) from e
@@ -312,7 +312,7 @@ async def get_distinct_attribute_values(etype: str, attr: str) -> dict[JsonVal,
     Works for all plain and observation data types except `dict` and `json`.
     """
     try:
-        return DB.get_distinct_val_count(etype, attr)
+        return DB.snapshots.get_distinct_val_count(etype, attr)
     except DatabaseError as e:
         raise HTTPException(status_code=400, detail=str(e)) from e
 
diff --git a/dp3/database/database.py b/dp3/database/database.py
index 69df79bf..688a1faa 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -1,118 +1,34 @@
 import inspect
 import logging
-import re
 import threading
 import time
-import urllib
 from collections import defaultdict
-from collections.abc import Generator, Iterable, Iterator
-from datetime import datetime, timedelta
-from ipaddress import IPv4Address, IPv6Address
-from typing import Any, Callable, Literal, Optional, Union
+from collections.abc import Generator, Iterator
+from datetime import datetime
+from typing import Any, Callable, Optional
 
-import bson
 import pymongo
-from bson.binary import USER_DEFINED_SUBTYPE, Binary
-from bson.codec_options import TypeDecoder
 from event_count_logger import DummyEventGroup
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel
 from pymongo import ReplaceOne, UpdateMany, UpdateOne, WriteConcern
 from pymongo.collection import Collection
 from pymongo.command_cursor import CommandCursor
 from pymongo.cursor import Cursor
-from pymongo.errors import BulkWriteError, DocumentTooLarge, OperationFailure, WriteError
+from pymongo.database import Database
+from pymongo.errors import OperationFailure
 from pymongo.results import DeleteResult, UpdateResult
 
 from dp3.common.attrspec import AttrType, timeseries_types
 from dp3.common.config import HierarchicalDict, ModelSpec
 from dp3.common.datapoint import DataPointBase
-from dp3.common.mac_address import MACAddress
 from dp3.common.scheduler import Scheduler
 from dp3.common.task import HASH
 from dp3.common.types import EventGroupType
-from dp3.common.utils import int2bytes
-from dp3.database.magic import search_and_replace
+from dp3.database.config import MongoConfig, MongoReplicaConfig, MongoStandaloneConfig
+from dp3.database.encodings import get_codec_options
+from dp3.database.exceptions import DatabaseError
 from dp3.database.schema_cleaner import SchemaCleaner
-
-BSON_OBJECT_TOO_LARGE = 10334
-
-BSON_IPV4_SUBTYPE = USER_DEFINED_SUBTYPE + 1
-BSON_IPV6_SUBTYPE = BSON_IPV4_SUBTYPE + 1
-BSON_MAC_SUBTYPE = BSON_IPV6_SUBTYPE + 1
-
-
-def fallback_encoder(value):
-    if isinstance(value, IPv4Address):
-        return Binary(value.packed, BSON_IPV4_SUBTYPE)
-    if isinstance(value, IPv6Address):
-        return Binary(value.packed, BSON_IPV6_SUBTYPE)
-    if isinstance(value, MACAddress):
-        return Binary(value.mac, BSON_MAC_SUBTYPE)
-    return value
-
-
-class DP3BinaryDecoder(TypeDecoder):
-    bson_type = Binary
-
-    def transform_bson(self, value):
-        if value.subtype == BSON_IPV4_SUBTYPE:
-            return IPv4Address(value)
-        if value.subtype == BSON_IPV6_SUBTYPE:
-            return IPv6Address(value)
-        if value.subtype == BSON_MAC_SUBTYPE:
-            return MACAddress(value)
-        return value
-
-
-class DatabaseError(Exception):
-    pass
-
-
-class MissingTableError(DatabaseError):
-    pass
-
-
-class MongoHostConfig(BaseModel, extra="forbid"):
-    """MongoDB host."""
-
-    address: str = "localhost"
-    port: int = 27017
-
-
-class MongoStandaloneConfig(BaseModel, extra="forbid"):
-    """MongoDB standalone configuration."""
-
-    mode: Literal["standalone"]
-    host: MongoHostConfig = MongoHostConfig()
-
-
-class MongoReplicaConfig(BaseModel, extra="forbid"):
-    """MongoDB replica set configuration."""
-
-    mode: Literal["replica"]
-    replica_set: str = "dp3"
-    hosts: list[MongoHostConfig]
-
-
-class StorageConfig(BaseModel, extra="forbid"):
-    """Storage configuration."""
-
-    snapshot_bucket_size: int = 32
-
-
-class MongoConfig(BaseModel, extra="forbid"):
-    """Database configuration."""
-
-    db_name: str = "dp3"
-    username: str = "dp3"
-    password: str = "dp3"
-    connection: Union[MongoStandaloneConfig, MongoReplicaConfig] = Field(..., discriminator="mode")
-    storage: StorageConfig = StorageConfig()
-
-    @field_validator("username", "password")
-    @classmethod
-    def url_safety(cls, v):
-        return urllib.parse.quote_plus(v)
+from dp3.database.snapshots import SnapshotCollectionContainer
 
 
 def get_caller_id():
@@ -133,7 +49,7 @@ class EntityDatabase:
     Initializes database schema based on database configuration.
 
     Args:
-        db_conf: configuration of database connection (content of database.yml)
+        config: configuration of database connection (content of database.yml)
         model_spec: ModelSpec object, configuration of data model (entities and attributes)
         process_index: index of worker process - used for sharding metadata
         num_processes: number of worker processes
@@ -175,20 +91,17 @@ def __init__(
         self._num_processes = num_processes
         self._process_index = process_index
 
-        self._normal_snapshot_eids = defaultdict(set)
-        self._oversized_snapshot_eids = defaultdict(set)
-        self._snapshot_bucket_size = db_config.storage.snapshot_bucket_size
-        self._bucket_delta = self._get_snapshot_bucket_delta(config)
-
         # Init and switch to correct database
-        self._db = self._db[db_config.db_name]
-        type_registry = bson.codec_options.TypeRegistry(
-            [DP3BinaryDecoder()], fallback_encoder=fallback_encoder
+        codec_opts = get_codec_options()
+        self._db = Database(self._db, db_config.db_name, codec_options=codec_opts)
+
+        self.snapshots = SnapshotCollectionContainer(
+            self._db, db_config, model_spec, config.get("snapshots", {})
         )
-        self._codec_opts = bson.codec_options.CodecOptions(type_registry=type_registry)
+        self._snapshot_bucket_size = db_config.storage.snapshot_bucket_size
 
         if process_index == 0:
-            self._init_database_schema(db_config.db_name)
+            self._init_database_schema()
 
         self.schema_cleaner = SchemaCleaner(
             self._db, self.get_module_cache("Schema"), self._db_schema_config, config, self.log
@@ -211,50 +124,6 @@ def __init__(
 
         self.log.info("Database successfully initialized!")
 
-    def _get_snapshot_bucket_delta(self, config) -> timedelta:
-        """Returns how long it takes to fill a snapshot bucket.
-
-        This depends on the frequency of snapshot creation and the bucket size.
-        """
-        creation_rate = config.get("snapshots.creation_rate", {"minute": "*/30"})
-        if len(creation_rate) > 1:
-            raise ValueError("Only one snapshot creation rate is supported.")
-
-        time_to_sec = {
-            "second": 1,
-            "minute": 60,
-            "hour": 60 * 60,
-            "day": 60 * 60 * 24,
-        }
-
-        time_to_sec_shifted = {
-            "second": 60,
-            "minute": 60 * 60,
-            "hour": 60 * 60 * 24,
-            "day": 60 * 60 * 24 * 30,
-        }
-
-        for key, value in creation_rate.items():
-            if value.startswith("*/"):
-                seconds = time_to_sec.get(key)
-                if seconds is None:
-                    raise ValueError(f"Unsupported snapshot creation rate: {creation_rate}")
-
-                snapshot_interval = int(value[2:])
-                bucket_delta = seconds * snapshot_interval
-            else:
-                seconds = time_to_sec_shifted.get(key)
-                if seconds is None:
-                    raise ValueError(f"Unsupported snapshot creation rate: {creation_rate}")
-
-                count = len(value.split(","))
-                bucket_delta = seconds // count
-            break
-        else:
-            raise ValueError(f"Unsupported snapshot creation rate: {creation_rate}")
-
-        return timedelta(seconds=bucket_delta * self._snapshot_bucket_size)
-
     @staticmethod
     def connect(config: MongoConfig) -> pymongo.MongoClient:
         if isinstance(config.connection, MongoStandaloneConfig):
@@ -301,25 +170,21 @@ def _master_col(self, entity: str, **kwargs) -> Collection:
 
         **kwargs: additional arguments for `pymongo.Database.get_collection`.
         """
-        return self._db.get_collection(f"{entity}#master", codec_options=self._codec_opts, **kwargs)
+        return self._db.get_collection(f"{entity}#master", **kwargs)
 
     def _snapshot_col(self, entity: str, **kwargs) -> Collection:
         """Returns snapshots collection for `entity`.
 
         **kwargs: additional arguments for `pymongo.Database.get_collection`.
         """
-        return self._db.get_collection(
-            f"{entity}#snapshots", codec_options=self._codec_opts, **kwargs
-        )
+        return self._db.get_collection(f"{entity}#snapshots", **kwargs)
 
     def _oversized_snapshots_col(self, entity: str, **kwargs) -> Collection:
         """Returns oversized snapshots collection for `entity`.
 
         **kwargs: additional arguments for `pymongo.Database.get_collection`.
         """
-        return self._db.get_collection(
-            f"{entity}#snapshots_oversized", codec_options=self._codec_opts, **kwargs
-        )
+        return self._db.get_collection(f"{entity}#snapshots_oversized", **kwargs)
 
     @staticmethod
     def _raw_col_name(entity: str) -> str:
@@ -331,7 +196,7 @@ def _raw_col(self, entity: str, **kwargs) -> Collection:
 
         **kwargs: additional arguments for `pymongo.Database.get_collection`.
         """
-        return self._db.get_collection(f"{entity}#raw", codec_options=self._codec_opts, **kwargs)
+        return self._db.get_collection(f"{entity}#raw", **kwargs)
 
     @staticmethod
     def _get_new_archive_col_name(entity: str) -> str:
@@ -348,9 +213,9 @@ def _archive_cols(self, entity: str, **kwargs) -> Generator[Collection, None, No
         **kwargs: additional arguments for `pymongo.Database.get_collection`.
         """
         for col_name in self._archive_col_names(entity):
-            yield self._db.get_collection(col_name, codec_options=self._codec_opts, **kwargs)
+            yield self._db.get_collection(col_name, **kwargs)
 
-    def _init_database_schema(self, db_name) -> None:
+    def _init_database_schema(self) -> None:
         """Runs full check and update of database schema.
 
         Checks whether `db_name` database exists. If not, creates it.
@@ -686,8 +551,6 @@ def remove_expired_ttls(self, etype: str, expired_eid_ttls: dict[str, list[str]]
     def delete_eids(self, etype: str, eids: list[str]):
         """Delete master record and all snapshots of `etype`:`eids`."""
         master_col = self._master_col(etype)
-        snapshot_col = self._snapshot_col(etype)
-        os_snapshot_col = self._oversized_snapshots_col(etype)
         try:
             res = master_col.delete_many({"_id": {"$in": eids}})
             self.log.debug(
@@ -696,16 +559,8 @@ def delete_eids(self, etype: str, eids: list[str]):
             self.elog.log("record_removed", count=res.deleted_count)
         except Exception as e:
             raise DatabaseError(f"Delete of master record failed: {e}\n{eids}") from e
-        try:
-            res = snapshot_col.delete_many(self._snapshot_bucket_eids_filter(etype, eids))
-            del_cnt = res.deleted_count * self._snapshot_bucket_size
-            self.log.debug("Deleted %s snapshots of %s (%s).", del_cnt, etype, len(eids))
-            res = os_snapshot_col.delete_many({"eid": {"$in": eids}})
-            self.log.debug(
-                "Deleted %s oversized snapshots of %s (%s).", res.deleted_count, etype, len(eids)
-            )
-        except Exception as e:
-            raise DatabaseError(f"Delete of snapshots failed: {e}\n{eids}") from e
+        self.snapshots.delete_eids(etype, eids)
+
         for f in self._on_entity_delete_many:
             try:
                 f(etype, eids)
@@ -715,24 +570,14 @@ def delete_eids(self, etype: str, eids: list[str]):
     def delete_eid(self, etype: str, eid: Any):
         """Delete master record and all snapshots of `etype`:`eid`."""
         master_col = self._master_col(etype)
-        snapshot_col = self._snapshot_col(etype)
-        os_snapshot_col = self._oversized_snapshots_col(etype)
         try:
             master_col.delete_one({"_id": eid})
             self.log.debug("Deleted master record of %s/%s.", etype, eid)
             self.elog.log("record_removed")
         except Exception as e:
             raise DatabaseError(f"Delete of master record failed: {e}\n{eid}") from e
-        try:
-            res = snapshot_col.delete_many(self._snapshot_bucket_eid_filter(etype, eid))
-            del_cnt = res.deleted_count * self._snapshot_bucket_size
-            self.log.debug("deleted %s snapshots of %s/%s.", del_cnt, etype, eid)
-            res = os_snapshot_col.delete_many({"eid": eid})
-            self.log.debug(
-                "Deleted %s oversized snapshots of %s/%s.", res.deleted_count, etype, eid
-            )
-        except Exception as e:
-            raise DatabaseError(f"Delete of snapshots failed: {e}\n{eid}") from e
+        self.snapshots.delete_eid(etype, eid)
+
         for f in self._on_entity_delete_one:
             try:
                 f(etype, eid)
@@ -918,266 +763,6 @@ def get_worker_master_records(
             {"#hash": {"$mod": [worker_cnt, worker_index]}, **query_filter}, **kwargs
         )
 
-    def get_latest_snapshot(self, etype: str, eid: Any) -> dict:
-        """Get latest snapshot of given etype/eid.
-
-        If doesn't exist, returns {}.
-        """
-        # Check `etype`
-        self._assert_etype_exists(etype)
-
-        snapshot_col = self._snapshot_col(etype)
-        return (
-            snapshot_col.find_one(
-                self._snapshot_bucket_eid_filter(etype, eid), {"last": 1}, sort=[("_id", -1)]
-            )
-            or {}
-        ).get("last", {})
-
-    def _get_latest_snapshots_date(self) -> Optional[datetime]:
-        """Get date of newest snapshot set.
-
-        Queries snapshots metadata to find latest fully completed set of snapshots
-        and returns it's date.
-        If it doesn't exist, returns `None`.
-
-        Note: This doesn't take into account etypes, so it may be inaccurate for new entity types.
-        """
-        # Get latest fully completed snapshot metadata
-        lfcsm = self._db["#metadata"].find_one(
-            {
-                "#module": "SnapShooter",
-                "workers_finished": self._num_processes,
-                "linked_finished": True,
-            },
-            sort=[("#time_created", -1)],
-        )
-
-        if lfcsm is None:
-            return None
-
-        # Extract date
-        return lfcsm["#time_created"]
-
-    def get_latest_snapshots(
-        self,
-        etype: str,
-        fulltext_filters: Optional[dict[str, str]] = None,
-        generic_filter: Optional[dict[str, Any]] = None,
-    ) -> tuple[Cursor, int]:
-        """Get latest snapshots of given `etype`.
-
-        This method is useful for displaying data on web.
-
-        Returns only documents matching `generic_filter` and `fulltext_filters`
-        (dictionary attribute - fulltext filter).
-        Fulltext filters are interpreted as regular expressions.
-        Only string values may be filtered this way. There's no validation that queried attribute
-        can be fulltext filtered.
-        Only plain and observation attributes with string-based data types can be queried.
-        Array and set data types are supported as well as long as they are not multi value
-        at the same time.
-        If you need to filter EIDs, ensure the EID is string, then use attribute `eid`.
-        Otherwise, use generic filter.
-
-        Generic filter allows filtering using generic MongoDB query (including `$and`, `$or`,
-        `$lt`, etc.).
-        For querying non-JSON-native types, you can use magic strings, such as
-        `"$$IPv4{<ip address>}"` for IPv4 addresses. The full spec with examples is in the
-        [magic strings module][dp3.database.magic].
-
-        Generic and fulltext filters are merged - fulltext overrides conflicting keys.
-
-        Also returns total document count (after filtering).
-
-        May raise `DatabaseError` if query is invalid.
-        """
-        # Check `etype`
-        self._assert_etype_exists(etype)
-
-        snapshot_col = self._snapshot_col(etype)
-
-        if not fulltext_filters:
-            fulltext_filters = {}
-
-        if not generic_filter:
-            generic_filter = {}
-
-        # Create base of query
-        try:
-            query = search_and_replace(generic_filter)
-        except ValueError as e:
-            raise DatabaseError(f"Invalid generic filter: {str(e)}") from e
-        query["latest"] = True
-
-        # Process fulltext filters
-        for attr in fulltext_filters:
-            fulltext_filter = {"$regex": fulltext_filters[attr], "$options": "i"}
-
-            # EID filter
-            if attr == "eid":
-                query["_id"] = fulltext_filter
-                continue
-
-            # Check if attribute exists
-            try:
-                attr_spec = self._db_schema_config.attr(etype, attr)
-            except KeyError as e:
-                raise DatabaseError(f"Attribute '{attr}' in fulltext filter doesn't exist") from e
-
-            # Correctly handle link<...> data type
-            if attr_spec.t in AttrType.PLAIN | AttrType.OBSERVATIONS and attr_spec.is_relation:
-                query["last." + attr + ".eid"] = fulltext_filter
-            else:
-                query["last." + attr] = fulltext_filter
-
-        try:
-            return snapshot_col.find(query, {"last": 1}).sort(
-                [("_id", pymongo.ASCENDING)]
-            ), snapshot_col.count_documents(query)
-        except OperationFailure as e:
-            raise DatabaseError(f"Query is invalid: {e}") from e
-
-    def _snapshot_bucket_id(self, etype: str, eid: Any, ctime: datetime) -> Union[str, Binary]:
-        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
-        ts_int = int(ctime.timestamp())
-
-        if entity_data_type == "string":
-            return f"{eid}_#{ts_int}"
-        if entity_data_type == "int" and isinstance(eid, int):
-            return self._pack_binary_snapshot_bucket_id(int2bytes(eid), ts_int)
-        if entity_data_type == "ipv4" and isinstance(eid, IPv4Address):
-            return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
-        if entity_data_type == "ipv6" and isinstance(eid, IPv6Address):
-            return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
-        if entity_data_type == "mac" and isinstance(eid, MACAddress):
-            return self._pack_binary_snapshot_bucket_id(eid.packed, ts_int)
-        raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}', {eid}")
-
-    @staticmethod
-    def _pack_binary_snapshot_bucket_id(eid: bytes, ts_int: int) -> Binary:
-        return Binary(b"".join([eid, ts_int.to_bytes(8, "big", signed=True)]))
-
-    def _snapshot_bucket_eid_filter(self, etype: str, eid: Any) -> dict:
-        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
-
-        if entity_data_type == "string":
-            eid = str(eid)
-            return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
-        if entity_data_type == "int":
-            eid = int(eid)
-            return {"_id": self._binary_snapshot_bucket_range(int2bytes(eid))}
-        if entity_data_type == "ipv4":
-            eid = IPv4Address(eid)
-            return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
-        if entity_data_type == "ipv6":
-            eid = IPv6Address(eid)
-            return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
-        if entity_data_type == "mac":
-            eid = MACAddress(eid)
-            return {"_id": self._binary_snapshot_bucket_range(eid.packed)}
-        raise ValueError(
-            f"Unsupported data type '{entity_data_type}' for entity {etype} '{eid}' ({type(eid)})"
-        )
-
-    def _snapshot_eid_filter_from_bid(self, etype: str, b_id: Union[bytes, str]) -> dict:
-        """Returns filter for snapshots with same eid as given bucket document _id.
-        Args:
-            etype: entity type
-            b_id: the _id of the snapshot bucket, type depends on etype's data type
-        """
-        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
-
-        if entity_data_type == "string":
-            eid = b_id.rsplit("_#", maxsplit=1)[0]
-            return {"_id": {"$regex": f"^{re.escape(eid)}_#"}}
-        if entity_data_type in ["int", "ipv4", "ipv6", "mac"]:
-            eid_bytes = b_id[:-8]
-            return {"_id": self._binary_snapshot_bucket_range(eid_bytes)}
-        raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}'")
-
-    def _binary_snapshot_bucket_range(self, eid: bytes) -> dict:
-        return {
-            "$gte": self._pack_binary_snapshot_bucket_id(eid, 0),
-            "$lt": self._pack_binary_snapshot_bucket_id(eid, -1),
-        }
-
-    def _snapshot_bucket_eids_filter(self, etype: str, eids: Iterable[Any]) -> dict:
-        entity_data_type = self._db_schema_config.entities[etype].id_data_type.root
-        if entity_data_type == "string":
-            return {"_id": {"$regex": "|".join([f"^{re.escape(eid)}_#" for eid in eids])}}
-        if entity_data_type in ["int", "ipv4", "ipv6", "mac"]:
-            return {"$or": [self._snapshot_bucket_eid_filter(etype, eid) for eid in eids]}
-        raise ValueError(f"Unsupported data type '{entity_data_type}' for entity '{etype}'")
-
-    def get_snapshots(
-        self, etype: str, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
-    ) -> Union[Cursor, CommandCursor]:
-        """Get all (or filtered) snapshots of given `eid`.
-
-        This method is useful for displaying `eid`'s history on web.
-
-        Args:
-            etype: entity type
-            eid: id of entity, to which data-points correspond
-            t1: left value of time interval (inclusive)
-            t2: right value of time interval (inclusive)
-        """
-        # Check `etype`
-        self._assert_etype_exists(etype)
-
-        snapshot_col = self._snapshot_col(etype)
-
-        # Find out if the snapshot is oversized
-        doc = (
-            snapshot_col.find(self._snapshot_bucket_eid_filter(etype, eid), {"oversized": 1})
-            .sort([("_id", -1)])
-            .limit(1)
-        )
-        doc = next(doc, None)
-        if doc and doc.get("oversized", False):
-            return self._get_snapshots_oversized(etype, eid, t1, t2)
-
-        query = {"_time_created": {}}
-        pipeline = [
-            {"$match": self._snapshot_bucket_eid_filter(etype, eid)},
-            {"$unwind": "$history"},
-            {"$replaceRoot": {"newRoot": "$history"}},
-        ]
-
-        # Filter by date
-        if t1:
-            query["_time_created"]["$gte"] = t1
-        if t2:
-            query["_time_created"]["$lte"] = t2
-
-        # Unset if empty
-        if query["_time_created"]:
-            pipeline.append({"$match": query})
-        pipeline.append({"$sort": {"_time_created": pymongo.ASCENDING}})
-        return snapshot_col.aggregate(pipeline)
-
-    def _get_snapshots_oversized(
-        self, etype: str, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
-    ) -> Cursor:
-        """Get all (or filtered) snapshots of given `eid` from oversized snapshots collection."""
-        snapshot_col = self._oversized_snapshots_col(etype)
-        query = {"eid": eid, "_time_created": {}}
-
-        # Filter by date
-        if t1:
-            query["_time_created"]["$gte"] = t1
-        if t2:
-            query["_time_created"]["$lte"] = t2
-
-        # Unset if empty
-        if not query["_time_created"]:
-            del query["_time_created"]
-
-        return snapshot_col.find(query, projection={"_id": False}).sort(
-            [("_time_created", pymongo.ASCENDING)]
-        )
-
     def get_value_or_history(
         self,
         etype: str,
@@ -1208,7 +793,7 @@ def get_value_or_history(
                 self.get_master_record(etype, eid).get(attr_name, {}).get("v", None)
             )
         elif attr_spec.t == AttrType.OBSERVATIONS:
-            result["current_value"] = self.get_latest_snapshot(etype, eid).get(attr_name, None)
+            result["current_value"] = self.snapshots.get_latest_one(etype, eid).get(attr_name, None)
 
         # Add history
         if attr_spec.t == AttrType.OBSERVATIONS:
@@ -1226,259 +811,6 @@ def estimate_count_eids(self, etype: str) -> int:
         master_col = self._master_col(etype)
         return master_col.estimated_document_count({})
 
-    def _migrate_to_oversized_snapshot(self, etype: str, eid: Any, snapshot: dict):
-        snapshot_col = self._snapshot_col(etype)
-        os_col = self._oversized_snapshots_col(etype)
-
-        try:
-            move_to_oversized = []
-            last_id = None
-            for doc in snapshot_col.find(self._snapshot_bucket_eid_filter(etype, eid)).sort(
-                {"_id": 1}
-            ):
-                move_to_oversized.extend(doc.get("history", []))
-                last_id = doc["_id"]
-            move_to_oversized.insert(0, snapshot)
-
-            os_col.insert_many(move_to_oversized)
-            snapshot_col.update_one(
-                {"_id": last_id},
-                {
-                    "$set": {"oversized": True, "last": snapshot, "count": 0},
-                    "$unset": {"history": ""},
-                },
-            )
-            snapshot_col.delete_many(
-                self._snapshot_bucket_eid_filter(etype, eid) | {"oversized": False}
-            )
-        except Exception as e:
-            raise DatabaseError(f"Update of snapshot {eid} failed: {e}, {snapshot}") from e
-
-    def save_snapshot(self, etype: str, snapshot: dict, ctime: datetime):
-        """Saves snapshot to specified entity of current master document.
-
-        Will move snapshot to oversized snapshots if the maintained bucket is too large.
-        """
-        # Check `etype`
-        self._assert_etype_exists(etype)
-
-        if "eid" not in snapshot:
-            self.log.error("Snapshot is missing 'eid' field: %s", snapshot)
-            return
-        eid = snapshot["eid"]
-        snapshot["_time_created"] = ctime
-
-        snapshot_col = self._snapshot_col(etype)
-        os_col = self._oversized_snapshots_col(etype)
-
-        # Find out if the snapshot is oversized
-        normal, oversized = self._get_snapshot_state(etype, {eid})
-        if normal:
-            try:
-                res = snapshot_col.update_one(
-                    self._snapshot_bucket_eid_filter(etype, eid)
-                    | {"count": {"$lt": self._snapshot_bucket_size}},
-                    {
-                        "$set": {"last": snapshot},
-                        "$push": {"history": {"$each": [snapshot], "$position": 0}},
-                        "$inc": {"count": 1},
-                        "$setOnInsert": {
-                            "_id": self._snapshot_bucket_id(etype, eid, ctime),
-                            "_time_created": ctime,
-                            "oversized": False,
-                            "latest": True,
-                        },
-                    },
-                    upsert=True,
-                )
-
-                if res.upserted_id is not None:
-                    snapshot_col.update_many(
-                        self._snapshot_bucket_eid_filter(etype, eid)
-                        | {"latest": True, "count": self._snapshot_bucket_size},
-                        {"$unset": {"latest": 1}},
-                    )
-            except (WriteError, OperationFailure, DocumentTooLarge) as e:
-                if e.code != BSON_OBJECT_TOO_LARGE:
-                    raise e
-                # The snapshot is too large, move it to oversized snapshots
-                self.log.info(f"Snapshot of {eid} is too large: {e}, marking as oversized.")
-                self._migrate_to_oversized_snapshot(etype, eid, snapshot)
-                self._cache_snapshot_state(etype, set(), normal)
-            except Exception as e:
-                raise DatabaseError(f"Insert of snapshot {eid} failed: {e}, {snapshot}") from e
-            return
-        elif oversized:
-            # Snapshot is already marked as oversized
-            snapshot_col.update_one(
-                self._snapshot_bucket_eid_filter(etype, eid), {"$set": {"last": snapshot}}
-            )
-            os_col.insert_one(snapshot)
-            return
-
-    def _get_snapshot_state(self, etype: str, eids: set[str]) -> tuple[set, set]:
-        """Get current state of snapshot of given `eid`."""
-        unknown = eids
-        normal = self._normal_snapshot_eids[etype] & unknown
-        oversized = self._oversized_snapshot_eids[etype] & unknown
-        unknown = unknown - normal - oversized
-
-        if not unknown:
-            return normal, oversized
-
-        snapshot_col = self._snapshot_col(etype)
-        new_oversized = set()
-        for doc in snapshot_col.find(
-            self._snapshot_bucket_eids_filter(etype, unknown) | {"oversized": True},
-            {"oversized": 1},
-        ):
-            eid = doc["_id"].rsplit("_#", maxsplit=1)[0]
-            new_oversized.add(eid)
-
-        unknown = unknown - new_oversized
-        self._normal_snapshot_eids[etype] |= unknown
-        self._oversized_snapshot_eids[etype] |= new_oversized
-
-        return normal | unknown, oversized | new_oversized
-
-    def _cache_snapshot_state(self, etype: str, normal: set, oversized: set):
-        """Cache snapshot state for given `etype`."""
-        self._normal_snapshot_eids[etype] |= normal
-
-        self._normal_snapshot_eids[etype] -= oversized
-        self._oversized_snapshot_eids[etype] |= oversized
-
-    def save_snapshots(self, etype: str, snapshots: list[dict], ctime: datetime):
-        """
-        Saves a list of snapshots of current master documents.
-
-        All snapshots must belong to same entity type.
-
-        Will move snapshots to oversized snapshots if the maintained bucket is too large.
-        For better understanding, see `save_snapshot()`.
-        """
-        # Check `etype`
-        self._assert_etype_exists(etype)
-
-        for snapshot in snapshots:
-            snapshot["_time_created"] = ctime
-
-        snapshot_col = self._snapshot_col(etype)
-        os_col = self._oversized_snapshots_col(etype)
-
-        snapshots_by_eid = defaultdict(list)
-        for snapshot in snapshots:
-            if "eid" not in snapshot:
-                continue
-            if "_id" in snapshot:
-                del snapshot["_id"]
-            snapshots_by_eid[snapshot["eid"]].append(snapshot)
-
-        # Find out if any of the snapshots are oversized
-        normal, oversized = self._get_snapshot_state(etype, set(snapshots_by_eid.keys()))
-
-        upserts = []
-        update_originals: list[list[dict]] = []
-        oversized_inserts = []
-        oversized_updates = []
-
-        # A normal snapshot, shift the last snapshot to history and update last
-        for eid in normal:
-            upserts.append(
-                UpdateOne(
-                    self._snapshot_bucket_eid_filter(etype, eid)
-                    | {"count": {"$lt": self._snapshot_bucket_size}},
-                    {
-                        "$set": {"last": snapshots_by_eid[eid][-1]},
-                        "$push": {"history": {"$each": snapshots_by_eid[eid], "$position": 0}},
-                        "$inc": {"count": len(snapshots_by_eid[eid])},
-                        "$setOnInsert": {
-                            "_id": self._snapshot_bucket_id(etype, eid, ctime),
-                            "_time_created": ctime,
-                            "oversized": False,
-                            "latest": True,
-                        },
-                    },
-                    upsert=True,
-                )
-            )
-            update_originals.append(snapshots_by_eid[eid])
-
-        # Snapshot is already marked as oversized
-        for eid in oversized:
-            oversized_inserts.extend(snapshots_by_eid[eid])
-            oversized_updates.append(
-                UpdateOne(
-                    self._snapshot_bucket_eid_filter(etype, eid),
-                    {"$set": {"last": snapshots_by_eid[eid][-1]}},
-                )
-            )
-
-        new_oversized = set()
-
-        if upserts:
-            try:
-                res = snapshot_col.bulk_write(upserts, ordered=False)
-
-                # Unset latest snapshots if new snapshots were inserted
-                if res.upserted_count > 0:
-                    unset_latest_updates = []
-                    for upsert_id in res.upserted_ids.values():
-                        unset_latest_updates.append(
-                            UpdateMany(
-                                self._snapshot_eid_filter_from_bid(etype, upsert_id)
-                                | {"latest": True, "count": self._snapshot_bucket_size},
-                                {"$unset": {"latest": 1}},
-                            )
-                        )
-                    up_res = snapshot_col.bulk_write(unset_latest_updates)
-                    if up_res.modified_count != res.upserted_count:
-                        self.log.info(
-                            "Upserted the first snapshot for %d entities.",
-                            res.upserted_count - up_res.modified_count,
-                        )
-
-                if res.modified_count + res.upserted_count != len(upserts):
-                    self.log.error(
-                        "Some snapshots were not updated, %s != %s",
-                        res.modified_count + res.upserted_count,
-                        len(upserts),
-                    )
-            except (BulkWriteError, OperationFailure) as e:
-                self.log.info("Update of snapshots failed, will retry with oversize.")
-                failed_indexes = [
-                    err["index"]
-                    for err in e.details["writeErrors"]
-                    if err["code"] == BSON_OBJECT_TOO_LARGE
-                ]
-                failed_snapshots = (update_originals[i] for i in failed_indexes)
-                for eid_snapshots in failed_snapshots:
-                    eid = eid_snapshots[0]["eid"]
-                    failed_snapshots = sorted(
-                        eid_snapshots, key=lambda s: s["_time_created"], reverse=True
-                    )
-                    self._migrate_to_oversized_snapshot(etype, eid, failed_snapshots[0])
-                    oversized_inserts.extend(failed_snapshots[1:])
-                    new_oversized.add(eid)
-
-                if any(err["code"] != BSON_OBJECT_TOO_LARGE for err in e.details["writeErrors"]):
-                    # Some other error occurred
-                    raise e
-            except Exception as e:
-                raise DatabaseError(f"Upsert of snapshots failed: {str(e)[:2048]}") from e
-
-        # Update the oversized snapshots
-        if oversized_inserts:
-            try:
-                if oversized_updates:
-                    snapshot_col.bulk_write(oversized_updates)
-                os_col.insert_many(oversized_inserts)
-            except Exception as e:
-                raise DatabaseError(f"Insert of snapshots failed: {str(e)[:2048]}") from e
-
-        # Cache the new state
-        self._cache_snapshot_state(etype, set(), new_oversized)
-
     def _get_metadata_id(self, module: str, time: datetime, worker_id: Optional[int] = None) -> str:
         """Generates unique metadata id based on `module`, `time` and the worker index."""
         worker_id = self._process_index if worker_id is None else worker_id
@@ -1666,64 +998,6 @@ def _split_timeseries_dps(
 
         return result
 
-    def get_distinct_val_count(self, etype: str, attr: str) -> dict[Any, int]:
-        """Counts occurences of distinct values of given attribute in snapshots.
-
-        Returns dictionary mapping value -> count.
-
-        Works for all plain and observation data types except `dict` and `json`.
-        """
-        self._assert_etype_exists(etype)
-
-        snapshot_col = self._snapshot_col(etype)
-
-        # Get attribute specification
-        try:
-            attr_spec = self._db_schema_config.attr(etype, attr)
-        except KeyError as e:
-            raise DatabaseError(f"Attribute '{attr}' does not exist") from e
-
-        if attr_spec.t not in AttrType.PLAIN | AttrType.OBSERVATIONS:
-            raise DatabaseError(f"Attribute '{attr}' isn't plain or observations")
-
-        # Attribute data type must be primitive, array<T> or set<T>
-        if any(needle in attr_spec.data_type.root for needle in ("dict", "json")):
-            raise DatabaseError(
-                f"Data type '{attr_spec.data_type}' of attribute '{attr}' is not processable"
-            )
-
-        # Build aggregation query
-        attr_path = "$last." + attr
-        unwinding = []
-
-        # Unwind array-like and multi value attributes
-        # If attribute is multi value array, unwind twice
-        if "array" in attr_spec.data_type.root or "set" in attr_spec.data_type.root:
-            unwinding.append({"$unwind": attr_path})
-        if attr_spec.t == AttrType.OBSERVATIONS and attr_spec.multi_value:
-            unwinding.append({"$unwind": attr_path})
-
-        # Group
-        agg_query_group_id = attr_path
-        if "link" in attr_spec.data_type.root:
-            agg_query_group_id += ".eid"
-
-        agg_query = [
-            {"$match": {"latest": True}},
-            *unwinding,
-            {"$group": {"_id": agg_query_group_id, "count": {"$sum": 1}}},
-            {"$sort": {"_id": 1, "count": -1}},
-        ]
-        # Run aggregation
-        distinct_counts_cur = snapshot_col.aggregate(agg_query)
-
-        distinct_counts = {x["_id"]: x["count"] for x in distinct_counts_cur}
-
-        if None in distinct_counts:
-            del distinct_counts[None]
-
-        return distinct_counts
-
     def move_raw_to_archive(self, etype: str):
         """Rename the current raw collection to archive collection.
 
@@ -1807,25 +1081,6 @@ def drop_empty_archives(self, etype: str) -> int:
                 raise DatabaseError(f"Drop of empty archive failed: {e}") from e
         return dropped_count
 
-    def delete_old_snapshots(self, etype: str, t_old: datetime) -> int:
-        """Delete old snapshots.
-
-        Periodically called for all `etype`s from HistoryManager.
-        """
-        snapshot_col = self._snapshot_col(etype)
-        os_snapshot_col = self._oversized_snapshots_col(etype)
-        deleted = 0
-        try:
-            res = snapshot_col.delete_many(
-                {"_time_created": {"$lte": t_old - self._bucket_delta}},
-            )
-            deleted += res.deleted_count * self._snapshot_bucket_size
-            res = os_snapshot_col.delete_many({"_time_created": {"$lt": t_old}})
-            deleted += res.deleted_count
-        except Exception as e:
-            raise DatabaseError(f"Delete of olds snapshots failed: {e}") from e
-        return deleted
-
     def get_module_cache(self, override_called_id: Optional[str] = None):
         """Return a persistent cache collection for given module name.
 
diff --git a/dp3/history_management/history_manager.py b/dp3/history_management/history_manager.py
index 7619d89d..b6291403 100644
--- a/dp3/history_management/history_manager.py
+++ b/dp3/history_management/history_manager.py
@@ -165,12 +165,10 @@ def delete_old_snapshots(self):
         self.log.debug("Deleting all snapshots before %s", t_old)
 
         deleted_total = 0
-        for etype in self.model_spec.entities:
-            try:
-                result = self.db.delete_old_snapshots(etype, t_old)
-                deleted_total += result
-            except DatabaseError as e:
-                self.log.exception(e)
+        try:
+            deleted_total = self.db.snapshots.delete_old(t_old)
+        except DatabaseError as e:
+            self.log.exception(e)
         self.log.debug("Deleted %s snapshots in total.", deleted_total)
 
     def archive_old_dps(self):
diff --git a/dp3/snapshots/snapshooter.py b/dp3/snapshots/snapshooter.py
index 3acc2ccf..88d647af 100644
--- a/dp3/snapshots/snapshooter.py
+++ b/dp3/snapshots/snapshooter.py
@@ -427,11 +427,11 @@ def make_linkless_snapshots(
             entity_cnt += 1
             snapshots.append(self.make_linkless_snapshot(etype, master_record, time))
             if len(snapshots) >= DB_SEND_CHUNK:
-                self.db.save_snapshots(etype, snapshots, time)
+                self.db.snapshots.save_many(etype, snapshots, time)
                 snapshots.clear()
 
         if snapshots:
-            self.db.save_snapshots(etype, snapshots, time)
+            self.db.snapshots.save_many(etype, snapshots, time)
             snapshots.clear()
         return entity_cnt
 
@@ -516,7 +516,7 @@ def make_snapshot(self, task: Snapshot):
         for (rtype, _rid), record in entity_values.items():
             if len(record) == 1 and not self.config.keep_empty:
                 continue
-            self.db.save_snapshot(rtype, record, task.time)
+            self.db.snapshots.save_one(rtype, record, task.time)
 
         if task.final:
             self.db.update_metadata(task.time, metadata={"linked_finished": True}, worker_id=0)

From fdee8d973fb4641eee08710e81288ba42a431523 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 12 Nov 2024 09:27:05 +0100
Subject: [PATCH 08/22] Config: Fix Link destination eid data type.

---
 dp3/api/internal/entity_response_models.py | 31 +------------
 dp3/common/attrspec.py                     | 22 +++------
 dp3/common/config.py                       | 54 ++++++++++++++--------
 dp3/common/context.py                      | 23 +++++++++
 dp3/common/datatype.py                     | 22 +++++++--
 dp3/common/entityspec.py                   | 16 -------
 tests/test_api/test_list_entities.py       | 18 +++++++-
 7 files changed, 101 insertions(+), 85 deletions(-)
 create mode 100644 dp3/common/context.py

diff --git a/dp3/api/internal/entity_response_models.py b/dp3/api/internal/entity_response_models.py
index e1b0c5d8..51cbb080 100644
--- a/dp3/api/internal/entity_response_models.py
+++ b/dp3/api/internal/entity_response_models.py
@@ -1,11 +1,10 @@
 from datetime import datetime
 from typing import Annotated, Any, Optional, Union
 
-from pydantic import BaseModel, Field, NonNegativeInt, PlainSerializer, model_validator
+from pydantic import BaseModel, Field, NonNegativeInt, PlainSerializer
 
-from dp3.common.attrspec import AttrSpec, AttrSpecGeneric, AttrSpecType, AttrType
+from dp3.common.attrspec import AttrSpecType, AttrType
 from dp3.common.datapoint import to_json_friendly
-from dp3.common.entityspec import EntitySpec, entity_context
 
 
 class EntityState(BaseModel):
@@ -21,32 +20,6 @@ class EntityState(BaseModel):
     attribs: dict[str, Annotated[AttrSpecType, Field(discriminator="type")]]
     eid_estimate_count: NonNegativeInt
 
-    @model_validator(mode="before")
-    def _validate_attribs(cls, data: Any) -> Any:
-        if not isinstance(data, dict):
-            return data
-        if any(x not in data for x in ["id", "name"]):
-            return data
-
-        entity_spec = EntitySpec.model_validate(
-            {
-                "id": data["id"],
-                "name": data["name"],
-                "id_data_type": data.get("id_data_type", "string"),
-                "snapshot": True,
-            }
-        )
-        print(entity_spec)
-
-        assert isinstance(data["attribs"], dict), "'attribs' must be a dictionary"
-        with entity_context(entity_spec):
-            data["attribs"] = {
-                attr_id: AttrSpec(attr_id, spec) if not isinstance(spec, AttrSpecGeneric) else spec
-                for attr_id, spec in data["attribs"].items()
-            }
-
-        return data
-
 
 # This is necessary to allow for non-JSON-serializable types in the model
 JsonVal = Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")]
diff --git a/dp3/common/attrspec.py b/dp3/common/attrspec.py
index 1ac9849b..79536837 100644
--- a/dp3/common/attrspec.py
+++ b/dp3/common/attrspec.py
@@ -13,6 +13,7 @@
 )
 from pydantic_core.core_schema import FieldValidationInfo
 
+from dp3.common.context import get_entity_context
 from dp3.common.datapoint import (
     DataPointBase,
     DataPointObservationsBase,
@@ -24,7 +25,7 @@
     dp_ts_v_validator,
 )
 from dp3.common.datatype import DataType, ReadOnly
-from dp3.common.entityspec import SpecModel, _init_attr_spec_context_var
+from dp3.common.entityspec import SpecModel
 from dp3.common.types import ParsedTimedelta
 
 # Regex of attribute and series id's
@@ -186,10 +187,7 @@ class AttrSpecPlain(AttrSpecClassic):
     def __init__(self, **data):
         super().__init__(**data)
 
-        entity_spec = _init_attr_spec_context_var.get()
-        if entity_spec is None:
-            raise ValueError("Entity specification not found in context")
-
+        entity_spec = get_entity_context()["self"]
         self._dp_model = create_model(
             f"DataPointPlain_{self.id}",
             __base__=DataPointPlainBase,
@@ -204,10 +202,7 @@ class AttrSpecReadOnly(AttrSpecPlain):
     def __init__(self, **data):
         super().__init__(**data)
 
-        entity_spec = _init_attr_spec_context_var.get()
-        if entity_spec is None:
-            raise ValueError("Entity specification not found in context")
-
+        entity_spec = get_entity_context()["self"]
         self._dp_model = create_model(
             f"DataPointReadOnly_{self.id}",
             __base__=DataPointPlainBase,
@@ -231,10 +226,7 @@ def __init__(self, **data):
         super().__init__(**data)
 
         value_validator = self.data_type.data_type
-        entity_spec = _init_attr_spec_context_var.get()
-        if entity_spec is None:
-            raise ValueError("Entity specification not found in context")
-
+        entity_spec = get_entity_context()["self"]
         self._dp_model = create_model(
             f"DataPointObservations_{self.id}",
             __base__=DataPointObservationsBase,
@@ -256,9 +248,7 @@ class AttrSpecTimeseries(AttrSpecGeneric):
     def __init__(self, **data):
         super().__init__(**data)
 
-        entity_spec = _init_attr_spec_context_var.get()
-        if entity_spec is None:
-            raise ValueError("Entity specification not found in context")
+        entity_spec = get_entity_context()["self"]
 
         # Typing of `v` field
         dp_value_typing = {}
diff --git a/dp3/common/config.py b/dp3/common/config.py
index 23351e0d..28e50c40 100644
--- a/dp3/common/config.py
+++ b/dp3/common/config.py
@@ -25,7 +25,8 @@
     AttrSpecReadOnly,
     AttrSpecType,
 )
-from dp3.common.entityspec import EntitySpec, entity_context
+from dp3.common.context import entity_context
+from dp3.common.entityspec import EntitySpec
 
 
 class NoDefault:
@@ -200,23 +201,6 @@ class EntitySpecDict(BaseModel):
     def __getitem__(self, item):
         return self.__getattribute__(item)
 
-    @model_validator(mode="before")
-    def _validate_attr_spec(cls, data: Any) -> Any:
-        if not isinstance(data, dict) or "entity" not in data or "attribs" not in data:
-            return data
-
-        if not isinstance(data["entity"], EntitySpec):
-            data["entity"] = EntitySpec.model_validate(data["entity"])
-
-        assert isinstance(data["attribs"], dict), "'attribs' must be a dictionary"
-        with entity_context(data["entity"]):
-            data["attribs"] = {
-                attr_id: AttrSpec(attr_id, spec) if not isinstance(spec, AttrSpecGeneric) else spec
-                for attr_id, spec in data["attribs"].items()
-            }
-
-        return data
-
 
 class ModelSpec(BaseModel):
     """
@@ -267,6 +251,40 @@ def __init__(self, config: HierarchicalDict):
             config=config, entities={}, attributes={}, entity_attributes={}, relations={}
         )
 
+    @model_validator(mode="before")
+    def _validate_config(cls, data: Any) -> Any:
+        if not isinstance(data, dict):
+            return data
+        if not all(isinstance(entity_dict, dict) for entity_dict in data.values()):
+            return data
+
+        data["entities"] = {}
+        if "config" not in data:
+            return data
+        config = data["config"]
+
+        # First validate all entities for global context
+        for entity_id, entity_dict in config.items():
+            if "entity" not in entity_dict or "attribs" not in entity_dict:
+                raise ValueError(f"Invalid entity specification for '{entity_id}'")
+            if not isinstance(entity_dict["entity"], EntitySpec):
+                entity_dict["entity"] = EntitySpec.model_validate(entity_dict["entity"])
+            data["entities"][entity_id] = entity_dict["entity"]
+
+        # Then validate all attributes
+        for entity_id, entity_dict in config.items():
+            if "attribs" not in entity_dict:
+                continue
+            with entity_context(data["entities"][entity_id], data["entities"]):
+                entity_dict["attribs"] = {
+                    attr_id: (
+                        AttrSpec(attr_id, spec) if not isinstance(spec, AttrSpecGeneric) else spec
+                    )
+                    for attr_id, spec in entity_dict["attribs"].items()
+                }
+
+        return data
+
     @field_validator("entities")
     def _fill_entities(cls, v, info: FieldValidationInfo):
         if "config" not in info.data:
diff --git a/dp3/common/context.py b/dp3/common/context.py
new file mode 100644
index 00000000..c6614894
--- /dev/null
+++ b/dp3/common/context.py
@@ -0,0 +1,23 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+_init_attr_spec_context_var = ContextVar("_init_attr_spec_context_var", default=None)
+
+
+@contextmanager
+def entity_context(self_spec, entities: dict) -> Iterator[None]:
+    """Context manager for AttrSpec initialization."""
+    token = _init_attr_spec_context_var.set({"self": self_spec, "entities": entities})
+    try:
+        yield
+    finally:
+        _init_attr_spec_context_var.reset(token)
+
+
+def get_entity_context() -> dict:
+    """Get entity spec context."""
+    cxt = _init_attr_spec_context_var.get()
+    if cxt is None or not isinstance(cxt, dict):
+        raise ValueError("Entity spec context is not set")
+    return cxt
diff --git a/dp3/common/datatype.py b/dp3/common/datatype.py
index bef7cb56..ec7fde98 100644
--- a/dp3/common/datatype.py
+++ b/dp3/common/datatype.py
@@ -6,7 +6,6 @@
 
 from pydantic import (
     BaseModel,
-    Extra,
     Json,
     PrivateAttr,
     RootModel,
@@ -14,6 +13,7 @@
     model_validator,
 )
 
+from dp3.common.context import get_entity_context
 from dp3.common.mac_address import MACAddress
 
 # Regular expressions for parsing various data types
@@ -53,8 +53,8 @@ def fail(cls, data):
         raise ValueError("Cannot instantiate a read-only attribute datapoint.")
 
 
-class Link(BaseModel, extra=Extra.forbid):
-    eid: str
+class Link(BaseModel, extra="forbid"):
+    eid: Any
 
 
 class DataType(RootModel):
@@ -145,13 +145,25 @@ def _determine_value_validator(self):
             self._mirror_as = mirrored if mirrored else None
             self._type_info = f"link<{etype},{data}>"
 
+            context = get_entity_context()
+            if etype not in context["entities"]:
+                raise ValueError(f"Entity type '{etype}' is not defined")
+            entity_spec = context["entities"][etype]
+
             if etype and data:
                 value_type = DataType(root=data)
                 data_type = create_model(
-                    f"Link<{data}>", __base__=Link, data=(value_type.data_type, ...)
+                    f"Link<{data}>",
+                    __base__=Link,
+                    eid=(entity_spec.id_data_type.data_type, ...),
+                    data=(value_type.data_type, ...),
                 )
             else:
-                data_type = Link
+                data_type = create_model(
+                    f"Link<{etype}>",
+                    __base__=Link,
+                    eid=(entity_spec.id_data_type.data_type, ...),
+                )
 
         elif re.match(re_dict, str_type):
             # Dict
diff --git a/dp3/common/entityspec.py b/dp3/common/entityspec.py
index 077dc4dc..fe00cdb9 100644
--- a/dp3/common/entityspec.py
+++ b/dp3/common/entityspec.py
@@ -1,6 +1,3 @@
-from collections.abc import Iterator
-from contextlib import contextmanager
-from contextvars import ContextVar
 from typing import Literal, Union
 
 from pydantic import BaseModel, Field, PrivateAttr, model_validator
@@ -81,16 +78,3 @@ def fill_validator(self):
 
     def validate_eid(self, v):
         return self._eid_validator(v)
-
-
-_init_attr_spec_context_var = ContextVar("_init_attr_spec_context_var", default=None)
-
-
-@contextmanager
-def entity_context(entity_spec: EntitySpec) -> Iterator[None]:
-    """Context manager for setting the `model_spec` context variable."""
-    token = _init_attr_spec_context_var.set(entity_spec)
-    try:
-        yield
-    finally:
-        _init_attr_spec_context_var.reset(token)
diff --git a/tests/test_api/test_list_entities.py b/tests/test_api/test_list_entities.py
index 1caea6dd..80ed388a 100644
--- a/tests/test_api/test_list_entities.py
+++ b/tests/test_api/test_list_entities.py
@@ -1,6 +1,8 @@
 import common
 
 from dp3.api.internal.entity_response_models import EntityState
+from dp3.common.context import entity_context
+from dp3.common.entityspec import EntitySpec
 
 
 class ListEntities(common.APITest):
@@ -10,6 +12,20 @@ def test_get_entities(self):
 
         received = response.json()
         self.assertIsInstance(received, dict)
+
+        entities = {}
+        # Validate received entity data
+        for key, item in received.items():
+            self.assertIsInstance(key, str)
+            self.assertIsInstance(item, dict)
+
+            entity_keys = set(item.keys()) & {"id", "id_data_type", "name"}
+            self.assertSetEqual(entity_keys, {"id", "id_data_type", "name"})
+            entities[key] = EntitySpec.model_validate(
+                {k: v for k, v in item.items() if k in entity_keys} | {"snapshot": True}
+            )
+
         for key, item in received.items():
-            received[key] = EntityState.model_validate(item)
+            with entity_context(entities[key], entities):
+                received[key] = EntityState.model_validate(item)
         self.assertListEqual(list(received.keys()), list(common.MODEL_SPEC.entities.keys()))

From 1265fabf1047acf545c59c1095880c030673bc4d Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 12 Nov 2024 09:33:42 +0100
Subject: [PATCH 09/22] Snapshot: Fix types in entity tuple messages.

---
 dp3/common/config.py         | 26 +++++++++++++++++++++++
 dp3/common/task.py           | 35 ++++++++++++++++++++++++++++---
 dp3/snapshots/snapshooter.py | 40 ++++++++++++++++++++----------------
 3 files changed, 80 insertions(+), 21 deletions(-)

diff --git a/dp3/common/config.py b/dp3/common/config.py
index 28e50c40..3b1f8eef 100644
--- a/dp3/common/config.py
+++ b/dp3/common/config.py
@@ -3,6 +3,9 @@
 """
 
 import os
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
 from typing import Annotated, Any, Optional, Union
 
 import yaml
@@ -443,3 +446,26 @@ def __repr__(self):
             f"num_processes={self.num_processes}, "
             f"process_index={self.process_index})"
         )
+
+
+_init_entity_type_context_var = ContextVar("init_entity_type_context_var", default=None)
+
+
+@contextmanager
+def entity_type_context(model_spec: ModelSpec) -> Iterator[None]:
+    """Context manager for AttrSpec initialization."""
+    token = _init_entity_type_context_var.set(
+        {entity: spec.id_data_type.root for entity, spec in model_spec.entities.items()}
+    )
+    try:
+        yield
+    finally:
+        _init_entity_type_context_var.reset(token)
+
+
+def get_entity_type_context() -> dict:
+    """Get entity spec context."""
+    cxt = _init_entity_type_context_var.get()
+    if cxt is None or not isinstance(cxt, dict):
+        raise ValueError("Entity type context is not set")
+    return cxt
diff --git a/dp3/common/task.py b/dp3/common/task.py
index 22601bf6..321b8129 100644
--- a/dp3/common/task.py
+++ b/dp3/common/task.py
@@ -5,20 +5,24 @@
 from contextvars import ContextVar
 from datetime import datetime
 from enum import Enum
-from typing import Annotated, Any, Optional
+from ipaddress import IPv4Address, IPv6Address
+from typing import Annotated, Any, Callable, Optional, Union
 
 from pydantic import (
     AfterValidator,
     BaseModel,
     BeforeValidator,
+    Discriminator,
     PlainSerializer,
+    Tag,
     ValidationError,
     field_validator,
 )
 from pydantic_core.core_schema import FieldValidationInfo
 
-from dp3.common.config import ModelSpec
+from dp3.common.config import ModelSpec, entity_type_context, get_entity_type_context
 from dp3.common.datapoint import DataPointBase, to_json_friendly
+from dp3.common.mac_address import MACAddress
 
 _init_context_var = ContextVar("_init_context_var", default=None)
 
@@ -186,6 +190,23 @@ class SnapshotMessageType(Enum):
     run_end = "run_end"
 
 
+def get_discriminator_value(entity_tuple: tuple[str, Any]) -> str:
+    name_2_type_name = get_entity_type_context()
+    return name_2_type_name[entity_tuple[0]]
+
+
+EntityTuple = Annotated[
+    Union[
+        Annotated[tuple[str, str], Tag("string")],
+        Annotated[tuple[str, int], Tag("int")],
+        Annotated[tuple[str, IPv4Address], Tag("ipv4")],
+        Annotated[tuple[str, IPv6Address], Tag("ipv6")],
+        Annotated[tuple[str, MACAddress], Tag("mac")],
+    ],
+    Discriminator(get_discriminator_value),
+]
+
+
 class Snapshot(Task):
     """Snapshot
 
@@ -200,7 +221,7 @@ class Snapshot(Task):
         final: If True, this is the last linked snapshot for the given time
     """
 
-    entities: list[tuple[str, str]] = []
+    entities: list[EntityTuple] = []
     time: datetime
     type: SnapshotMessageType
     final: bool = False
@@ -210,3 +231,11 @@ def routing_key(self):
 
     def as_message(self) -> str:
         return self.model_dump_json()
+
+    @staticmethod
+    def get_validator(model_spec: ModelSpec) -> Callable[[Union[str, bytes]], "Snapshot"]:
+        def json_validator(serialized: Union[str, bytes]) -> Snapshot:
+            with entity_type_context(model_spec):
+                return Snapshot.model_validate_json(serialized)
+
+        return json_validator
diff --git a/dp3/snapshots/snapshooter.py b/dp3/snapshots/snapshooter.py
index 88d647af..434220f2 100644
--- a/dp3/snapshots/snapshooter.py
+++ b/dp3/snapshots/snapshooter.py
@@ -32,7 +32,7 @@
     AttrType,
     ObservationsHistoryParams,
 )
-from dp3.common.config import CronExpression, PlatformConfig
+from dp3.common.config import CronExpression, PlatformConfig, entity_type_context
 from dp3.common.scheduler import Scheduler
 from dp3.common.task import DataPointTask, Snapshot, SnapshotMessageType, task_context
 from dp3.common.types import EventGroupType
@@ -89,7 +89,7 @@ def __init__(
         queue = f"{platform_config.app_name}-worker-{platform_config.process_index}-snapshots"
         self.snapshot_queue_reader = TaskQueueReader(
             callback=self.process_snapshot_task,
-            parse_task=Snapshot.model_validate_json,
+            parse_task=Snapshot.get_validator(self.model_spec),
             app_name=platform_config.app_name,
             worker_index=platform_config.process_index,
             rabbit_config=platform_config.config.get("processing_core.msg_broker", {}),
@@ -240,16 +240,18 @@ def make_snapshots(self):
             },
         )
         # Broadcast run start
-        self.snapshot_queue_writer.broadcast_task(
-            task=Snapshot(type=SnapshotMessageType.run_start, time=time)
-        )
+        with entity_type_context(self.model_spec):
+            self.snapshot_queue_writer.broadcast_task(
+                task=Snapshot(type=SnapshotMessageType.run_start, time=time)
+            )
 
         # distribute list of possibly linked entities to all workers
         cached = self.get_cached_link_entity_ids()
         self.log.debug("Broadcasting %s cached linked entities", len(cached))
-        self.snapshot_queue_writer.broadcast_task(
-            task=Snapshot(entities=cached, time=time, type=SnapshotMessageType.linked_entities)
-        )
+        with entity_type_context(self.model_spec):
+            self.snapshot_queue_writer.broadcast_task(
+                task=Snapshot(entities=cached, time=time, type=SnapshotMessageType.linked_entities)
+            )
 
         # Load links only for a reduced set of entities
         self.log.debug("Loading linked entities.")
@@ -264,14 +266,15 @@ def make_snapshots(self):
                 counts["entities"] += len(linked_entities_component)
                 counts["components"] += 1
 
-                self.snapshot_queue_writer.put_task(
-                    task=Snapshot(
-                        entities=linked_entities_component,
-                        time=time,
-                        type=SnapshotMessageType.task,
-                        final=(i + 1 == len(linked_entities)),
+                with entity_type_context(self.model_spec):
+                    self.snapshot_queue_writer.put_task(
+                        task=Snapshot(
+                            entities=linked_entities_component,
+                            time=time,
+                            type=SnapshotMessageType.task,
+                            final=(i + 1 == len(linked_entities)),
+                        )
                     )
-                )
 
             if len(linked_entities) == 0:
                 self.db.update_metadata(time, metadata={"linked_finished": True})
@@ -286,9 +289,10 @@ def make_snapshots(self):
             )
 
         # Broadcast run end
-        self.snapshot_queue_writer.broadcast_task(
-            task=Snapshot(type=SnapshotMessageType.run_end, time=time)
-        )
+        with entity_type_context(self.model_spec):
+            self.snapshot_queue_writer.broadcast_task(
+                task=Snapshot(type=SnapshotMessageType.run_end, time=time)
+            )
 
     def get_cached_link_entity_ids(self):
         used = [f"{etype}#{attr}" for etype, attr in self._correlation_hooks.used_links]

From f93af1f4d396026007f8573e1b7fdec446afe924 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 12 Nov 2024 09:57:11 +0100
Subject: [PATCH 10/22] Fix typehints and docstrings.

---
 dp3/common/callback_registrar.py  | 17 +++++++++--------
 dp3/common/datatype.py            | 13 ++++++++++---
 dp3/core/collector.py             |  9 +++++----
 dp3/core/link_manager.py          | 13 +++++++------
 dp3/database/database.py          | 19 ++++++++++---------
 dp3/database/magic.py             |  2 +-
 dp3/database/schema_cleaner.py    |  5 +++--
 dp3/database/snapshots.py         | 27 ++++++++++++++-------------
 dp3/task_processing/task_hooks.py |  7 ++++---
 9 files changed, 63 insertions(+), 49 deletions(-)

diff --git a/dp3/common/callback_registrar.py b/dp3/common/callback_registrar.py
index efb24308..ba13c299 100644
--- a/dp3/common/callback_registrar.py
+++ b/dp3/common/callback_registrar.py
@@ -8,6 +8,7 @@
 from dp3.common.attrspec import AttrType
 from dp3.common.config import ModelSpec, PlatformConfig, read_config_dir
 from dp3.common.datapoint import DataPointType
+from dp3.common.datatype import AnyEidT
 from dp3.common.scheduler import Scheduler
 from dp3.common.state import SharedFlag
 from dp3.common.task import DataPointTask
@@ -38,7 +39,7 @@ def write_datapoints_into_record(model_spec: ModelSpec, tasks: list[DataPointTas
 def on_entity_creation_in_snapshots(
     model_spec: ModelSpec,
     run_flag: SharedFlag,
-    original_hook: Callable[[str, DataPointTask], list[DataPointTask]],
+    original_hook: Callable[[AnyEidT, DataPointTask], list[DataPointTask]],
     etype: str,
     record: dict,
 ) -> list[DataPointTask]:
@@ -55,7 +56,7 @@ def on_entity_creation_in_snapshots(
 def on_attr_change_in_snapshots(
     model_spec: ModelSpec,
     run_flag: SharedFlag,
-    original_hook: Callable[[str, DataPointTask], Union[list[DataPointTask], None]],
+    original_hook: Callable[[AnyEidT, DataPointTask], Union[list[DataPointTask], None]],
     etype: str,
     record: dict,
 ) -> list[DataPointTask]:
@@ -194,7 +195,7 @@ def register_task_hook(self, hook_type: str, hook: Callable):
 
     def register_allow_entity_creation_hook(
         self,
-        hook: Callable[[str, DataPointTask], bool],
+        hook: Callable[[AnyEidT, DataPointTask], bool],
         entity: str,
     ):
         """
@@ -210,7 +211,7 @@ def register_allow_entity_creation_hook(
 
     def register_on_entity_creation_hook(
         self,
-        hook: Callable[[str, DataPointTask], list[DataPointTask]],
+        hook: Callable[[AnyEidT, DataPointTask], list[DataPointTask]],
         entity: str,
         refresh: SharedFlag = None,
         may_change: list[list[str]] = None,
@@ -259,7 +260,7 @@ def register_entity_hook(self, hook_type: str, hook: Callable, entity: str):
 
     def register_on_new_attr_hook(
         self,
-        hook: Callable[[str, DataPointType], Union[None, list[DataPointTask]]],
+        hook: Callable[[AnyEidT, DataPointType], Union[None, list[DataPointTask]]],
         entity: str,
         attr: str,
         refresh: SharedFlag = None,
@@ -269,7 +270,7 @@ def register_on_new_attr_hook(
         Registers passed hook to be called on new attribute datapoint.
 
         Args:
-            hook: `hook` callable should expect eid, Task and attribute value as arguments.
+            hook: `hook` callable should expect eid and a datapoint as arguments.
                 Can optionally return a list of DataPointTasks to perform.
             entity: specifies entity type
             attr: specifies attribute name
@@ -391,7 +392,7 @@ def register_snapshot_finalize_hook(self, hook: Callable[[], list[DataPointTask]
 
     def register_periodic_update_hook(
         self,
-        hook: Callable[[str, str, dict], list[DataPointTask]],
+        hook: Callable[[str, AnyEidT, dict], list[DataPointTask]],
         hook_id: str,
         entity_type: str,
         period: ParsedTimedelta,
@@ -412,7 +413,7 @@ def register_periodic_update_hook(
 
     def register_periodic_eid_update_hook(
         self,
-        hook: Callable[[str, str], list[DataPointTask]],
+        hook: Callable[[str, AnyEidT], list[DataPointTask]],
         hook_id: str,
         entity_type: str,
         period: ParsedTimedelta,
diff --git a/dp3/common/datatype.py b/dp3/common/datatype.py
index ec7fde98..fe8b80ce 100644
--- a/dp3/common/datatype.py
+++ b/dp3/common/datatype.py
@@ -1,7 +1,7 @@
-import ipaddress
 import re
 from datetime import datetime
 from enum import Enum
+from ipaddress import IPv4Address, IPv6Address
 from typing import Any, Optional, Union
 
 from pydantic import (
@@ -33,8 +33,8 @@
     "int": int,
     "int64": int,
     "float": float,
-    "ipv4": ipaddress.IPv4Address,
-    "ipv6": ipaddress.IPv6Address,
+    "ipv4": IPv4Address,
+    "ipv6": IPv6Address,
     "mac": MACAddress,
     "time": datetime,
     "special": Any,
@@ -43,6 +43,13 @@
 
 eid_data_types = ["string", "int", "ipv4", "ipv6", "mac"]
 
+AnyEidT = Union[str, int, IPv4Address, IPv6Address, MACAddress]
+"""Type alias for any of possible entity ID data types.
+
+Note that the type is determined based on the loaded entity configuration
+and in most cases is only one of the options, based on what entity is being processed.
+"""
+
 
 class ReadOnly(BaseModel):
     """The ReadOnly data_type is used to avoid datapoint insertion for an attribute."""
diff --git a/dp3/core/collector.py b/dp3/core/collector.py
index 3afb41d3..08a60fce 100644
--- a/dp3/core/collector.py
+++ b/dp3/core/collector.py
@@ -13,6 +13,7 @@
 from dp3.common.callback_registrar import CallbackRegistrar
 from dp3.common.config import CronExpression, PlatformConfig
 from dp3.common.datapoint import DataPointBase, DataPointObservationsBase, DataPointTimeseriesBase
+from dp3.common.datatype import AnyEidT
 from dp3.common.task import DataPointTask
 from dp3.database.database import EntityDatabase
 
@@ -93,7 +94,7 @@ def _get_inverse_relations(self) -> dict[str, set[tuple[str, str]]]:
         return inverse_relations
 
     def extend_ttl_on_create(
-        self, eid: str, task: DataPointTask, base_ttl: timedelta
+        self, eid: AnyEidT, task: DataPointTask, base_ttl: timedelta
     ) -> list[DataPointTask]:
         """Extends the TTL of the entity by the specified timedelta."""
         task = DataPointTask(
@@ -276,7 +277,7 @@ def collect_ttl(self, etype: str):
         )
 
     def extend_plain_ttl(
-        self, eid: str, dp: DataPointBase, extend_by: timedelta
+        self, eid: AnyEidT, dp: DataPointBase, extend_by: timedelta
     ) -> list[DataPointTask]:
         """Extends the TTL of the entity by the specified timedelta."""
         now = datetime.utcnow()
@@ -288,7 +289,7 @@ def extend_plain_ttl(
         return [task]
 
     def extend_observations_ttl(
-        self, eid: str, dp: DataPointObservationsBase, extend_by: timedelta
+        self, eid: AnyEidT, dp: DataPointObservationsBase, extend_by: timedelta
     ) -> list[DataPointTask]:
         """Extends the TTL of the entity by the specified timedelta."""
         task = DataPointTask(
@@ -299,7 +300,7 @@ def extend_observations_ttl(
         return [task]
 
     def extend_timeseries_ttl(
-        self, eid: str, dp: DataPointTimeseriesBase, extend_by: timedelta
+        self, eid: AnyEidT, dp: DataPointTimeseriesBase, extend_by: timedelta
     ) -> list[DataPointTask]:
         """Extends the TTL of the entity by the specified timedelta."""
         task = DataPointTask(
diff --git a/dp3/core/link_manager.py b/dp3/core/link_manager.py
index c41649e4..4d6d9726 100644
--- a/dp3/core/link_manager.py
+++ b/dp3/core/link_manager.py
@@ -12,6 +12,7 @@
 from dp3.common.callback_registrar import CallbackRegistrar
 from dp3.common.config import PlatformConfig
 from dp3.common.datapoint import DataPointBase, DataPointObservationsBase
+from dp3.common.datatype import AnyEidT
 from dp3.database.database import EntityDatabase
 
 
@@ -70,7 +71,7 @@ def _setup_cache_indexes(self):
         self.cache.create_index("using_attr", background=True)
         self.cache.create_index("ttl", expireAfterSeconds=0, background=True)
 
-    def add_plain_to_link_cache(self, etype_to: str, eid: str, dp: DataPointBase):
+    def add_plain_to_link_cache(self, etype_to: str, eid: AnyEidT, dp: DataPointBase):
         self.cache.update_one(
             {
                 "to": f"{etype_to}#{dp.v.eid}",
@@ -81,7 +82,7 @@ def add_plain_to_link_cache(self, etype_to: str, eid: str, dp: DataPointBase):
             upsert=True,
         )
 
-    def add_iterable_plain_to_link_cache(self, etype_to: str, eid: str, dp: DataPointBase):
+    def add_iterable_plain_to_link_cache(self, etype_to: str, eid: AnyEidT, dp: DataPointBase):
         linked_eids = [v.eid for v in dp.v]
         self.cache.update_many(
             {
@@ -94,7 +95,7 @@ def add_iterable_plain_to_link_cache(self, etype_to: str, eid: str, dp: DataPoin
         )
 
     def add_observation_to_link_cache(
-        self, etype_to: str, post_validity: timedelta, eid: str, dp: DataPointObservationsBase
+        self, etype_to: str, post_validity: timedelta, eid: AnyEidT, dp: DataPointObservationsBase
     ):
         self.cache.update_one(
             {
@@ -107,7 +108,7 @@ def add_observation_to_link_cache(
         )
 
     def add_iterable_observation_to_link_cache(
-        self, etype_to: str, post_validity: timedelta, eid: str, dp: DataPointObservationsBase
+        self, etype_to: str, post_validity: timedelta, eid: AnyEidT, dp: DataPointObservationsBase
     ):
         linked_eids = [v.eid for v in dp.v]
         self.cache.update_many(
@@ -120,7 +121,7 @@ def add_iterable_observation_to_link_cache(
             upsert=True,
         )
 
-    def remove_link_cache_of_deleted(self, etype: str, eid: str):
+    def remove_link_cache_of_deleted(self, etype: str, eid: AnyEidT):
         # Delete from master records
         # Get all links to the deleted entity grouped by attribute
         links = self.cache.aggregate(
@@ -140,7 +141,7 @@ def remove_link_cache_of_deleted(self, etype: str, eid: str):
             [DeleteMany({"from": f"{etype}#{eid}"}), DeleteMany({"to": f"{etype}#{eid}"})]
         )
 
-    def remove_link_cache_of_many_deleted(self, etype: str, eids: list[str]):
+    def remove_link_cache_of_many_deleted(self, etype: str, eids: list[AnyEidT]):
         # Delete from master records
         # Get all links to the deleted entities grouped by attribute
         links = self.cache.aggregate(
diff --git a/dp3/database/database.py b/dp3/database/database.py
index 688a1faa..4863af7e 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -5,7 +5,7 @@
 from collections import defaultdict
 from collections.abc import Generator, Iterator
 from datetime import datetime
-from typing import Any, Callable, Optional
+from typing import Callable, Optional
 
 import pymongo
 from event_count_logger import DummyEventGroup
@@ -21,6 +21,7 @@
 from dp3.common.attrspec import AttrType, timeseries_types
 from dp3.common.config import HierarchicalDict, ModelSpec
 from dp3.common.datapoint import DataPointBase
+from dp3.common.datatype import AnyEidT
 from dp3.common.scheduler import Scheduler
 from dp3.common.task import HASH
 from dp3.common.types import EventGroupType
@@ -360,7 +361,7 @@ def await_updated_schema(self):
         self.schema_cleaner.await_updated()
 
     def insert_datapoints(
-        self, eid: Any, dps: list[DataPointBase], new_entity: bool = False
+        self, eid: AnyEidT, dps: list[DataPointBase], new_entity: bool = False
     ) -> None:
         """Inserts datapoint to raw data collection and updates master record.
 
@@ -509,7 +510,7 @@ def update_master_records(self, etype: str, eids: list[str], records: list[dict]
         except Exception as e:
             raise DatabaseError(f"Update of master records failed: {e}\n{records}") from e
 
-    def extend_ttl(self, etype: str, eid: Any, ttl_tokens: dict[str, datetime]):
+    def extend_ttl(self, etype: str, eid: AnyEidT, ttl_tokens: dict[str, datetime]):
         """Extends TTL of given `etype`:`eid` by `ttl_tokens`."""
         master_col = self._master_col(etype)
         try:
@@ -567,7 +568,7 @@ def delete_eids(self, etype: str, eids: list[str]):
             except Exception as e:
                 self.log.exception("Error in on_entity_delete_many callback %s: %s", f, e)
 
-    def delete_eid(self, etype: str, eid: Any):
+    def delete_eid(self, etype: str, eid: AnyEidT):
         """Delete master record and all snapshots of `etype`:`eid`."""
         master_col = self._master_col(etype)
         try:
@@ -726,7 +727,7 @@ def delete_many_link_dps(
         except Exception as e:
             raise DatabaseError(f"Delete of link datapoints failed: {e}") from e
 
-    def get_master_record(self, etype: str, eid: Any, **kwargs) -> dict:
+    def get_master_record(self, etype: str, eid: AnyEidT, **kwargs) -> dict:
         """Get current master record for etype/eid.
 
         If doesn't exist, returns {}.
@@ -738,7 +739,7 @@ def get_master_record(self, etype: str, eid: Any, **kwargs) -> dict:
 
         return master_col.find_one({"_id": entity_spec.validate_eid(eid)}, **kwargs) or {}
 
-    def ekey_exists(self, etype: str, eid: Any) -> bool:
+    def ekey_exists(self, etype: str, eid: AnyEidT) -> bool:
         """Checks whether master record for etype/eid exists"""
         return bool(self.get_master_record(etype, eid))
 
@@ -767,7 +768,7 @@ def get_value_or_history(
         self,
         etype: str,
         attr_name: str,
-        eid: Any,
+        eid: AnyEidT,
         t1: Optional[datetime] = None,
         t2: Optional[datetime] = None,
     ) -> dict:
@@ -851,7 +852,7 @@ def get_observation_history(
         self,
         etype: str,
         attr_name: str,
-        eid: Any,
+        eid: AnyEidT,
         t1: datetime = None,
         t2: datetime = None,
         sort: int = None,
@@ -894,7 +895,7 @@ def get_timeseries_history(
         self,
         etype: str,
         attr_name: str,
-        eid: Any,
+        eid: AnyEidT,
         t1: datetime = None,
         t2: datetime = None,
         sort: int = None,
diff --git a/dp3/database/magic.py b/dp3/database/magic.py
index e0ba99b1..2d9e5e5e 100644
--- a/dp3/database/magic.py
+++ b/dp3/database/magic.py
@@ -149,7 +149,7 @@ def search_and_replace(query: dict[str, Any]) -> dict[str, Any]:
     """Search and replace all occurrences of magic strings in the query.
 
     A helper function to
-    [`get_latest_snapshots`][dp3.database.database.EntityDatabase.get_latest_snapshots]
+    [`snapshots.get_latest`][dp3.database.snapshots.TypedSnapshotCollection.get_latest]
     """
     if isinstance(query, dict):
         for key, value in query.items():
diff --git a/dp3/database/schema_cleaner.py b/dp3/database/schema_cleaner.py
index af314fbe..08598712 100644
--- a/dp3/database/schema_cleaner.py
+++ b/dp3/database/schema_cleaner.py
@@ -3,6 +3,7 @@
 from collections import defaultdict
 from datetime import datetime
 from logging import Logger
+from typing import Callable
 
 import pymongo
 from pymongo import DeleteOne, InsertOne
@@ -56,7 +57,7 @@ def __init__(
             "snapshot_bucket_size": self._config.get("database.storage.snapshot_bucket_size", 32)
         }
 
-        self.migrations = {2: self.migrate_schema_2_to_3}
+        self.migrations: dict[int, Callable[[dict], dict]] = {2: self.migrate_schema_2_to_3}
 
     def get_current_schema_doc(self, infer: bool = False) -> dict:
         """
@@ -520,7 +521,7 @@ def migrate(self, schema_doc: dict) -> dict:
 
         return schema_doc
 
-    def migrate_schema_2_to_3(self, schema) -> dict:
+    def migrate_schema_2_to_3(self, schema: dict) -> dict:
         """
         Migrates schema from version 2 to version 3.
 
diff --git a/dp3/database/snapshots.py b/dp3/database/snapshots.py
index eabb9118..9e34e06f 100644
--- a/dp3/database/snapshots.py
+++ b/dp3/database/snapshots.py
@@ -19,6 +19,7 @@
 
 from dp3.common.attrspec import AttrSpecType, AttrType
 from dp3.common.config import ModelSpec
+from dp3.common.datatype import AnyEidT
 from dp3.common.mac_address import MACAddress
 from dp3.common.utils import int2bytes
 from dp3.database.config import MongoConfig
@@ -124,7 +125,7 @@ def _binary_bucket_range(self, eid: bytes) -> dict:
         }
 
     @abc.abstractmethod
-    def _bucket_id(self, eid: Any, ctime: datetime) -> Union[str, Binary]:
+    def _bucket_id(self, eid: AnyEidT, ctime: datetime) -> Union[str, Binary]:
         """Returns `_id` for snapshot bucket document."""
 
     @abc.abstractmethod
@@ -135,14 +136,14 @@ def _filter_from_bid(self, b_id: Union[bytes, str]) -> dict:
         """
 
     @abc.abstractmethod
-    def _filter_from_eid(self, eid: Any) -> dict:
+    def _filter_from_eid(self, eid: AnyEidT) -> dict:
         """Returns filter for snapshots with given eid."""
 
     @abc.abstractmethod
     def _filter_from_eids(self, eids: Iterable[Any]) -> dict:
         """Returns filter for snapshots with given eids."""
 
-    def get_latest_one(self, eid: Any) -> dict:
+    def get_latest_one(self, eid: AnyEidT) -> dict:
         """Get latest snapshot of given eid.
 
         If doesn't exist, returns {}.
@@ -229,7 +230,7 @@ def get_latest(
             raise SnapshotCollectionError(f"Query is invalid: {e}") from e
 
     def get_by_eid(
-        self, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
+        self, eid: AnyEidT, t1: Optional[datetime] = None, t2: Optional[datetime] = None
     ) -> Union[Cursor, CommandCursor]:
         """Get all (or filtered) snapshots of given `eid`.
 
@@ -326,7 +327,7 @@ def get_distinct_val_count(self, attr: str) -> dict[Any, int]:
         return distinct_counts
 
     def _get_oversized(
-        self, eid: Any, t1: Optional[datetime] = None, t2: Optional[datetime] = None
+        self, eid: AnyEidT, t1: Optional[datetime] = None, t2: Optional[datetime] = None
     ) -> Cursor:
         """Get all (or filtered) snapshots of given `eid` from oversized snapshots collection."""
         snapshot_col = self._os_col()
@@ -346,7 +347,7 @@ def _get_oversized(
             [("_time_created", pymongo.ASCENDING)]
         )
 
-    def _migrate_to_oversized(self, eid: Any, snapshot: dict):
+    def _migrate_to_oversized(self, eid: AnyEidT, snapshot: dict):
         snapshot_col = self._col()
         os_col = self._os_col()
 
@@ -607,7 +608,7 @@ def delete_old(self, t_old: datetime) -> int:
             raise SnapshotCollectionError(f"Delete of olds snapshots failed: {e}") from e
         return deleted
 
-    def delete_eid(self, eid: Any):
+    def delete_eid(self, eid: AnyEidT):
         """Delete all snapshots of `eid`."""
         try:
             res = self._col().delete_many(self._filter_from_eid(eid))
@@ -655,18 +656,18 @@ def _filter_from_bid(self, b_id: str) -> dict:
 
 class BinaryEidSnapshots(TypedSnapshotCollection, ABC):
     @abc.abstractmethod
-    def _get_packed(self, eid: Any) -> bytes:
+    def _get_packed(self, eid: AnyEidT) -> bytes:
         """Return the packed bytes representation of the given eid."""
 
     def _filter_from_bid(self, b_id: bytes) -> dict:
         eid_bytes = b_id[:-8]
         return {"_id": self._binary_bucket_range(eid_bytes)}
 
-    def _bucket_id(self, eid: Any, ctime: datetime) -> Binary:
+    def _bucket_id(self, eid: AnyEidT, ctime: datetime) -> Binary:
         ts = int(ctime.timestamp())
         return self._pack_binary_bucket_id(self._get_packed(eid), ts)
 
-    def _filter_from_eid(self, eid: Any) -> dict:
+    def _filter_from_eid(self, eid: AnyEidT) -> dict:
         return {"_id": self._binary_bucket_range(self._get_packed(eid))}
 
     def _filter_from_eids(self, eids: Iterable[Any]) -> dict:
@@ -735,7 +736,7 @@ def save_many(self, entity_type: str, snapshots: list[dict], ctime: datetime) ->
         """
         return self[entity_type].save_many(snapshots, ctime)
 
-    def get_latest_one(self, entity_type: str, eid: Any) -> dict:
+    def get_latest_one(self, entity_type: str, eid: AnyEidT) -> dict:
         """Get latest snapshot of given eid.
 
         If doesn't exist, returns {}.
@@ -780,7 +781,7 @@ def get_latest(
     def get_by_eid(
         self,
         entity_type: str,
-        eid: Any,
+        eid: AnyEidT,
         t1: Optional[datetime] = None,
         t2: Optional[datetime] = None,
     ) -> Union[Cursor, CommandCursor]:
@@ -820,7 +821,7 @@ def delete_old(self, t_old: datetime) -> int:
                 raise SnapshotCollectionError(f"Delete of old snapshots failed: {e}") from e
         return deleted_total
 
-    def delete_eid(self, entity_type: str, eid: Any) -> int:
+    def delete_eid(self, entity_type: str, eid: AnyEidT) -> int:
         """Delete snapshots of given `eids`."""
         return self[entity_type].delete_eid(eid)
 
diff --git a/dp3/task_processing/task_hooks.py b/dp3/task_processing/task_hooks.py
index 6fcc127e..0cba01f0 100644
--- a/dp3/task_processing/task_hooks.py
+++ b/dp3/task_processing/task_hooks.py
@@ -4,6 +4,7 @@
 from dp3.common.attrspec import AttrType
 from dp3.common.config import ModelSpec
 from dp3.common.datapoint import DataPointBase
+from dp3.common.datatype import AnyEidT
 from dp3.common.task import DataPointTask, task_context
 from dp3.common.types import EventGroupType
 from dp3.common.utils import get_func_name
@@ -72,7 +73,7 @@ def register(self, hook_type: str, hook: Callable):
 
         self.log.debug(f"Added '{hook_type}' hook: {get_func_name(hook)}")
 
-    def run_allow_creation(self, eid: str, task: DataPointTask):
+    def run_allow_creation(self, eid: AnyEidT, task: DataPointTask):
         for hook in self._allow_creation:
             try:
                 if not hook(eid, task):
@@ -86,7 +87,7 @@ def run_allow_creation(self, eid: str, task: DataPointTask):
 
         return True
 
-    def run_on_creation(self, eid: str, task: DataPointTask):
+    def run_on_creation(self, eid: AnyEidT, task: DataPointTask):
         new_tasks = []
 
         with task_context(self.model_spec):
@@ -150,7 +151,7 @@ def register(self, hook_type: str, hook: Callable):
 
         self.log.debug(f"Added '{hook_type}' hook: {get_func_name(hook)}")
 
-    def run_on_new(self, eid: str, dp: DataPointBase):
+    def run_on_new(self, eid: AnyEidT, dp: DataPointBase):
         new_tasks = []
 
         with task_context(self.model_spec):

From 169ddf346952f3964fb6b66e44c59ff74b46ea76 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 12 Nov 2024 16:43:33 +0100
Subject: [PATCH 11/22] Ensure runtime type checking.

Updated tests in the process.
---
 dp3/api/internal/config.py                    |  6 +-
 dp3/api/internal/models.py                    | 35 ++++++++++-
 dp3/api/routers/entity.py                     | 62 +++++++++++--------
 dp3/common/config.py                          |  9 +++
 dp3/common/entityspec.py                      |  4 ++
 dp3/database/database.py                      | 23 +++++--
 dp3/snapshots/snapshooter.py                  |  5 +-
 tests/modules/test_module.py                  |  2 +-
 .../test_api/test_get_distinct_attr_values.py |  2 +
 tests/test_api/test_get_entity_eid_data.py    |  2 +
 tests/test_api/test_get_entity_eids.py        | 26 +++++---
 tests/test_api/test_list_entities.py          |  6 ++
 tests/test_api/test_snapshots.py              | 62 ++++++++++++-------
 tests/test_config/db_entities/A.yml           |  1 +
 tests/test_config/db_entities/D.yml           |  4 +-
 15 files changed, 177 insertions(+), 72 deletions(-)

diff --git a/dp3/api/internal/config.py b/dp3/api/internal/config.py
index 95b90e98..49b6821e 100644
--- a/dp3/api/internal/config.py
+++ b/dp3/api/internal/config.py
@@ -2,7 +2,11 @@
 import sys
 from enum import Enum
 
-from pydantic import BaseModel, ValidationError, field_validator
+from pydantic import (
+    BaseModel,
+    ValidationError,
+    field_validator,
+)
 
 from dp3.api.internal.dp_logger import DPLogger
 from dp3.common.config import ModelSpec, read_config_dir
diff --git a/dp3/api/internal/models.py b/dp3/api/internal/models.py
index 47500a3b..b60133b3 100644
--- a/dp3/api/internal/models.py
+++ b/dp3/api/internal/models.py
@@ -1,8 +1,9 @@
 from datetime import datetime
-from typing import Annotated, Any, Optional
+from typing import Annotated, Any, Literal, Optional, Union
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, TypeAdapter, create_model, model_validator
 
+from dp3.api.internal.config import MODEL_SPEC
 from dp3.api.internal.helpers import api_to_dp3_datapoint
 from dp3.common.types import T2Datetime
 
@@ -23,7 +24,7 @@ class DataPoint(BaseModel):
     """
 
     type: str
-    id: str
+    id: Any
     attr: str
     v: Any
     t1: Optional[datetime] = None
@@ -40,3 +41,31 @@ def validate_against_attribute(self):
             raise ValueError(f"Missing key: {e}") from e
 
         return self
+
+
+class EntityId(BaseModel):
+    """Dummy model for entity id
+
+    Attributes:
+        type: Entity type
+        id: Entity ID
+    """
+
+    type: Literal["entity_type"]
+    id: Any
+
+
+entity_id_models = []
+for entity_type, entity_spec in MODEL_SPEC.entities.items():
+    dtype = entity_spec.eid_type
+    entity_id_models.append(
+        create_model(
+            f"EntityId{{{entity_type}}}",
+            __base__=BaseModel,
+            type=(Literal[entity_type], Field(..., alias="etype")),
+            id=(dtype, Field(..., alias="eid")),
+        )
+    )
+
+EntityId = Annotated[Union[tuple(entity_id_models)], Field(discriminator="type")]  # noqa: F811
+EntityIdAdapter = TypeAdapter(EntityId)
diff --git a/dp3/api/routers/entity.py b/dp3/api/routers/entity.py
index 237bb3d5..7130455e 100644
--- a/dp3/api/routers/entity.py
+++ b/dp3/api/routers/entity.py
@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, Optional
+from typing import Annotated, Any, Optional
 
 from fastapi import APIRouter, Depends, HTTPException, Query, Request
 from pydantic import Json, NonNegativeInt, ValidationError
@@ -15,9 +15,7 @@
     JsonVal,
 )
 from dp3.api.internal.helpers import api_to_dp3_datapoint
-from dp3.api.internal.models import (
-    DataPoint,
-)
+from dp3.api.internal.models import DataPoint, EntityId, EntityIdAdapter
 from dp3.api.internal.response_models import ErrorResponse, RequestValidationError, SuccessResponse
 from dp3.common.attrspec import AttrType
 from dp3.common.task import DataPointTask, task_context
@@ -31,35 +29,46 @@ async def check_etype(etype: str):
     return etype
 
 
+async def parse_eid(etype: str, eid: str):
+    """Middleware to parse EID"""
+    try:
+        return EntityIdAdapter.validate_python({"etype": etype, "eid": eid})
+    except ValidationError as e:
+        raise RequestValidationError(["path", "eid"], e.errors()[0]["msg"]) from e
+
+
+ParsedEid = Annotated[EntityId, Depends(parse_eid)]
+
+
 def get_eid_master_record_handler(
-    etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: EntityId, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
 ):
     """Handler for getting master record of EID"""
     # TODO: This is probably not the most efficient way. Maybe gather only
     # plain data from master record and then call `get_timeseries_history`
     # for timeseries.
     master_record = DB.get_master_record(
-        etype, eid, projection={"_id": False, "#hash": False, "#min_t2s": False}
+        e.type, e.id, projection={"_id": False, "#hash": False, "#min_t2s": False}
     )
 
-    entity_attribs = MODEL_SPEC.attribs(etype)
+    entity_attribs = MODEL_SPEC.attribs(e.type)
 
     # Get filtered timeseries data
     for attr in master_record:
         # Check for no longer existing attributes
         if attr in entity_attribs and entity_attribs[attr].t == AttrType.TIMESERIES:
             master_record[attr] = DB.get_timeseries_history(
-                etype, attr, eid, t1=date_from, t2=date_to
+                e.type, attr, e.id, t1=date_from, t2=date_to
             )
 
     return master_record
 
 
 def get_eid_snapshots_handler(
-    etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: EntityId, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
 ) -> list[dict[str, Any]]:
     """Handler for getting snapshots of EID"""
-    snapshots = list(DB.snapshots.get_by_eid(etype, eid, t1=date_from, t2=date_to))
+    snapshots = list(DB.snapshots.get_by_eid(e.type, e.id, t1=date_from, t2=date_to))
 
     return snapshots
 
@@ -200,7 +209,7 @@ async def list_entity_type_eids(
 
 @router.get("/{etype}/{eid}")
 async def get_eid_data(
-    etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: ParsedEid, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
 ) -> EntityEidData:
     """Get data of `etype`'s `eid`.
 
@@ -209,8 +218,8 @@ async def get_eid_data(
 
     Combines function of `/{etype}/{eid}/master` and `/{etype}/{eid}/snapshots`.
     """
-    master_record = get_eid_master_record_handler(etype, eid, date_from, date_to)
-    snapshots = get_eid_snapshots_handler(etype, eid, date_from, date_to)
+    master_record = get_eid_master_record_handler(e, date_from, date_to)
+    snapshots = get_eid_snapshots_handler(e, date_from, date_to)
 
     # Whether this eid contains any data
     empty = not master_record and len(snapshots) == 0
@@ -220,24 +229,23 @@ async def get_eid_data(
 
 @router.get("/{etype}/{eid}/master")
 async def get_eid_master_record(
-    etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: ParsedEid, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
 ) -> EntityEidMasterRecord:
     """Get master record of `etype`'s `eid`."""
-    return get_eid_master_record_handler(etype, eid, date_from, date_to)
+    return get_eid_master_record_handler(e, date_from, date_to)
 
 
 @router.get("/{etype}/{eid}/snapshots")
 async def get_eid_snapshots(
-    etype: str, eid: str, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: ParsedEid, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
 ) -> EntityEidSnapshots:
     """Get snapshots of `etype`'s `eid`."""
-    return get_eid_snapshots_handler(etype, eid, date_from, date_to)
+    return get_eid_snapshots_handler(e, date_from, date_to)
 
 
 @router.get("/{etype}/{eid}/get/{attr}")
 async def get_eid_attr_value(
-    etype: str,
-    eid: str,
+    e: ParsedEid,
     attr: str,
     date_from: Optional[datetime] = None,
     date_to: Optional[datetime] = None,
@@ -250,12 +258,14 @@ async def get_eid_attr_value(
     - history: in case of timeseries attribute
     """
     # Check if attribute exists
-    if attr not in MODEL_SPEC.attribs(etype):
+    if attr not in MODEL_SPEC.attribs(e.type):
         raise RequestValidationError(["path", "attr"], f"Attribute '{attr}' doesn't exist")
 
-    value_or_history = DB.get_value_or_history(etype, attr, eid, t1=date_from, t2=date_to)
+    value_or_history = DB.get_value_or_history(e.type, attr, e.id, t1=date_from, t2=date_to)
 
-    return EntityEidAttrValueOrHistory(attr_type=MODEL_SPEC.attr(etype, attr).t, **value_or_history)
+    return EntityEidAttrValueOrHistory(
+        attr_type=MODEL_SPEC.attr(e.type, attr).t, **value_or_history
+    )
 
 
 @router.post("/{etype}/{eid}/set/{attr}")
@@ -318,11 +328,11 @@ async def get_distinct_attribute_values(etype: str, attr: str) -> dict[JsonVal,
 
 
 @router.post("/{etype}/{eid}/ttl")
-async def extend_eid_ttls(etype: str, eid: str, body: dict[str, datetime]) -> SuccessResponse:
+async def extend_eid_ttls(e: ParsedEid, body: dict[str, datetime]) -> SuccessResponse:
     """Extend TTLs of the specified entity"""
     # Construct task
     with task_context(MODEL_SPEC):
-        task = DataPointTask(etype=etype, eid=eid, ttl_tokens=body)
+        task = DataPointTask(etype=e.type, eid=e.id, ttl_tokens=body)
 
     # Push tasks to task queue
     TASK_WRITER.put_task(task, False)
@@ -331,7 +341,7 @@ async def extend_eid_ttls(etype: str, eid: str, body: dict[str, datetime]) -> Su
 
 
 @router.delete("/{etype}/{eid}")
-async def delete_eid_record(etype: str, eid: str) -> SuccessResponse:
+async def delete_eid_record(e: ParsedEid) -> SuccessResponse:
     """Delete the master record and snapshots of the specified entity.
 
     Notice that this does not delete any raw datapoints,
@@ -339,7 +349,7 @@ async def delete_eid_record(etype: str, eid: str) -> SuccessResponse:
     """
     # Create a "delete" task and push it to task queue
     with task_context(MODEL_SPEC):
-        task = DataPointTask(etype=etype, eid=eid, delete=True)
+        task = DataPointTask(etype=e.type, eid=e.id, delete=True)
     TASK_WRITER.put_task(task, False)
 
     return SuccessResponse()
diff --git a/dp3/common/config.py b/dp3/common/config.py
index 3b1f8eef..fcc7bb02 100644
--- a/dp3/common/config.py
+++ b/dp3/common/config.py
@@ -29,6 +29,7 @@
     AttrSpecType,
 )
 from dp3.common.context import entity_context
+from dp3.common.datatype import AnyEidT
 from dp3.common.entityspec import EntitySpec
 
 
@@ -272,6 +273,11 @@ def _validate_config(cls, data: Any) -> Any:
                 raise ValueError(f"Invalid entity specification for '{entity_id}'")
             if not isinstance(entity_dict["entity"], EntitySpec):
                 entity_dict["entity"] = EntitySpec.model_validate(entity_dict["entity"])
+            if entity_id != entity_dict["entity"].id:
+                spec_id = entity_dict["entity"].id
+                raise ValueError(
+                    f"Entity id '{entity_id}' does not match entity id '{spec_id}' in spec."
+                )
             data["entities"][entity_id] = entity_dict["entity"]
 
         # Then validate all attributes
@@ -387,6 +393,9 @@ def items(self):
     def keys(self):
         return self.config.keys()
 
+    def parse_eid(self, entity_type: str, eid: str) -> AnyEidT:
+        return self.entities[entity_type].validate_eid(eid)
+
     def __contains__(self, item):
         return item in self.config
 
diff --git a/dp3/common/entityspec.py b/dp3/common/entityspec.py
index fe00cdb9..f0e0fba7 100644
--- a/dp3/common/entityspec.py
+++ b/dp3/common/entityspec.py
@@ -78,3 +78,7 @@ def fill_validator(self):
 
     def validate_eid(self, v):
         return self._eid_validator(v)
+
+    @property
+    def eid_type(self):
+        return self._eid_validator
diff --git a/dp3/database/database.py b/dp3/database/database.py
index 4863af7e..2cc90135 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -339,6 +339,18 @@ def _assert_etype_exists(self, etype: str):
         if etype not in self._db_schema_config.entities:
             raise DatabaseError(f"Entity '{etype}' does not exist")
 
+    def _assert_eid_correct_dtype(self, etype: str, eid: AnyEidT):
+        """Asserts `etype` existence.
+
+        If doesn't exist, raises `DatabaseError`.
+        """
+        entity_spec = self._db_schema_config.entities[etype]
+        if not isinstance(eid, entity_spec.eid_type):
+            raise DatabaseError(
+                f"Invalid entity ID type: {type(eid)}, expected {entity_spec.eid_type}"
+            )
+        return entity_spec
+
     def update_schema(self):
         """
         Checks whether schema saved in database is up-to-date and updates it if necessary.
@@ -374,6 +386,7 @@ def insert_datapoints(
 
         # Check `etype`
         self._assert_etype_exists(etype)
+        self._assert_eid_correct_dtype(etype, eid)
 
         # Insert raw datapoints
         dps_dicts = [dp.model_dump(exclude={"attr_type"}) for dp in dps]
@@ -732,12 +745,12 @@ def get_master_record(self, etype: str, eid: AnyEidT, **kwargs) -> dict:
 
         If doesn't exist, returns {}.
         """
-        # Check `etype`
         self._assert_etype_exists(etype)
-        entity_spec = self._db_schema_config.entities[etype]
+        self._assert_eid_correct_dtype(etype, eid)
+
         master_col = self._master_col(etype)
 
-        return master_col.find_one({"_id": entity_spec.validate_eid(eid)}, **kwargs) or {}
+        return master_col.find_one({"_id": eid}, **kwargs) or {}
 
     def ekey_exists(self, etype: str, eid: AnyEidT) -> bool:
         """Checks whether master record for etype/eid exists"""
@@ -745,7 +758,6 @@ def ekey_exists(self, etype: str, eid: AnyEidT) -> bool:
 
     def get_master_records(self, etype: str, **kwargs) -> Cursor:
         """Get cursor to current master records of etype."""
-        # Check `etype`
         self._assert_etype_exists(etype)
 
         master_col = self._master_col(etype)
@@ -781,8 +793,8 @@ def get_value_or_history(
 
         Returns dict with two keys: `current_value` and `history` (list of values).
         """
-        # Check `etype`
         self._assert_etype_exists(etype)
+        self._assert_eid_correct_dtype(etype, eid)
 
         attr_spec = self._db_schema_config.attr(etype, attr_name)
 
@@ -806,7 +818,6 @@ def get_value_or_history(
 
     def estimate_count_eids(self, etype: str) -> int:
         """Estimates count of `eid`s in given `etype`"""
-        # Check `etype`
         self._assert_etype_exists(etype)
 
         master_col = self._master_col(etype)
diff --git a/dp3/snapshots/snapshooter.py b/dp3/snapshots/snapshooter.py
index 434220f2..3cc7da95 100644
--- a/dp3/snapshots/snapshooter.py
+++ b/dp3/snapshots/snapshooter.py
@@ -312,9 +312,10 @@ def get_linked_entities(self, time: datetime, cached_linked_entities: list[tuple
         entity_to_component = {}
         linked_components = []
         for etype, eid in cached_linked_entities:
+            typed_eid = self.model_spec.parse_eid(etype, eid)
             master_record = self.db.get_master_record(
-                etype, eid, projection=self.entity_relation_attrs[etype]
-            ) or {"_id": eid}
+                etype, typed_eid, projection=self.entity_relation_attrs[etype]
+            ) or {"_id": typed_eid}
 
             if not self.config.keep_empty and len(master_record) == 1:
                 continue
diff --git a/tests/modules/test_module.py b/tests/modules/test_module.py
index 3d90b737..3640a3e4 100644
--- a/tests/modules/test_module.py
+++ b/tests/modules/test_module.py
@@ -7,7 +7,7 @@
 
 
 def copy_linked(_: str, record: dict, link: str, attr: str):
-    record[attr] = record.get(link, {}).get(attr, None)
+    record[attr] = record.get(link, {}).get("record", {}).get(attr, None)
 
 
 def modify_value(_: str, record: dict, attr: str, value):
diff --git a/tests/test_api/test_get_distinct_attr_values.py b/tests/test_api/test_get_distinct_attr_values.py
index 916134de..a3a6e55c 100644
--- a/tests/test_api/test_get_distinct_attr_values.py
+++ b/tests/test_api/test_get_distinct_attr_values.py
@@ -34,6 +34,8 @@ def setUpClass(cls) -> None:
                 ]
             )
             print(res.content.decode("utf-8"), file=sys.stderr)
+            if res.status_code != 200:
+                raise Exception(f"Failed to push datapoints: {res.status_code}")
         sleep(9)
 
         # Make snapshots
diff --git a/tests/test_api/test_get_entity_eid_data.py b/tests/test_api/test_get_entity_eid_data.py
index 024ea9b4..13a7c34c 100644
--- a/tests/test_api/test_get_entity_eid_data.py
+++ b/tests/test_api/test_get_entity_eid_data.py
@@ -38,6 +38,8 @@ def setUp(self) -> None:
         print(dps, file=sys.stderr)
         res = self.push_datapoints(dps)
         print(res.content.decode("utf-8"), file=sys.stderr)
+        if res.status_code != 200:
+            raise Exception(f"Failed to push datapoints: {res.status_code}")
 
     def test_get_entity_data(self):
         data = self.query_expected_value(
diff --git a/tests/test_api/test_get_entity_eids.py b/tests/test_api/test_get_entity_eids.py
index dd430d10..bc5fde46 100644
--- a/tests/test_api/test_get_entity_eids.py
+++ b/tests/test_api/test_get_entity_eids.py
@@ -14,9 +14,11 @@ def setUpClass(cls) -> None:
         dp_base = {"src": "setup@test", "attr": "data1", "type": "A"}
         for i in range(0, 100, 20):
             res = cls.push_datapoints(
-                [{**dp_base, "id": f"A{i}", "v": f"v{i}"} for i in range(i, i + 20)]
+                [{**dp_base, "id": i, "v": f"v{i}"} for i in range(i, i + 20)]
             )
-            print(res.content.decode("utf-8"), file=sys.stderr)
+            if res.status_code != 200:
+                print(json.dumps(res.json(), indent=2), file=sys.stderr)
+                raise Exception(f"Failed to push datapoints: {res.status_code}")
         sleep(8)
         cls.get_request("control/make_snapshots")
         sleep(6)
@@ -26,7 +28,7 @@ def test_get_entity_eids(self):
         self.assertEqual(20, len(eids.data))
 
     def test_get_entity_eids_pagination(self):
-        expected_eids = {f"A{i}" for i in range(0, 100)}
+        expected_eids = set(range(0, 100))
         received_eids = set()
 
         for i in range(0, 100, 10):
@@ -34,19 +36,27 @@ def test_get_entity_eids_pagination(self):
             self.assertEqual(10, len(eids.data), f"Failed at {i}")
             received_eids.update(x["eid"] for x in eids.data)
 
-        eids = self.get_entity_data("entity/A", EntityEidList, skip=101, limit=20)
+        eids = self.get_entity_data("entity/A", EntityEidList, skip=102, limit=20)
         self.assertEqual(0, len(eids.data))
         self.assertSetEqual(expected_eids, received_eids)
 
     def test_get_entity_eids_generic_filter(self):
         eids = self.get_entity_data(
-            "entity/A", EntityEidList, generic_filter=json.dumps({"last.eid": "A0"})
+            "entity/A", EntityEidList, generic_filter=json.dumps({"last.eid": 0})
         )
         self.assertEqual(1, len(eids.data))
-        self.assertEqual("A0", eids.data[0]["eid"])
+        self.assertEqual(0, eids.data[0]["eid"])
 
-    def test_get_entity_eids_fulltext_filters_eid(self):
+    def test_get_entity_eids_generic_filters_eid(self):
         eids = self.get_entity_data(
-            "entity/A", EntityEidList, fulltext_filters=json.dumps({"eid": "A5.*"})
+            "entity/A",
+            EntityEidList,
+            generic_filter=json.dumps(
+                {"$or": [{"last.eid": 5}, {"last.eid": {"$gte": 50, "$lt": 60}}]}
+            ),
         )
         self.assertEqual(11, len(eids.data))  # A5, A50 ... A59
+        self.assertEqual(
+            {5, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59},
+            {x["eid"] for x in eids.data},
+        )
diff --git a/tests/test_api/test_list_entities.py b/tests/test_api/test_list_entities.py
index 80ed388a..21ffa2b1 100644
--- a/tests/test_api/test_list_entities.py
+++ b/tests/test_api/test_list_entities.py
@@ -24,6 +24,12 @@ def test_get_entities(self):
             entities[key] = EntitySpec.model_validate(
                 {k: v for k, v in item.items() if k in entity_keys} | {"snapshot": True}
             )
+            for k in entity_keys:
+                self.assertEqual(
+                    entities[key].__getattribute__(k),
+                    common.MODEL_SPEC.entities[key].__getattribute__(k),
+                    f"Mismatch in entity '{key}' data at key '{k}'",
+                )
 
         for key, item in received.items():
             with entity_context(entities[key], entities):
diff --git a/tests/test_api/test_snapshots.py b/tests/test_api/test_snapshots.py
index b78f7162..8bdd4119 100644
--- a/tests/test_api/test_snapshots.py
+++ b/tests/test_api/test_snapshots.py
@@ -1,4 +1,6 @@
 import datetime
+import json
+import sys
 from time import sleep
 
 import common
@@ -13,35 +15,47 @@ def setUpClass(cls) -> None:
         t1 = (now - datetime.timedelta(minutes=30)).strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-4]
         t2 = (now + datetime.timedelta(minutes=30)).strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-4]
 
+        def make_dp(type, id, attr, v, time=False):
+            base = {"type": type, "id": id, "attr": attr, "v": v}
+            if time:
+                base |= {"t1": t1, "t2": t2, "c": 1.0}
+            return base
+
         entity_datapoints = [
-            # For test_linked_entities_loaded
-            {"type": "B", "id": "b1", "attr": "as", "v": "a1", "t1": t1, "t2": t2, "c": 1.0},
-            {"type": "B", "id": "b1", "attr": "data1", "v": "initb"},
-            {"type": "B", "id": "b1", "attr": "data2", "v": "initb"},
-            {"type": "A", "id": "a1", "attr": "bs", "v": "b1", "t1": t1, "t2": t2, "c": 1.0},
-            {"type": "A", "id": "a1", "attr": "data1", "v": "inita"},
-            {"type": "A", "id": "a1", "attr": "data2", "v": "inita"},
-            # For test_weak_component_detection
-            {"type": "B", "id": "b2", "attr": "as", "v": "a1", "t1": t1, "t2": t2, "c": 1.0},
-            {"type": "B", "id": "b2", "attr": "data1", "v": "initb"},
-            {"type": "B", "id": "b2", "attr": "data2", "v": "initb"},
-            {"type": "A", "id": "a2", "attr": "data1", "v": "inita"},
-            {"type": "A", "id": "a2", "attr": "data2", "v": "inita"},
-            # For test_hook_dependency_value_forwarding
-            {"type": "D", "id": "d1", "attr": "cs", "v": "c1", "t1": t1, "t2": t2, "c": 1.0},
-            {"type": "D", "id": "d1", "attr": "data1", "v": "initb"},
-            {"type": "D", "id": "d1", "attr": "data2", "v": "initb"},
-            {"type": "C", "id": "c1", "attr": "ds", "v": "d1", "t1": t1, "t2": t2, "c": 1.0},
-            {"type": "C", "id": "c1", "attr": "data1", "v": "inita"},
-            {"type": "C", "id": "c1", "attr": "data2", "v": "inita"},
+            # For test_linked_entities_loaded (B-"b1" <-> A-420)
+            make_dp("B", "b1", "as", {"eid": 420}, time=True),
+            make_dp("B", "b1", "data1", "initb"),
+            make_dp("B", "b1", "data2", "initb"),
+            make_dp("A", 420, "bs", {"eid": "b1"}, time=True),
+            make_dp("A", 420, "data1", "inita"),
+            make_dp("A", 420, "data2", "inita"),
+            # For test_weak_component_detection (B-"b2" -> A-422)
+            make_dp("B", "b2", "as", {"eid": 422}, time=True),
+            make_dp("B", "b2", "data1", "initb"),
+            make_dp("B", "b2", "data2", "initb"),
+            make_dp("A", 422, "data1", "inita"),
+            make_dp("A", 422, "data2", "inita"),
+            # For test_hook_dependency_value_forwarding (D-"d1" <-> C-"c1")
+            make_dp("D", "d1", "cs", {"eid": "c1"}, time=True),
+            make_dp("D", "d1", "data1", "initb"),
+            make_dp("D", "d1", "data2", "initb"),
+            make_dp("C", "c1", "ds", {"eid": "d1"}, time=True),
+            make_dp("C", "c1", "data1", "inita"),
+            make_dp("C", "c1", "data2", "inita"),
         ]
-        cls.push_datapoints(entity_datapoints)
+        res = cls.push_datapoints(entity_datapoints)
+        if res.status_code != 200:
+            print(json.dumps(res.json(), indent=2), sys.stderr)
+            raise Exception(f"Failed to push datapoints: {res.status_code}")
+
+        sleep(3)
         cls.get_request("control/make_snapshots")
         sleep(3)
 
     def test_linked_entities_loaded(self):
-        for entity, eid in [("A", "a1"), ("B", "b1")]:
+        for entity, eid in [("A", 420), ("B", "b1")]:
             data = self.get_entity_data(f"entity/{entity}/{eid}", EntityEidData)
+            self.assertGreater(len(data.snapshots), 0)
             for snapshot in data.snapshots:
                 self.assertEqual(snapshot["data1"], "initb")
                 self.assertEqual(snapshot["data2"], "inita")
@@ -51,8 +65,9 @@ def test_weak_component_detection(self):
         Remove link from the first entity to be processed, so that an elementary BFS starting
         from this entity is not enough to load the entire weakly connected component.
         """
-        for entity, eid in [("A", "a2"), ("B", "b2")]:
+        for entity, eid in [("B", "b2")]:
             data = self.get_entity_data(f"entity/{entity}/{eid}", EntityEidData)
+            self.assertGreater(len(data.snapshots), 0)
             for snapshot in data.snapshots:
                 self.assertEqual(snapshot["data1"], "initb")
                 self.assertEqual(snapshot["data2"], "inita")
@@ -60,6 +75,7 @@ def test_weak_component_detection(self):
     def test_hook_dependency_value_forwarding(self):
         for entity, eid in [("D", "d1"), ("C", "c1")]:
             data = self.get_entity_data(f"entity/{entity}/{eid}", EntityEidData)
+            self.assertGreater(len(data.snapshots), 0)
             for snapshot in data.snapshots:
                 self.assertEqual(snapshot["data1"], "modifd")
                 self.assertEqual(snapshot["data2"], "modifc")
diff --git a/tests/test_config/db_entities/A.yml b/tests/test_config/db_entities/A.yml
index 05a11a40..d602cd2f 100644
--- a/tests/test_config/db_entities/A.yml
+++ b/tests/test_config/db_entities/A.yml
@@ -1,5 +1,6 @@
 entity:
   id: A
+  id_data_type: int
   name: test entity A
   snapshot: true
 
diff --git a/tests/test_config/db_entities/D.yml b/tests/test_config/db_entities/D.yml
index e9f2af33..15485aaf 100644
--- a/tests/test_config/db_entities/D.yml
+++ b/tests/test_config/db_entities/D.yml
@@ -1,6 +1,6 @@
 entity:
-  id: C
-  name: test entity C
+  id: D
+  name: test entity D
   snapshot: true
 
 attribs:

From 50ab7ce8c4d3a55b21c9e86aabd6489e5cbd654e Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Wed, 13 Nov 2024 16:45:30 +0100
Subject: [PATCH 12/22] DB: Fix deleting and Link cache access.

---
 dp3/api/internal/entity_response_models.py    |   2 +-
 dp3/common/task.py                            |  40 ++
 dp3/core/collector.py                         |   4 +-
 dp3/core/link_manager.py                      |   7 +-
 dp3/database/database.py                      |  72 ++--
 dp3/snapshots/snapshooter.py                  |  22 +-
 tests/test_api/common.py                      |  19 +-
 tests/test_api/test_garbage_collector.py      | 358 ++++++++++++++++++
 tests/test_config/db_entities/del_anchor.yml  |  26 ++
 .../db_entities/del_anchor_int.yml            |  26 ++
 tests/test_config/db_entities/del_ttl.yml     |  23 ++
 tests/test_config/db_entities/del_ttl_ip.yml  |  23 ++
 tests/test_config/db_entities/del_weak.yml    |  20 +
 .../test_config/db_entities/del_weak_mac.yml  |  20 +
 tests/test_config/garbage_collector.yml       |   5 +
 15 files changed, 620 insertions(+), 47 deletions(-)
 create mode 100644 tests/test_api/test_garbage_collector.py
 create mode 100644 tests/test_config/db_entities/del_anchor.yml
 create mode 100644 tests/test_config/db_entities/del_anchor_int.yml
 create mode 100644 tests/test_config/db_entities/del_ttl.yml
 create mode 100644 tests/test_config/db_entities/del_ttl_ip.yml
 create mode 100644 tests/test_config/db_entities/del_weak.yml
 create mode 100644 tests/test_config/db_entities/del_weak_mac.yml
 create mode 100644 tests/test_config/garbage_collector.yml

diff --git a/dp3/api/internal/entity_response_models.py b/dp3/api/internal/entity_response_models.py
index 51cbb080..0e480d9d 100644
--- a/dp3/api/internal/entity_response_models.py
+++ b/dp3/api/internal/entity_response_models.py
@@ -73,7 +73,7 @@ class EntityEidAttrValueOrHistory(BaseModel):
     """
 
     attr_type: AttrType
-    current_value: JsonVal = None
+    current_value: Dp3Val = None
     history: HistoryVal = []
 
 
diff --git a/dp3/common/task.py b/dp3/common/task.py
index 321b8129..308de043 100644
--- a/dp3/common/task.py
+++ b/dp3/common/task.py
@@ -15,6 +15,7 @@
     Discriminator,
     PlainSerializer,
     Tag,
+    TypeAdapter,
     ValidationError,
     field_validator,
 )
@@ -22,6 +23,7 @@
 
 from dp3.common.config import ModelSpec, entity_type_context, get_entity_type_context
 from dp3.common.datapoint import DataPointBase, to_json_friendly
+from dp3.common.datatype import AnyEidT
 from dp3.common.mac_address import MACAddress
 
 _init_context_var = ContextVar("_init_context_var", default=None)
@@ -239,3 +241,41 @@ def json_validator(serialized: Union[str, bytes]) -> Snapshot:
                 return Snapshot.model_validate_json(serialized)
 
         return json_validator
+
+
+EntityTuplesAdapter = TypeAdapter(list[EntityTuple])
+
+
+def validate_entities(model_spec: ModelSpec, v: list[tuple[str, str]]) -> list[EntityTuple]:
+    with entity_type_context(model_spec):
+        return EntityTuplesAdapter.validate_python(v)
+
+
+def parse_eids_from_cache(model_spec: ModelSpec, link_entity_entries: list[str]) -> list[AnyEidT]:
+    """Parses entity IDs from the "Link" cache.
+
+    Args:
+        model_spec: Model specification.
+        link_entity_entries: List of entity entries from the cache.
+
+    Returns:
+        List of entity IDs.
+    """
+    raw = [entry.split("#", maxsplit=1) for entry in link_entity_entries]
+    return [eid for etype, eid in validate_entities(model_spec, raw)]
+
+
+def parse_eid_tuples_from_cache(
+    model_spec: ModelSpec, link_entity_entries: list[str]
+) -> list[tuple[str, AnyEidT]]:
+    """Parses entity IDs from the "Link" cache.
+
+    Args:
+        model_spec: Model specification.
+        link_entity_entries: List of entity entries from the cache.
+
+    Returns:
+        List of tuples (entity type, entity ID).
+    """
+    raw = [entry.split("#", maxsplit=1) for entry in link_entity_entries]
+    return validate_entities(model_spec, raw)
diff --git a/dp3/core/collector.py b/dp3/core/collector.py
index 08a60fce..bb5a2df9 100644
--- a/dp3/core/collector.py
+++ b/dp3/core/collector.py
@@ -14,7 +14,7 @@
 from dp3.common.config import CronExpression, PlatformConfig
 from dp3.common.datapoint import DataPointBase, DataPointObservationsBase, DataPointTimeseriesBase
 from dp3.common.datatype import AnyEidT
-from dp3.common.task import DataPointTask
+from dp3.common.task import DataPointTask, parse_eids_from_cache
 from dp3.database.database import EntityDatabase
 
 DB_SEND_CHUNK = 1000
@@ -174,7 +174,7 @@ def collect_weak(self, etype: str):
                 {"$group": {"_id": "$to"}},
             ]
         )
-        have_references = [doc["_id"].split("#", maxsplit=1)[1] for doc in aggregated]
+        have_references = parse_eids_from_cache(self.model_spec, [doc["_id"] for doc in aggregated])
         entities += len(have_references)
 
         to_delete = []
diff --git a/dp3/core/link_manager.py b/dp3/core/link_manager.py
index 4d6d9726..4750b467 100644
--- a/dp3/core/link_manager.py
+++ b/dp3/core/link_manager.py
@@ -13,6 +13,7 @@
 from dp3.common.config import PlatformConfig
 from dp3.common.datapoint import DataPointBase, DataPointObservationsBase
 from dp3.common.datatype import AnyEidT
+from dp3.common.task import parse_eids_from_cache
 from dp3.database.database import EntityDatabase
 
 
@@ -133,7 +134,7 @@ def remove_link_cache_of_deleted(self, etype: str, eid: AnyEidT):
         # Delete the retrieved references
         for link in links:
             etype_from, attr = link["_id"].split("#", maxsplit=1)
-            eids_from = [l_eid.split("#", maxsplit=1)[1] for l_eid in link["eids_from"]]
+            eids_from = parse_eids_from_cache(self.model_spec, link["eids_from"])
             self.db.delete_link_dps(etype_from, eids_from, attr, eid)
 
         # Delete from cache
@@ -167,8 +168,8 @@ def remove_link_cache_of_many_deleted(self, etype: str, eids: list[AnyEidT]):
             etype_from, attr = link["_id"].split("#", maxsplit=1)
             etypes_from.append(etype_from)
             attrs.append(attr)
-            affected_eids.append([l_eid.split("#", maxsplit=1)[1] for l_eid in link["eids_from"]])
-            eids_to.append([l_eid.split("#", maxsplit=1)[1] for l_eid in link["eids_to"]])
+            affected_eids.append(parse_eids_from_cache(self.model_spec, link["eids_from"]))
+            eids_to.append(parse_eids_from_cache(self.model_spec, link["eids_to"]))
 
         self.db.delete_many_link_dps(etypes_from, affected_eids, attrs, eids_to)
 
diff --git a/dp3/database/database.py b/dp3/database/database.py
index 2cc90135..35b17302 100644
--- a/dp3/database/database.py
+++ b/dp3/database/database.py
@@ -160,7 +160,7 @@ def stop(self) -> None:
         self._push_master()
 
     def register_on_entity_delete(
-        self, f_one: Callable[[str, str], None], f_many: Callable[[str, list[str]], None]
+        self, f_one: Callable[[str, AnyEidT], None], f_many: Callable[[str, list[AnyEidT]], None]
     ):
         """Registers function to be called when entity is forcibly deleted."""
         self._on_entity_delete_one.append(f_one)
@@ -473,13 +473,18 @@ def _push_master(self):
                 updates = [
                     UpdateOne(
                         {"_id": eid},
-                        {
-                            "$push": {
-                                attr: push_dps[0] if len(push_dps) == 1 else {"$each": push_dps}
-                                for attr, push_dps in changes["pushes"].items()
-                            },
-                            "$set": changes["$set"],
-                        },
+                        (
+                            {
+                                "$push": {
+                                    attr: push_dps[0] if len(push_dps) == 1 else {"$each": push_dps}
+                                    for attr, push_dps in changes.get("pushes", {}).items()
+                                }
+                            }
+                            if "pushes" in changes
+                            else {}
+                        )
+                        | ({"$set": changes["$set"]} if "$set" in changes else {})
+                        | ({"$max": changes["$max"]} if "$max" in changes else {}),
                         upsert=True,
                     )
                     for eid, changes in master_changes.items()
@@ -503,7 +508,7 @@ def _push_master(self):
                     f"Update of master records failed: {e}\n{master_changes}"
                 ) from e
 
-    def update_master_records(self, etype: str, eids: list[str], records: list[dict]) -> None:
+    def update_master_records(self, etype: str, eids: list[AnyEidT], records: list[dict]) -> None:
         """Replace master records of `etype`:`eid` with the provided `records`.
 
         Raises DatabaseError when update fails.
@@ -525,22 +530,22 @@ def update_master_records(self, etype: str, eids: list[str], records: list[dict]
 
     def extend_ttl(self, etype: str, eid: AnyEidT, ttl_tokens: dict[str, datetime]):
         """Extends TTL of given `etype`:`eid` by `ttl_tokens`."""
-        master_col = self._master_col(etype)
-        try:
-            master_col.update_one(
-                {"_id": eid},
-                {
-                    "$max": {
-                        f"#ttl.{token_name}": token_value
-                        for token_name, token_value in ttl_tokens.items()
-                    }
-                },
-            )
-            self.log.debug("Updated TTL of %s: %s.", etype, eid)
-        except Exception as e:
-            raise DatabaseError(f"TTL update failed: {e} ({ttl_tokens})") from e
+        extensions = {
+            f"#ttl.{token_name}": token_value for token_name, token_value in ttl_tokens.items()
+        }
+        with self._master_buffer_locks[etype]:
+            buf = self._master_buffers[etype]
+            if eid in buf:
+                if "$max" in buf[eid]:
+                    for ttl_name, this_val in extensions.items():
+                        curr_val = buf[eid]["$max"].get(ttl_name, datetime.min)
+                        buf[eid]["$max"][ttl_name] = max(curr_val, this_val)
+                else:
+                    self._master_buffers[etype][eid]["$max"] = extensions
+            else:
+                self._master_buffers[etype][eid] = {"$max": extensions}
 
-    def remove_expired_ttls(self, etype: str, expired_eid_ttls: dict[str, list[str]]):
+    def remove_expired_ttls(self, etype: str, expired_eid_ttls: dict[AnyEidT, list[str]]):
         """Removes expired TTL of given `etype`:`eid`."""
         master_col = self._master_col(etype)
         try:
@@ -562,9 +567,13 @@ def remove_expired_ttls(self, etype: str, expired_eid_ttls: dict[str, list[str]]
         except Exception as e:
             raise DatabaseError(f"TTL update failed: {e}") from e
 
-    def delete_eids(self, etype: str, eids: list[str]):
+    def delete_eids(self, etype: str, eids: list[AnyEidT]):
         """Delete master record and all snapshots of `etype`:`eids`."""
         master_col = self._master_col(etype)
+        with self._master_buffer_locks[etype]:
+            for eid in eids:
+                if eid in self._master_buffers[etype]:
+                    del self._master_buffers[etype][eid]
         try:
             res = master_col.delete_many({"_id": {"$in": eids}})
             self.log.debug(
@@ -584,6 +593,9 @@ def delete_eids(self, etype: str, eids: list[str]):
     def delete_eid(self, etype: str, eid: AnyEidT):
         """Delete master record and all snapshots of `etype`:`eid`."""
         master_col = self._master_col(etype)
+        with self._master_buffer_locks[etype]:
+            if eid in self._master_buffers[etype]:
+                del self._master_buffers[etype][eid]
         try:
             master_col.delete_one({"_id": eid})
             self.log.debug("Deleted master record of %s/%s.", etype, eid)
@@ -688,7 +700,7 @@ def delete_old_dps(self, etype: str, attr_name: str, t_old: datetime) -> UpdateR
             raise DatabaseError(f"Delete of old datapoints failed: {e}") from e
 
     def delete_link_dps(
-        self, etype: str, affected_eids: list[str], attr_name: str, eid_to: str
+        self, etype: str, affected_eids: list[AnyEidT], attr_name: str, eid_to: AnyEidT
     ) -> None:
         """Delete link datapoints from master collection.
 
@@ -712,9 +724,9 @@ def delete_link_dps(
     def delete_many_link_dps(
         self,
         etypes: list[str],
-        affected_eids: list[list[str]],
+        affected_eids: list[list[AnyEidT]],
         attr_names: list[str],
-        eids_to: list[list[str]],
+        eids_to: list[list[AnyEidT]],
     ) -> None:
         """Delete link datapoints from master collection.
 
@@ -754,7 +766,9 @@ def get_master_record(self, etype: str, eid: AnyEidT, **kwargs) -> dict:
 
     def ekey_exists(self, etype: str, eid: AnyEidT) -> bool:
         """Checks whether master record for etype/eid exists"""
-        return bool(self.get_master_record(etype, eid))
+        with self._master_buffer_locks[etype]:
+            in_cache = eid in self._master_buffers[etype]
+        return in_cache or bool(self.get_master_record(etype, eid, projection={"_id": 1}))
 
     def get_master_records(self, etype: str, **kwargs) -> Cursor:
         """Get cursor to current master records of etype."""
diff --git a/dp3/snapshots/snapshooter.py b/dp3/snapshots/snapshooter.py
index 3cc7da95..f7aeaa0f 100644
--- a/dp3/snapshots/snapshooter.py
+++ b/dp3/snapshots/snapshooter.py
@@ -34,7 +34,14 @@
 )
 from dp3.common.config import CronExpression, PlatformConfig, entity_type_context
 from dp3.common.scheduler import Scheduler
-from dp3.common.task import DataPointTask, Snapshot, SnapshotMessageType, task_context
+from dp3.common.task import (
+    DataPointTask,
+    Snapshot,
+    SnapshotMessageType,
+    parse_eids_from_cache,
+    task_context,
+    parse_eid_tuples_from_cache,
+)
 from dp3.common.types import EventGroupType
 from dp3.common.utils import get_func_name
 from dp3.database.database import EntityDatabase
@@ -299,12 +306,12 @@ def get_cached_link_entity_ids(self):
         result = self.cache.aggregate(
             [{"$match": {"using_attr": {"$in": used}}}, {"$group": {"_id": "$from"}}]
         )
-        links_from = {tuple(doc["_id"].split("#", maxsplit=1)) for doc in result}
+        links_from = parse_eid_tuples_from_cache(self.model_spec, [doc["_id"] for doc in result])
         result = self.cache.aggregate(
             [{"$match": {"using_attr": {"$in": used}}}, {"$group": {"_id": "$to"}}]
         )
-        links_to = {tuple(doc["_id"].split("#", maxsplit=1)) for doc in result}
-        return list(links_from | links_to)
+        links_to = parse_eid_tuples_from_cache(self.model_spec, [doc["_id"] for doc in result])
+        return list(set(links_from) | set(links_to))
 
     def get_linked_entities(self, time: datetime, cached_linked_entities: list[tuple[str, str]]):
         """Get weakly connected components from entity graph."""
@@ -479,9 +486,10 @@ def add_mirrored_links(self, entity_type: str, values: dict):
                         {"$project": {"from": 1}},
                     ]
                 )
-                values[mirror_name] = [
-                    {"eid": source["from"].split("#", maxsplit=1)[1]} for source in link_sources
-                ]
+                parsed = parse_eids_from_cache(
+                    self.model_spec, [doc["from"] for doc in link_sources]
+                )
+                values[mirror_name] = [{"eid": eid} for eid in parsed]
 
     def make_snapshot(self, task: Snapshot):
         """
diff --git a/tests/test_api/common.py b/tests/test_api/common.py
index 348c4139..671dc74f 100644
--- a/tests/test_api/common.py
+++ b/tests/test_api/common.py
@@ -103,16 +103,25 @@ def post_request(path, **kwargs) -> requests.Response:
         )
 
     def query_expected_value(
-        self, query: Callable[[], A], assertion: Callable[[A], bool], attempts=5, delay_s=0.5
+        self,
+        query: Callable[[], A],
+        assertion: Callable[[A], bool],
+        attempts=25,
+        delay_s=0.1,
+        msg: str = None,
     ):
         payload: A = None
-        for _ in range(attempts):
-            time.sleep(delay_s)
+        for i in range(attempts):
+            if i > 0:
+                time.sleep(delay_s)
             payload = query()
-            print(payload, file=sys.stderr)
             if assertion(payload):
                 break
-        self.assertTrue(assertion(payload))
+            print(payload, file=sys.stderr)
+        if msg:
+            self.assertTrue(assertion(payload), msg=msg)
+        else:
+            self.assertTrue(assertion(payload))
         return payload
 
     @classmethod
diff --git a/tests/test_api/test_garbage_collector.py b/tests/test_api/test_garbage_collector.py
new file mode 100644
index 00000000..402584b6
--- /dev/null
+++ b/tests/test_api/test_garbage_collector.py
@@ -0,0 +1,358 @@
+import json
+import sys
+from time import sleep
+from typing import Any
+
+import common
+import requests
+
+from dp3.api.internal.entity_response_models import EntityEidAttrValueOrHistory
+
+
+def make_dps(entity_dict: dict[str, dict[str, dict[str, Any]]]) -> list[dict[str, Any]]:
+    """Take a db "inventory" and produce a list of datapoints to create it."""
+    dps = []
+    for entity, e_dict in entity_dict.items():
+        for eid, record in e_dict.items():
+            for attr, value in record.items():
+                dps.append(
+                    {
+                        "type": entity,
+                        "id": eid,
+                        "attr": attr,
+                        "v": value,
+                    }
+                )
+    return dps
+
+
+GET_ATTR = "entity/{entity}/{eid}/get/{attr}"
+
+
+def ref(link: Any):
+    return {"eid": link}
+
+
+class GarbageCollectorTest(common.APITest):
+    test1_data = "test1_data"
+    test2_data = "test2_data"
+    test3_data = "test3_data"
+    test4_data = "test4_data"
+    test5_data = "test5_data"
+    test6_data = "test6_data"
+    test7_data = "test7_data"
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        d1 = "data1"  # Data attribute for all entities
+        r = "ref"  # Ref attribute to weak entity for all entities
+
+        # Simple delete test - create isolated anchor and delete via API
+        del_anchor = {"a1": {d1: cls.test1_data}}
+        del_anchor_int = {1: {d1: cls.test1_data}}
+
+        # Simple weak test - create an anchor and a weak entity
+        del_anchor["a2"] = {d1: cls.test2_data, r: ref("w1")}
+        del_weak = {"w1": {d1: cls.test2_data}}
+
+        del_anchor_int[2] = {d1: cls.test2_data, r: ref("00:00:00:00:00:01")}
+        del_weak_mac = {"00:00:00:00:00:01": {d1: cls.test2_data}}
+
+        # Weak chains test - an anchor and two weak entity chain
+        del_anchor["a3"] = {d1: cls.test4_data, r: ref("w2_1")}
+        del_weak["w2_1"] = {d1: cls.test4_data, r: ref("w2_2")}
+        del_weak["w2_2"] = {d1: cls.test4_data}
+
+        del_anchor_int[3] = {d1: cls.test4_data, r: ref("00:00:00:00:00:21")}
+        del_weak_mac["00:00:00:00:00:21"] = {d1: cls.test4_data, r: ref("00:00:00:00:00:22")}
+        del_weak_mac["00:00:00:00:00:22"] = {d1: cls.test4_data}
+
+        # Weak multiple ref test - two anchors, one weak entity
+        del_anchor["a4"] = {d1: cls.test5_data, r: ref("w3")}
+        del_anchor["a5"] = {d1: cls.test5_data, r: ref("w3")}
+        del_weak["w3"] = {d1: cls.test5_data}
+
+        del_anchor_int[4] = {d1: cls.test5_data, r: ref("00:00:00:00:00:31")}
+        del_anchor_int[5] = {d1: cls.test5_data, r: ref("00:00:00:00:00:31")}
+        del_weak_mac["00:00:00:00:00:31"] = {d1: cls.test5_data}
+
+        # No references after delete test - create an anchor and a weak entity, delete the weak
+        del_anchor["a6"] = {d1: cls.test6_data, r: ref("w4")}
+        del_weak["w4"] = {d1: cls.test6_data}
+
+        del_anchor_int[6] = {d1: cls.test6_data, r: ref("00:00:00:00:00:41")}
+        del_weak_mac["00:00:00:00:00:41"] = {d1: cls.test6_data}
+
+        inventory = {
+            "del_anchor": del_anchor,
+            "del_anchor_int": del_anchor_int,
+            "del_weak": del_weak,
+            "del_weak_mac": del_weak_mac,
+        }
+        cls.setup_inventory(inventory)
+
+    @classmethod
+    def setup_inventory(cls, inventory: dict[str, dict[str, dict[str, Any]]]) -> None:
+        dps = make_dps(inventory)
+        res = cls.push_datapoints(dps)
+        if res.status_code != 200:
+            print(json.dumps(res.json(), indent=2), sys.stderr)
+            raise Exception(f"Failed to push datapoints: {res.status_code}")
+        sleep(1)  # Wait for the data to be processed
+
+    @staticmethod
+    def delete_request(entity: str, eid: str) -> requests.Response:
+        return requests.delete(f"{common.base_url}/entity/{entity}/{eid}", timeout=5)
+
+    def get_attr(self, entity: str, eid: str, attr: str) -> EntityEidAttrValueOrHistory:
+        return self.get_entity_data(
+            GET_ATTR.format(entity=entity, eid=eid, attr=attr),
+            EntityEidAttrValueOrHistory,
+        )
+
+    def assertExpectedAttrValue(self, entity: str, eid: str, attr: str, expected: Any) -> None:
+        self.query_expected_value(
+            lambda: self.get_attr(entity, eid, attr),
+            lambda received: received == expected,
+            attempts=30,
+            delay_s=0.2,
+            msg=f"A value of {entity}/{eid} is not as expected",
+        )
+
+    def assertExpectedAttrValues(self, entities, attr: str, expected: Any) -> None:
+        for entity, eid in entities:
+            self.assertExpectedAttrValue(entity, eid, attr, expected)
+
+    def test_simple_delete(self):
+        """Delete isolated anchor via API"""
+        anchor = ("del_anchor", "a1")
+        int_anchor = ("del_anchor_int", 1)
+        both = [anchor, int_anchor]
+        # Ensure the anchor exists and is reachable via API
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test1_data)
+        self.assertExpectedAttrValues(both, "data1", expected)
+
+        # Delete the anchor
+        for anchor in both:
+            res = self.delete_request(*anchor)
+            self.assertEqual(res.status_code, 200)
+
+        # Ensure the anchor is deleted
+        sleep(1)
+        expected = EntityEidAttrValueOrHistory(attr_type=1)
+        with self.subTest("String anchor deleted"):
+            self.assertExpectedAttrValue(*anchor, "data1", expected)
+        with self.subTest("Int anchor deleted"):
+            self.assertExpectedAttrValue(*int_anchor, "data1", expected)
+
+    def test_weak_delete(self):
+        """Delete anchor to kill the weak"""
+        entities = [("del_anchor", "a2"), ("del_weak", "w1")]
+        typed_entities = [("del_anchor_int", 2), ("del_weak_mac", "00:00:00:00:00:01")]
+        all_entities = entities + typed_entities
+
+        # Ensure the entities exist
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test2_data)
+        self.assertExpectedAttrValues(all_entities, "data1", expected)
+
+        # Delete the anchors
+        res = self.delete_request(*entities[0])
+        self.assertEqual(res.status_code, 200)
+        res = self.delete_request(*typed_entities[0])
+        self.assertEqual(res.status_code, 200)
+
+        # Ensure all entities are deleted
+        sleep(1)
+        expected = EntityEidAttrValueOrHistory(attr_type=1)
+        with self.subTest("String entities deleted"):
+            self.assertExpectedAttrValues(entities, "data1", expected)
+        with self.subTest("Typed entities deleted"):
+            self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+    def test_ttl_deletes(self):
+        # Merged two tests into one in order to avoid waiting for TTL twice
+        simple_ttl = {"t1": {"data1": self.test3_data}}
+        weak_inventory = {
+            "del_ttl": {"t2": {"data1": self.test3_data, "ref": {"eid": "w2"}}} | simple_ttl,
+            "del_weak": {"w2": {"data1": self.test3_data}},
+        }
+        typed_weak_inventory = {
+            "del_ttl_ip": {
+                "3.0.0.1": {"data1": self.test3_data},
+                "3.0.0.2": {"data1": self.test3_data, "ref": ref("03:00:00:00:00:01")},
+            },
+            "del_weak_mac": {"03:00:00:00:00:01": {"data1": self.test3_data}},
+        }
+        self.setup_inventory(weak_inventory | typed_weak_inventory)
+        simple_ttl_entity = ("del_ttl", "t1")
+        entities = [("del_ttl", "t2"), ("del_weak", "w2")]
+        typed_ttl_entity = ("del_ttl_ip", "3.0.0.1")
+        typed_entities = [("del_ttl_ip", "3.0.0.2"), ("del_weak_mac", "03:00:00:00:00:01")]
+
+        # Ensure the entity exists
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test3_data)
+        self.assertExpectedAttrValue(*simple_ttl_entity, "data1", expected)
+        self.assertExpectedAttrValue(*typed_ttl_entity, "data1", expected)
+        self.assertExpectedAttrValues(entities, "data1", expected)
+        self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+        # Wait for TTL to expire
+        sleep(9.5)
+
+        # Ensure the entity is deleted
+        expected = EntityEidAttrValueOrHistory(attr_type=1)
+        with self.subTest("Simple TTL deleted"):
+            self.assertExpectedAttrValue(*simple_ttl_entity, "data1", expected)
+        with self.subTest("Typed TTL deleted"):
+            self.assertExpectedAttrValue(*typed_ttl_entity, "data1", expected)
+
+        # Ensure the weak entity is deleted
+        with self.subTest("Weak entity deleted"):
+            self.assertExpectedAttrValues(entities, "data1", expected)
+        with self.subTest("Typed weak entity deleted"):
+            self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+    def test_weak_chain(self):
+        # Ensure the entities exist
+        entities = [("del_anchor", "a3"), ("del_weak", "w2_1"), ("del_weak", "w2_2")]
+        typed_entities = [
+            ("del_anchor_int", 3),
+            ("del_weak_mac", "00:00:00:00:00:21"),
+            ("del_weak_mac", "00:00:00:00:00:22"),
+        ]
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test4_data)
+        self.assertExpectedAttrValues(entities, "data1", expected)
+        self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+        # Delete the anchor
+        res = self.delete_request(*entities[0])
+        self.assertEqual(res.status_code, 200)
+        res = self.delete_request(*typed_entities[0])
+        self.assertEqual(res.status_code, 200)
+
+        # Ensure all entities are deleted
+        sleep(2)
+        expected = EntityEidAttrValueOrHistory(attr_type=1)
+        with self.subTest("String entities deleted"):
+            self.assertExpectedAttrValues(entities, "data1", expected)
+        with self.subTest("Typed entities deleted"):
+            self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+    def test_weak_multiple_ref(self):
+        # Ensure the entities exist
+        entities = [("del_anchor", "a4"), ("del_anchor", "a5"), ("del_weak", "w3")]
+        typed_entities = [
+            ("del_anchor_int", 4),
+            ("del_anchor_int", 5),
+            ("del_weak_mac", "00:00:00:00:00:31"),
+        ]
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test5_data)
+        self.assertExpectedAttrValues(entities, "data1", expected)
+        self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+        # Delete one of the anchors
+        to_del = entities.pop(0)
+        res = self.delete_request(*to_del)
+        self.assertEqual(res.status_code, 200)
+        to_del = typed_entities.pop(0)
+        res = self.delete_request(*to_del)
+        self.assertEqual(res.status_code, 200)
+
+        # Ensure the weak entity is still present
+        sleep(1)
+        self.assertExpectedAttrValues(entities, "data1", expected)
+        self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+        # Delete the last anchor
+        to_del = entities.pop(0)
+        res = self.delete_request(*to_del)
+        self.assertEqual(res.status_code, 200)
+        to_del = typed_entities.pop(0)
+        res = self.delete_request(*to_del)
+        self.assertEqual(res.status_code, 200)
+
+        # Ensure the weak entity is deleted
+        sleep(1)
+        expected = EntityEidAttrValueOrHistory(attr_type=1)
+        with self.subTest("String entities deleted"):
+            self.assertExpectedAttrValues(entities, "data1", expected)
+        with self.subTest("Typed entities deleted"):
+            self.assertExpectedAttrValues(typed_entities, "data1", expected)
+
+    def test_no_ref_after_delete(self):
+        entities = [("del_anchor", "a6"), ("del_weak", "w4")]
+        typed_entities = [("del_anchor_int", 6), ("del_weak_mac", "00:00:00:00:00:41")]
+        all_entities = entities + typed_entities
+
+        # Ensure the entities exist
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test6_data)
+        self.assertExpectedAttrValues(all_entities, "data1", expected)
+
+        # Delete the anchors
+        res = self.delete_request(*entities[1])
+        self.assertEqual(res.status_code, 200)
+        res = self.delete_request(*typed_entities[1])
+        self.assertEqual(res.status_code, 200)
+
+        # Ensure weak entities are deleted
+        sleep(1)
+        expected = EntityEidAttrValueOrHistory(attr_type=1)
+        with self.subTest("String entities deleted"):
+            self.assertExpectedAttrValue(*entities[1], "data1", expected)
+        with self.subTest("Typed entities deleted"):
+            self.assertExpectedAttrValue(*typed_entities[1], "data1", expected)
+
+        # Ensure the anchor is still present, but has no references
+        expected_ref = expected
+        expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test6_data)
+        with self.subTest("String anchor has correct value"):
+            self.assertExpectedAttrValue(*entities[0], "data1", expected)
+            self.assertExpectedAttrValue(*entities[0], "ref", expected_ref)
+        with self.subTest("Typed anchor has correct value"):
+            self.assertExpectedAttrValue(*typed_entities[0], "data1", expected)
+            self.assertExpectedAttrValue(*typed_entities[0], "ref", expected_ref)
+
+    def test_no_ref_after_garbage_collection(self):
+        # No references after garbage collection test
+        del_anchor = {
+            "a7": {"data1": self.test7_data, "ref_ttl": ref("t5")},
+            "a8": {"data1": self.test7_data, "ref_ttl": ref("t5")},
+        }
+        del_anchor_int = {
+            7: {"data1": self.test7_data, "ref_ttl": ref("5.0.0.1")},
+            8: {"data1": self.test7_data, "ref_ttl": ref("5.0.0.1")},
+        }
+        inventory = {
+            "del_anchor": del_anchor,
+            "del_ttl": {"t5": {"data1": self.test7_data}},
+            "del_anchor_int": del_anchor_int,
+            "del_ttl_ip": {"5.0.0.1": {"data1": self.test7_data}},
+        }
+        entities = [("del_anchor", "a7"), ("del_anchor", "a8"), ("del_ttl", "t5")]
+        typed_entities = [("del_anchor_int", 7), ("del_anchor_int", 8), ("del_ttl_ip", "5.0.0.1")]
+        all_entities = entities + typed_entities
+        self.setup_inventory(inventory)
+
+        expected_data = EntityEidAttrValueOrHistory(attr_type=1, current_value=self.test7_data)
+        expected_none = EntityEidAttrValueOrHistory(attr_type=1)
+
+        # Ensure the entities exist
+        self.assertExpectedAttrValues(all_entities, "data1", expected_data)
+
+        # Await TTL expiration
+        sleep(9.5)
+
+        # Ensure the entities are deleted
+        with self.subTest("String TTL deleted"):
+            self.assertExpectedAttrValue(*entities[-1], "data1", expected_none)
+        with self.subTest("Typed TTL deleted"):
+            self.assertExpectedAttrValue(*typed_entities[-1], "data1", expected_none)
+
+        # Ensure the anchors are still present, but have no references
+        with self.subTest("String anchors has correct value"):
+            self.assertExpectedAttrValues(entities[:-1], "data1", expected_data)
+            self.assertExpectedAttrValues(entities[:-1], "ref_ttl", expected_none)
+        with self.subTest("Typed anchors has correct value"):
+            self.assertExpectedAttrValues(typed_entities[:-1], "data1", expected_data)
+            self.assertExpectedAttrValues(typed_entities[:-1], "ref_ttl", expected_none)
diff --git a/tests/test_config/db_entities/del_anchor.yml b/tests/test_config/db_entities/del_anchor.yml
new file mode 100644
index 00000000..0f5ae1a5
--- /dev/null
+++ b/tests/test_config/db_entities/del_anchor.yml
@@ -0,0 +1,26 @@
+entity:
+  id: del_anchor
+  id_data_type: string
+  name: Delete test anchor
+  snapshot: true
+  lifetime:
+    type: immortal
+
+attribs:
+  data1:
+    name: data1
+    description: entity data
+    type: plain
+    data_type: string
+
+  ref:
+    name: ref
+    description: reference to weak entity
+    type: plain
+    data_type: link<del_weak>
+
+  ref_ttl:
+    name: ref_ttl
+    description: reference to TTL entity
+    type: plain
+    data_type: link<del_ttl>
diff --git a/tests/test_config/db_entities/del_anchor_int.yml b/tests/test_config/db_entities/del_anchor_int.yml
new file mode 100644
index 00000000..72d19de7
--- /dev/null
+++ b/tests/test_config/db_entities/del_anchor_int.yml
@@ -0,0 +1,26 @@
+entity:
+  id: del_anchor_int
+  id_data_type: int
+  name: Delete test anchor
+  snapshot: true
+  lifetime:
+    type: immortal
+
+attribs:
+  data1:
+    name: data1
+    description: entity data
+    type: plain
+    data_type: string
+
+  ref:
+    name: ref
+    description: reference to weak entity
+    type: plain
+    data_type: link<del_weak_mac>
+
+  ref_ttl:
+    name: ref_ttl
+    description: reference to TTL entity
+    type: plain
+    data_type: link<del_ttl_ip>
diff --git a/tests/test_config/db_entities/del_ttl.yml b/tests/test_config/db_entities/del_ttl.yml
new file mode 100644
index 00000000..e18fc664
--- /dev/null
+++ b/tests/test_config/db_entities/del_ttl.yml
@@ -0,0 +1,23 @@
+entity:
+  id: del_ttl
+  id_data_type: string
+  name: Delete test ttl entity
+  snapshot: true
+  lifetime:
+    type: ttl
+    on_create: 5s
+    mirror_data: true
+
+attribs:
+  data1:
+    name: data1
+    description: entity data
+    type: plain
+    data_type: string
+    ttl: 10s
+
+  ref:
+    name: ref
+    description: reference to weak entity
+    type: plain
+    data_type: link<del_weak>
\ No newline at end of file
diff --git a/tests/test_config/db_entities/del_ttl_ip.yml b/tests/test_config/db_entities/del_ttl_ip.yml
new file mode 100644
index 00000000..90935249
--- /dev/null
+++ b/tests/test_config/db_entities/del_ttl_ip.yml
@@ -0,0 +1,23 @@
+entity:
+  id: del_ttl_ip
+  id_data_type: ipv4
+  name: Delete test ttl entity
+  snapshot: true
+  lifetime:
+    type: ttl
+    on_create: 5s
+    mirror_data: true
+
+attribs:
+  data1:
+    name: data1
+    description: entity data
+    type: plain
+    data_type: string
+    ttl: 10s
+
+  ref:
+    name: ref
+    description: reference to weak entity
+    type: plain
+    data_type: link<del_weak_mac>
\ No newline at end of file
diff --git a/tests/test_config/db_entities/del_weak.yml b/tests/test_config/db_entities/del_weak.yml
new file mode 100644
index 00000000..5f088b44
--- /dev/null
+++ b/tests/test_config/db_entities/del_weak.yml
@@ -0,0 +1,20 @@
+entity:
+  id: del_weak
+  id_data_type: string
+  name: Delete test weak entity
+  snapshot: true
+  lifetime:
+    type: weak
+
+attribs:
+  data1:
+    name: data1
+    description: entity data
+    type: plain
+    data_type: string
+
+  ref:
+    name: ref
+    description: reference to weak entity
+    type: plain
+    data_type: link<del_weak>
\ No newline at end of file
diff --git a/tests/test_config/db_entities/del_weak_mac.yml b/tests/test_config/db_entities/del_weak_mac.yml
new file mode 100644
index 00000000..f255d05b
--- /dev/null
+++ b/tests/test_config/db_entities/del_weak_mac.yml
@@ -0,0 +1,20 @@
+entity:
+  id: del_weak_mac
+  id_data_type: mac
+  name: Delete test weak entity
+  snapshot: true
+  lifetime:
+    type: weak
+
+attribs:
+  data1:
+    name: data1
+    description: entity data
+    type: plain
+    data_type: string
+
+  ref:
+    name: ref
+    description: reference to weak entity
+    type: plain
+    data_type: link<del_weak_mac>
\ No newline at end of file
diff --git a/tests/test_config/garbage_collector.yml b/tests/test_config/garbage_collector.yml
new file mode 100644
index 00000000..b25d3e02
--- /dev/null
+++ b/tests/test_config/garbage_collector.yml
@@ -0,0 +1,5 @@
+# Garbage collector configuration - cleans up expired entities
+
+# The rate at which the garbage collector will run - A CRON expression
+collection_rate:
+  second: "*/2"
\ No newline at end of file

From 32027b7418208a2357105230006952ebaf37eec3 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Thu, 14 Nov 2024 15:22:35 +0100
Subject: [PATCH 13/22] API: Fix serialization of nested types.

---
 dp3/api/internal/entity_response_models.py    |  8 ++-
 tests/test_api/common.py                      | 58 +++++++++++++++++--
 ...st_datapoints.py => test_01_datapoints.py} | 54 ++++++++++++-----
 ...ttr_value.py => test_02_get_attr_value.py} | 44 +++++++++++++-
 .../db_entities/test_entity_type.yml          | 57 ++++++++++++++++++
 5 files changed, 195 insertions(+), 26 deletions(-)
 rename tests/test_api/{test_datapoints.py => test_01_datapoints.py} (58%)
 rename tests/test_api/{test_get_attr_value.py => test_02_get_attr_value.py} (58%)

diff --git a/dp3/api/internal/entity_response_models.py b/dp3/api/internal/entity_response_models.py
index 0e480d9d..728f7701 100644
--- a/dp3/api/internal/entity_response_models.py
+++ b/dp3/api/internal/entity_response_models.py
@@ -24,10 +24,12 @@ class EntityState(BaseModel):
 # This is necessary to allow for non-JSON-serializable types in the model
 JsonVal = Annotated[Any, PlainSerializer(to_json_friendly, when_used="json")]
 
-PlainVal = dict[str, JsonVal]
-HistoryVal = list[dict[str, JsonVal]]
+LinkVal = dict[str, JsonVal]
+PlainVal = Union[LinkVal, JsonVal]
+MultiVal = list[PlainVal]
+HistoryVal = list[dict[str, PlainVal]]
 
-Dp3Val = Union[HistoryVal, PlainVal, JsonVal]
+Dp3Val = Union[HistoryVal, MultiVal, PlainVal]
 
 EntityEidMasterRecord = dict[str, Dp3Val]
 
diff --git a/tests/test_api/common.py b/tests/test_api/common.py
index 671dc74f..b874abe5 100644
--- a/tests/test_api/common.py
+++ b/tests/test_api/common.py
@@ -31,22 +31,25 @@ class ConfigEnv(BaseModel):
 
 values = {
     "valid": {
-        "binary": ["true"],
-        "int": ["123"],
-        "int64": ["123"],
+        "binary": [True],
+        "int": [123],
+        "int64": [123],
         "string": ['"xyz"'],
-        "float": ["1.0", "1"],
+        "float": [1.0],
         "ipv4": ["127.0.0.1"],
         "ipv6": ["2001:0db8:85a3:0000:0000:8a2e:0370:7334", "::1"],
         "mac": ["de:ad:be:ef:ba:be", "11:22:33:44:55:66"],
         "time": ["2020-01-01T00:00:00"],
-        "json": ['{"test": "test"}'],
+        "json": [{"test": "test"}],
         "category": ["cat1"],
         "array": [[1, 2, 3]],
         "set": [[1, 2, 3]],
         "dict": [{"key1": 1, "key2": "xyz"}],
         "link": [{"eid": "test_entity_2"}],
         "data_link": [{"eid": "test_entity_3", "data": 42}],
+        "data_link_int": [{"eid": 666, "data": "1.1.1.1"}],
+        "data_link_ip": [{"eid": "2.2.2.2", "data": "test"}],
+        "data_link_mac": [{"eid": "de:ad:be:ef:ba:be", "data": "test"}],
     },
     "invalid": {
         "binary": ["xyz"],
@@ -70,7 +73,50 @@ class ConfigEnv(BaseModel):
             {"eid": "test_entity_3"},
             {"eid": "test_entity_3", "data": [42]},
         ],
+        "data_link_int": [
+            666,
+            {"id": 666},
+            {"eid": 666},
+            {"eid": 666, "data": [42]},
+            {"eid": "int666", "data": 42},
+        ],
+        "data_link_ip": [
+            "2.2.2.2",
+            {"id": "2.2.2.2"},
+            {"eid": "2.2.2.2"},
+            {"eid": "2.2.2.2", "data": [42]},
+            {"eid": "ip2.2.2.2", "data": "test"},
+        ],
+        "data_link_mac": [
+            "de:ad:be:ef:ba:be",
+            {"id": "de:ad:be:ef:ba:be"},
+            {"eid": "de:ad:be:ef:ba:be"},
+            {"eid": "de:ad:be:ef:ba:be", "data": [42]},
+            {"eid": "macde:ad:be:ef:ba:be", "data": "wrong"},
+        ],
+    },
+}
+
+observation_values = {
+    "valid": {
+        "data_link_multi": [
+            {"eid": "test_entity_4", "data": 42},
+            {"eid": "test_entity_3", "data": 42},
+        ],
+        "data_link_int_multi": [
+            {"eid": 666, "data": "1.1.1.1"},
+            {"eid": 667, "data": "1.1.1.2"},
+        ],
+        "data_link_ip_multi": [
+            {"eid": "2.2.2.2", "data": "test"},
+            {"eid": "2.2.2.3", "data": "test"},
+        ],
+        "data_link_mac_multi": [
+            {"eid": "de:ad:be:ef:ba:be", "data": "test"},
+            {"eid": "11:22:33:44:55:66", "data": "test"},
+        ],
     },
+    "invalid": {},
 }
 
 
@@ -121,7 +167,7 @@ def query_expected_value(
         if msg:
             self.assertTrue(assertion(payload), msg=msg)
         else:
-            self.assertTrue(assertion(payload))
+            self.assertTrue(assertion(payload), f"with payload: {payload}")
         return payload
 
     @classmethod
diff --git a/tests/test_api/test_datapoints.py b/tests/test_api/test_01_datapoints.py
similarity index 58%
rename from tests/test_api/test_datapoints.py
rename to tests/test_api/test_01_datapoints.py
index 0dea92ed..f79a59da 100644
--- a/tests/test_api/test_datapoints.py
+++ b/tests/test_api/test_01_datapoints.py
@@ -1,4 +1,7 @@
+import json
 import sys
+from datetime import datetime
+from typing import Any
 
 import common
 from common import ACCEPTED_ERROR_CODES
@@ -54,30 +57,49 @@ def test_missing_value(self):
         )
         self.assertIn(response.status_code, ACCEPTED_ERROR_CODES)
 
-    def helper_test_datatype_value(self, data_type: str, value, expected_codes: set[int]):
-        response = self.push_datapoints(
-            [
-                {
-                    "type": "test_entity_type",
-                    "id": "test_entity_id",
-                    "attr": f"test_attr_{data_type}",
-                    "v": value,
-                }
-            ]
-        )
-        print(response.content.decode("utf-8"), file=sys.stderr)
+    @staticmethod
+    def make_datapoint(data_type: str, value: Any) -> dict[str, Any]:
+        return {
+            "type": "test_entity_type",
+            "id": "test_entity_id",
+            "attr": f"test_attr_{data_type}",
+            "v": value,
+        }
+
+    def make_observation_datapoint(self, data_type: str, value: Any) -> dict[str, Any]:
+        dp = self.make_datapoint(data_type, value)
+        dp["t1"] = datetime.utcnow().isoformat()
+        return dp
+
+    def helper_test_datatype_value(self, datapoint: dict, expected_codes: set[int]):
+        response = self.push_datapoints([datapoint])
+
+        if response.status_code not in expected_codes:
+            print(json.dumps(response.json(), indent=2), file=sys.stderr)
         self.assertIn(response.status_code, expected_codes)
 
     def test_data_type_values_valid(self):
         for data_type, valid in common.values["valid"].items():
             for value in valid:
                 with self.subTest(data_type=data_type, v=value):
-                    self.helper_test_datatype_value(data_type, value=value, expected_codes={200})
+                    dp = self.make_datapoint(data_type, value)
+                    self.helper_test_datatype_value(dp, expected_codes={200})
+
+        for data_type, valid in common.observation_values["valid"].items():
+            for value in valid:
+                with self.subTest(data_type=data_type, v=value):
+                    dp = self.make_observation_datapoint(data_type, value)
+                    self.helper_test_datatype_value(dp, expected_codes={200})
 
     def test_data_type_values_invalid(self):
         for data_type, valid in common.values["invalid"].items():
             for value in valid:
                 with self.subTest(data_type=data_type, v=value):
-                    self.helper_test_datatype_value(
-                        data_type, value=value, expected_codes=ACCEPTED_ERROR_CODES
-                    )
+                    dp = self.make_datapoint(data_type, value)
+                    self.helper_test_datatype_value(dp, expected_codes=ACCEPTED_ERROR_CODES)
+
+        for data_type, valid in common.observation_values["invalid"].items():
+            for value in valid:
+                with self.subTest(data_type=data_type, v=value):
+                    dp = self.make_observation_datapoint(data_type, value)
+                    self.helper_test_datatype_value(dp, expected_codes=ACCEPTED_ERROR_CODES)
diff --git a/tests/test_api/test_get_attr_value.py b/tests/test_api/test_02_get_attr_value.py
similarity index 58%
rename from tests/test_api/test_get_attr_value.py
rename to tests/test_api/test_02_get_attr_value.py
index 971af5d8..a92e1d31 100644
--- a/tests/test_api/test_get_attr_value.py
+++ b/tests/test_api/test_02_get_attr_value.py
@@ -18,7 +18,9 @@ class GetEidAttrValue(common.APITest):
     @classmethod
     def setUpClass(cls) -> None:
         super().setUpClass()
-        sleep(10)
+        sleep(1)
+        cls.get_request("control/make_snapshots")
+        sleep(6)
 
     def test_unknown_entity_type(self):
         response = self.get_request(
@@ -73,3 +75,43 @@ def test_valid_time_intervals(self):
             with self.subTest(msg=msg):
                 response = self.get_request(TESTED_HISTORY_PATH, **intervals)
                 self.assertEqual(200, response.status_code)
+
+    def get_history_attr_value(self, path: str):
+        result = self.get_entity_data(path, EntityEidAttrValueOrHistory)
+        result.history = []  # Ignore history for sake of test reliability
+        if isinstance(result.current_value, list):
+            result.current_value = sorted(result.current_value, key=lambda x: str(x))
+        return result
+
+    def test_attr_serialization(self):
+        for data_type, valid in common.values["valid"].items():
+            value = valid[-1]  # Plain attribute has the latest sent value
+            with self.subTest(data_type=data_type, v=value):
+                path = TESTED_PATH.format(
+                    entity="test_entity_type",
+                    eid="test_entity_id",
+                    attr=f"test_attr_{data_type}",
+                )
+                expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=value)
+                result = self.get_entity_data(path, EntityEidAttrValueOrHistory)
+                if expected != result:
+                    print(f"Expected: {expected}")
+                    print(f"Result:   {result}")
+                self.assertEqual(expected, result)
+
+        for data_type, valid in common.observation_values["valid"].items():
+            with self.subTest(data_type=data_type, v=value):
+                path = TESTED_PATH.format(
+                    entity="test_entity_type",
+                    eid="test_entity_id",
+                    attr=f"test_attr_{data_type}",
+                )
+                expected = EntityEidAttrValueOrHistory(
+                    attr_type=2, current_value=sorted(valid, key=lambda x: str(x))
+                )
+                self.query_expected_value(
+                    lambda: self.get_history_attr_value(path),
+                    lambda received: received == expected,
+                    attempts=10,
+                    delay_s=0.2,
+                )
diff --git a/tests/test_config/db_entities/test_entity_type.yml b/tests/test_config/db_entities/test_entity_type.yml
index 051da28b..0caeef40 100644
--- a/tests/test_config/db_entities/test_entity_type.yml
+++ b/tests/test_config/db_entities/test_entity_type.yml
@@ -93,3 +93,60 @@ attribs:
     name: test_attr_data_link
     type: plain
     data_type: link<test_entity_type, int>
+  test_attr_data_link_multi:
+    name: test_attr_data_link_multi
+    type: observations
+    data_type: link<test_entity_type, int>
+    multi_value: true
+    history_params:
+      max_age: 7d
+      pre_validity: 10m
+      post_validity: 10m
+      aggregate: true
+
+
+  test_attr_data_link_int:
+    name: test_attr_data_link_int
+    type: plain
+    data_type: link<del_anchor_int, ipv4>
+  test_attr_data_link_ip:
+    name: test_attr_data_link_ip
+    type: plain
+    data_type: link<del_ttl_ip, string>
+  test_attr_data_link_mac:
+    name: test_attr_data_link_mac
+    type: plain
+    data_type: link<del_weak_mac, string>
+
+  test_attr_data_link_int_multi:
+    name: test_attr_data_link_int_multi
+    type: observations
+    data_type: link<del_anchor_int, ipv4>
+    multi_value: true
+    history_params:
+      max_age: 7d
+      pre_validity: 10m
+      post_validity: 10m
+      aggregate: true
+
+  test_attr_data_link_ip_multi:
+    name: test_attr_data_link_ip_multi
+    type: observations
+    data_type: link<del_ttl_ip, string>
+    multi_value: true
+    history_params:
+      max_age: 7d
+      pre_validity: 10m
+      post_validity: 10m
+      aggregate: true
+
+  test_attr_data_link_mac_multi:
+    name: test_attr_data_link_mac_multi
+    type: observations
+    data_type: link<del_weak_mac, string>
+    multi_value: true
+    history_params:
+      max_age: 7d
+      pre_validity: 10m
+      post_validity: 10m
+      aggregate: true

From e73efdf1ccb6b4e282ff70db8d0eefdc2beaa459 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Mon, 18 Nov 2024 10:47:17 +0100
Subject: [PATCH 14/22] Docs: Updated info about types.

---
 docs/api.md                       |  7 ++--
 docs/configuration/db_entities.md | 19 +++++++--
 docs/modules.md                   | 67 +++++++++++++++++++++++++------
 dp3/common/mac_address.py         |  5 +++
 4 files changed, 80 insertions(+), 18 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 3998c08a..3aabdfd1 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -150,20 +150,20 @@ More details depends on the particular type of the attribute.
 Can be represented using both **plain** attributes and **observations**. The difference will be only
 in time specification. Two examples using observations:
 
-**no data - `link<mac>`**: just the eid is sent
+**no data - `link<mac>`**: Sent as a dictionary with a single `"eid"` key.
 
 ```json
 {
   "type": "ip",
   "id": "192.168.0.1",
   "attr": "mac_addrs",
-  "v": "AA:AA:AA:AA:AA",
+  "v": {"eid": "AA:AA:AA:AA:AA"},
   "t1": "2022-08-01T12:00:00",
   "t2": "2022-08-01T12:10:00"
 }
 ```
 
-**with additional data - `link<ip, int>`**: The eid and the data are sent as a dictionary.
+**with additional data - `link<ip, int>`**: Sent as a dictionary with `"eid"` and `"data"` keys.
 
 ```json
 {
@@ -410,6 +410,7 @@ Returns dictionary containing all entity types configured -- their simplified co
 {
   "<entity_id>": {
     "id": "<entity_id>",
+    "id_data_type": "<entity_spec.id_data_type>",
     "name": "<entity_spec.name>",
     "attribs": "<MODEL_SPEC.attribs(e_id)>",
     "eid_estimate_count": "<DB.estimate_count_eids(e_id)>"
diff --git a/docs/configuration/db_entities.md b/docs/configuration/db_entities.md
index eeeb8bc8..1dc7d6a1 100644
--- a/docs/configuration/db_entities.md
+++ b/docs/configuration/db_entities.md
@@ -86,10 +86,23 @@ Entity is described simply by:
 | Parameter      | Data-type           | Default value       | Description                                                                                                                               |
 |----------------|---------------------|---------------------|-------------------------------------------------------------------------------------------------------------------------------------------|
 | **`id`**       | string (identifier) | *(mandatory)*       | Short string identifying the entity type, it's machine name (must match regex `[a-zA-Z_][a-zA-Z0-9_-]*`). Lower-case only is recommended. |
+| `id_data_type` | string              | "string"            | Data type of the entity id (`eid`) value, see [Supported eid data types](#supported-entity-id-data-types).                                |
 | **`name`**     | string              | *(mandatory)*       | Attribute name for humans. May contain any symbols.                                                                                       |
 | **`snapshot`** | bool                | *(mandatory)*       | Whether to create snapshots of the entity. See [Architecture](../architecture.md#data-flow) for more details.                             |
 | `lifetime`     | `Lifetime Spec`     | `Immortal Lifetime` | Defines the lifetime of the entitiy, entities are never deleted by default. See the [Entity Lifetimes](lifetimes.md) for details.         |
 
+### Supported entity id data types
+
+Only a subset of [primitive data types](#primitive-types) is supported for entity ids. The supported data types are:
+
+- `string` (default)
+- `int`: 32-bit signed integer (range from -2147483648 to +2147483647)
+- `ipv4`: IPv4 address, represented as [IPv4Address][ipaddress.IPv4Address] (passed as dotted-decimal string)
+- `ipv6`: IPv6 address, represented as [IPv6Address][ipaddress.IPv6Address] (passed as string in short or full format)
+- `mac`: MAC address, represented as [MACAddress][dp3.common.mac_address.MACAddress]  (passed as string)
+
+Whenever writing a piece of code independent of a specific configuration,
+the [`AnyEidT`][dp3.common.datatype.AnyEidT] type alias should be used.
 
 ## Attributes
 
@@ -196,9 +209,9 @@ List of supported values for parameter `data_type`:
 - `int64`: 64-bit signed integer (use when the range of normal `int` is not sufficent)
 - `float`
 - `time`: Timestamp in `YYYY-MM-DD[T]HH:MM[:SS[.ffffff]][Z or [±]HH[:]MM]` format or timestamp since 1.1.1970 in seconds or milliseconds.
-- `ip4`: IPv4 address (passed as dotted-decimal string)
-- `ip6`: IPv6 address (passed as string in short or full format)
-- `mac`: MAC address (passed as string)
+- `ipv4`: IPv4 address, represented as [IPv4Address][ipaddress.IPv4Address] (passed as dotted-decimal string)
+- `ipv6`: IPv6 address, represented as [IPv6Address][ipaddress.IPv6Address] (passed as string in short or full format)
+- `mac`: MAC address, represented as [MACAddress][dp3.common.mac_address.MACAddress]  (passed as string)
 - `json`: Any JSON object can be stored, all processing is handled by user's code. This is here for special cases which can't be mapped to any other data type.
 
 #### Composite types
diff --git a/docs/modules.md b/docs/modules.md
index 5303baca..a5d646b3 100644
--- a/docs/modules.md
+++ b/docs/modules.md
@@ -122,6 +122,19 @@ and [`on_new_attr`](#attribute-hooks) callbacks,
 and you can enable it by passing the `refresh` keyword argument
 to the callback registration. See the Callbacks section for more details.
 
+## Type of `eid`
+
+!!! tip "Specifying the `eid` type"
+
+    At runtime, the `eid` will be exactly the type as specified in the entity specification.
+
+All the examples on this page will show the `eid` as a string, as that is the default type.
+The type of the `eid` is can be configured in the entity specification, as is
+detailed [here](configuration//db_entities.md#entity).
+
+The typehint of the `eid` used in callback registration definitions is the [
+`AnyEidT`][dp3.common.datatype.AnyEidT] type, which is a type alias of Union of all the allowed
+types of the `eid` in the entity specification.
 
 ## Callbacks
 
@@ -184,10 +197,13 @@ registrar.register_task_hook("on_task_start", task_hook)
 Receives eid and Task, may prevent entity record creation (by returning False).
 The callback is registered using the 
 [`register_allow_entity_creation_hook`][dp3.common.callback_registrar.CallbackRegistrar.register_allow_entity_creation_hook] method.
-Required signature is `Callable[[str, DataPointTask], bool]`.
+Required signature is `Callable[[AnyEidT, DataPointTask], bool]`.
 
 ```python
-def entity_creation(eid: str, task: DataPointTask) -> bool:
+def entity_creation(
+        eid: str,  # (1)! 
+        task: DataPointTask,
+) -> bool:
     return eid.startswith("1")
 
 registrar.register_allow_entity_creation_hook(
@@ -195,16 +211,22 @@ registrar.register_allow_entity_creation_hook(
 )
 ```
 
+1. `eid` may not be string, depending on the entity configuration, see [Type of
+   `eid`](#type-of-eid).
+
 #### Entity `on_entity_creation` hook
 
 Receives eid and Task, may return new DataPointTasks.
 
 Callbacks which are called once when an entity is created are registered using the 
 [`register_on_entity_creation_hook`][dp3.common.callback_registrar.CallbackRegistrar.register_on_entity_creation_hook] method.
-Required signature is `Callable[[str, DataPointTask], list[DataPointTask]]`.
+Required signature is `Callable[[AnyEidT, DataPointTask], list[DataPointTask]]`.
 
 ```python
-def processing_function(eid: str, task: DataPointTask) -> list[DataPointTask]:
+def processing_function(
+        eid: str,  # (1)! 
+        task: DataPointTask
+) -> list[DataPointTask]:
     output = does_work(task)
     return [DataPointTask(
         model_spec=task.model_spec,
@@ -224,6 +246,9 @@ registrar.register_on_entity_creation_hook(
 )
 ```
 
+1. `eid` may not be string, depending on the entity configuration, see [Type of
+   `eid`](#type-of-eid).
+
 The `register_on_entity_creation_hook` method also allows for refreshing of values derived 
 by the registered hook. This can be done using the `refresh` keyword argument, (expecting a [`SharedFlag`][dp3.common.state.SharedFlag] object, which is created by default for all modules)
 and the `may_change` keyword argument, which lists all the attributes that the hook may change.
@@ -243,10 +268,13 @@ registrar.register_on_entity_creation_hook(
 Callbacks that are called on every incoming datapoint of an attribute are registered using the 
 [`register_on_new_attr_hook`][dp3.common.callback_registrar.CallbackRegistrar.register_on_new_attr_hook] method.
 The callback allways receives eid, attribute and Task, and may return new DataPointTasks.
-The required signature is `Callable[[str, DataPointBase], Union[None, list[DataPointTask]]]`.
+The required signature is `Callable[[AnyEidT, DataPointBase], Union[None, list[DataPointTask]]]`.
 
 ```python
-def attr_hook(eid: str, dp: DataPointBase) -> list[DataPointTask]:
+def attr_hook(
+        eid: str,  # (1)!
+        dp: DataPointBase,
+) -> list[DataPointTask]:
     ...
     return []
 
@@ -255,6 +283,9 @@ registrar.register_on_new_attr_hook(
 )
 ```
 
+1. `eid` may not be string, depending on the entity configuration, see [Type of
+   `eid`](#type-of-eid).
+
 This hook can be refreshed on configuration changes if you feel like the attribute value may change too slowly
 to catch up naturally. 
 This can be done using the `refresh` keyword argument, (expecting a [`SharedFlag`][dp3.common.state.SharedFlag] object, which is created by default for all modules)
@@ -290,7 +321,6 @@ def timeseries_hook(
 ) -> list[DataPointTask]:
     ...
     return []
-    
 
 registrar.register_timeseries_hook(
     timeseries_hook, "test_entity_type", "test_attr_type",
@@ -371,8 +401,8 @@ Learn more about the updater module in the [updater configuration](configuration
 #### Periodic Update Hook
 
 The [`register_periodic_update_hook`][dp3.common.callback_registrar.CallbackRegistrar.register_periodic_update_hook]
-method expects a callable with the following signature: 
-`Callable[[str, str, dict], list[DataPointTask]]`, where the arguments are the entity type,
+method expects a callable with the following signature:
+`Callable[[str, AnyEidT, dict], list[DataPointTask]]`, where the arguments are the entity type,
 entity ID and master record. 
 The callable should return a list of DataPointTask objects to perform (possibly empty).
 
@@ -382,7 +412,11 @@ The following example shows how to register a periodic update hook for an entity
 The hook will be called for all entities of this type every day.
 
 ```python
-def periodic_update_hook(entity_type: str, eid: str, record: dict) -> list[DataPointTask]:
+def periodic_update_hook(
+        entity_type: str,
+        eid: str,  # (1)!
+        record: dict,
+) -> list[DataPointTask]:
     ...
     return []
 
@@ -391,6 +425,9 @@ registrar.register_periodic_update_hook(
 )
 ```
 
+1. `eid` may not be string, depending on the entity configuration, see [Type of
+   `eid`](#type-of-eid).
+
 !!! warning "Set a Realistic Update Period"
 
     Try to configure the period to match the real execution time of the registered hooks, 
@@ -405,12 +442,15 @@ This hook is useful when the entity record is not needed for the update, meaning
 
 The [`register_periodic_eid_update_hook`][dp3.common.callback_registrar.CallbackRegistrar.register_periodic_eid_update_hook]
 method expects a callable with the following signature:
-`Callable[[str, str], list[DataPointTask]]`, where the first argument is the entity type and the second is the entity ID.
+`Callable[[str, AnyEidT], list[DataPointTask]]`, where the first argument is the entity type and the second is the entity ID.
 The callable should return a list of DataPointTask objects to perform (possibly empty).
 All other arguments are the same as for the [periodic update hook](#periodic-update-hook).
 
 ```python
-def periodic_eid_update_hook(entity_type: str, eid: str) -> list[DataPointTask]:
+def periodic_eid_update_hook(
+        entity_type: str,
+        eid: str,  # (1)!
+) -> list[DataPointTask]:
     ...
     return []
 
@@ -419,6 +459,9 @@ registrar.register_periodic_eid_update_hook(
 )
 ```
 
+1. `eid` may not be string, depending on the entity configuration, see [Type of
+   `eid`](#type-of-eid).
+
 ## Running module code in a separate thread
 
 The module is free to run its own code in separate threads or processes.
diff --git a/dp3/common/mac_address.py b/dp3/common/mac_address.py
index 6d9d85f2..28100d48 100644
--- a/dp3/common/mac_address.py
+++ b/dp3/common/mac_address.py
@@ -5,6 +5,11 @@
 
 
 class MACAddress:
+    """Represents a MAC Address.
+
+    Can be initialized from colon or comma separated string, or from raw bytes.
+    """
+
     def __init__(self, mac: Union[bytes, str, "MACAddress"]):
         if isinstance(mac, self.__class__):
             mac = mac.mac  # type: ignore

From c6386c8ee6dba2c23c35ae846466a0ba5854d32e Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Mon, 18 Nov 2024 13:34:08 +0100
Subject: [PATCH 15/22] SchemaUpdate: Bumped schema version, added migration.

---
 dp3/bin/schema_update.py       |  27 ++++++-
 dp3/database/schema_cleaner.py | 130 +++++++++++++++++++++++++++------
 2 files changed, 130 insertions(+), 27 deletions(-)

diff --git a/dp3/bin/schema_update.py b/dp3/bin/schema_update.py
index 1e4526b8..020c9515 100644
--- a/dp3/bin/schema_update.py
+++ b/dp3/bin/schema_update.py
@@ -50,11 +50,17 @@ def main(args):
         config.get("processing_core.worker_processes"),
     )
 
-    prev_schema, config_schema, updates, deleted_entites = db.schema_cleaner.get_schema_status()
+    (
+        prev_schema,
+        config_schema,
+        eid_updates,
+        updates,
+        deleted_entites,
+    ) = db.schema_cleaner.get_schema_status()
     if prev_schema["version"] != config_schema["version"]:
         log.info(
             f"Schema version mismatch: {prev_schema['version']} (DB) "
-            f"!= {config_schema['version']} (config)"
+            f"!= {config_schema['version']} (DP3 required)"
         )
         if confirm_changes("Are you sure you want to perform a migration now? (y/[n]): "):
             db.schema_cleaner.migrate(prev_schema)
@@ -67,6 +73,21 @@ def main(args):
         )
         if confirm_changes("Are you sure you want to change the storage now? (y/[n]): "):
             db.schema_cleaner.update_storage(prev_schema["storage"], config_schema["storage"])
+    elif eid_updates:
+        prev_eid_types = prev_schema["entity_id_types"]
+        curr_eid_types = config_schema["entity_id_types"]
+
+        log.info("Mismatch in entity ID types:")
+        for entity, entity_updates in sorted(eid_updates.items(), key=lambda x: x[0]):
+            prev = prev_eid_types.get(entity, "None")
+            curr = curr_eid_types.get(entity, "None")
+            log.info(
+                f"- {entity}: Type changed from {prev} to {curr} "
+                f"- all relevant collections must be dropped"
+            )
+            if entity not in deleted_entites:
+                deleted_entites.append(entity)
+        log.info("")
     elif prev_schema["schema"] == config_schema["schema"]:
         log.info("Schema is OK!")
         return
@@ -78,7 +99,7 @@ def main(args):
 
     if deleted_entites:
         log.info("Suggested removal of entities:")
-        for entity in deleted_entites:
+        for entity in sorted(deleted_entites):
             log.info(f"- {entity}")
 
     if updates:
diff --git a/dp3/database/schema_cleaner.py b/dp3/database/schema_cleaner.py
index 08598712..def02c87 100644
--- a/dp3/database/schema_cleaner.py
+++ b/dp3/database/schema_cleaner.py
@@ -17,7 +17,7 @@
 # number of seconds to wait for the i-th attempt to reconnect after error
 RECONNECT_DELAYS = [1, 2, 5, 10, 30]
 # current database schema version
-SCHEMA_VERSION = 3
+SCHEMA_VERSION = 4
 
 # Collections belonging to entity
 # Used when deleting no-longer existing entity.
@@ -57,7 +57,10 @@ def __init__(
             "snapshot_bucket_size": self._config.get("database.storage.snapshot_bucket_size", 32)
         }
 
-        self.migrations: dict[int, Callable[[dict], dict]] = {2: self.migrate_schema_2_to_3}
+        self.migrations: dict[int, Callable[[dict], dict]] = {
+            2: self.migrate_schema_2_to_3,
+            3: self.migrate_schema_3_to_4,
+        }
 
     def get_current_schema_doc(self, infer: bool = False) -> dict:
         """
@@ -70,8 +73,10 @@ def get_current_schema_doc(self, infer: bool = False) -> dict:
         if schema_doc is None and infer:
             self.log.info("No schema found, inferring initial schema from database")
             schema = self.infer_current_schema()
+            entity_id_types = self.construct_entity_types()
             schema_doc = {
                 "_id": 0,
+                "entity_id_types": entity_id_types,
                 "schema": schema,
                 "storage": self.storage,
                 "version": SCHEMA_VERSION,
@@ -93,7 +98,11 @@ def safe_update_schema(self):
         ```py
             {
                 "_id": int,
-                "schema": { ... }, # (1)!
+                entity_id_types: {
+                    <entity_type>: <entity_id_type>,
+                    ...
+                },
+                "schema": { ... },  # (1)!
                 "storage": {
                     "snapshot_bucket_size": int
                 },
@@ -106,7 +115,7 @@ def safe_update_schema(self):
         Raises:
             ValueError: If conflicting changes are detected.
         """
-        db_schema, config_schema, updates, deleted_entites = self.get_schema_status()
+        db_schema, config_schema, eid_updates, updates, deleted_entites = self.get_schema_status()
 
         if db_schema["version"] != config_schema["version"]:
             self.log.warning(
@@ -124,6 +133,16 @@ def safe_update_schema(self):
             )
             return self._error_on_conflict()
 
+        if eid_updates:
+            for entity in eid_updates.keys():
+                self.log.warning(
+                    "Entity ID type mismatch for %s: %s (DB) -> %s (config)",
+                    entity,
+                    db_schema["entity_id_types"][entity],
+                    config_schema["entity_id_types"][entity],
+                )
+            return self._error_on_conflict()
+
         if db_schema["schema"] == config_schema["schema"]:
             self.log.info("Schema OK!")
         elif not updates and not deleted_entites:
@@ -137,7 +156,7 @@ def _error_on_conflict(self):
         self.log.warning("Please run `dp3 schema-update` to make the changes.")
         raise ValueError("Schema update failed: Conflicting changes detected.")
 
-    def get_schema_status(self) -> tuple[dict, dict, dict, list]:
+    def get_schema_status(self) -> tuple[dict, dict, dict, dict, list]:
         """
         Gets the current schema status.
         `database_schema` is the schema document from the database.
@@ -150,42 +169,61 @@ def get_schema_status(self) -> tuple[dict, dict, dict, list]:
         schema_doc = self.get_current_schema_doc(infer=True)
 
         current_schema = self.construct_schema_doc()
-
-        if (
-            schema_doc["schema"] == current_schema
-            and schema_doc["version"] == SCHEMA_VERSION
-            and schema_doc["storage"] == self.storage
-        ):
-            return schema_doc, schema_doc, {}, []
+        current_entity_types = self.construct_entity_types()
 
         new_schema = {
-            "_id": schema_doc["_id"] + 1,
+            "_id": schema_doc["_id"],
+            "entity_id_types": current_entity_types,
             "schema": current_schema,
             "storage": self.storage,
             "version": SCHEMA_VERSION,
         }
-        updates, deleted_entites = self.detect_changes(schema_doc, current_schema)
-        return schema_doc, new_schema, updates, deleted_entites
+
+        if schema_doc == new_schema:
+            return schema_doc, schema_doc, {}, {}, []
+
+        new_schema["_id"] += 1
+
+        eid_updates, updates, deleted_entites = self.detect_changes(
+            schema_doc, current_entity_types, current_schema
+        )
+        return schema_doc, new_schema, eid_updates, updates, deleted_entites
 
     def detect_changes(
-        self, db_schema_doc: dict, current_schema: dict
-    ) -> tuple[dict[str, dict[str, dict[str, str]]], list[str]]:
+        self, db_schema_doc: dict, current_entity_types: dict, current_schema: dict
+    ) -> tuple[dict[str, str], dict[str, dict[str, dict[str, str]]], list[str]]:
         """
         Detects changes between configured schema and the one saved in the database.
 
         Args:
             db_schema_doc: Schema document from the database.
+            current_entity_types: Entity ID types from the configuration.
             current_schema: Schema from the configuration.
 
         Returns:
             Tuple of required updates to each entity and list of deleted entites.
         """
-        if db_schema_doc["schema"] == current_schema:
-            return {}, []
-
         if db_schema_doc["version"] != SCHEMA_VERSION:
             self.log.info("Schema version changed, skipping detecting changes.")
-            return {}, []
+            return {}, {}, []
+
+        eid_updates = {}
+        for entity, id_type in db_schema_doc["entity_id_types"].items():
+            if entity not in current_schema:
+                self.log.info("Schema breaking change: Entity %s was deleted", entity)
+                continue
+
+            if id_type != current_entity_types[entity]:
+                self.log.info(
+                    "Schema breaking change: Entity %s ID type changed: %s -> %s",
+                    entity,
+                    id_type,
+                    current_entity_types[entity],
+                )
+                eid_updates[entity] = "$drop"
+
+        if db_schema_doc["schema"] == current_schema:
+            return eid_updates, {}, []
 
         updates = {}
         deleted_entites = []
@@ -228,7 +266,7 @@ def detect_changes(
             if entity_updates:
                 updates[entity] = entity_updates
 
-        return updates, deleted_entites
+        return eid_updates, updates, deleted_entites
 
     def execute_updates(
         self, updates: dict[str, dict[str, dict[str, str]]], deleted_entites: list[str]
@@ -240,6 +278,17 @@ def execute_updates(
                     col = col_placeholder.format(entity)
                     self.log.info("%s: Deleting collection: %s", entity, col)
                     self._db[col].drop()
+
+                # Delete entity from Link cache
+                self.log.info("%s: Deleting entity entries from Link cache", entity)
+                self._db["#cache#Link"].delete_many(
+                    {
+                        "$or": [
+                            {"from": {"$regex": f"^{entity}#"}},
+                            {"to": {"$regex": f"^{entity}#"}},
+                        ]
+                    }
+                )
             except Exception as e:
                 raise ValueError(f"Schema update failed: {e}") from e
 
@@ -412,6 +461,18 @@ def _construct_spec_dict(spec: AttrSpecType):
             res["series"] = {s: sspec.data_type.type_info for s, sspec in spec.series.items()}
         return res
 
+    def construct_entity_types(self) -> dict:
+        """
+        Constructs entity ID type dictionary from model specification.
+
+        Returns:
+            Dictionary with entity ID types.
+        """
+        return {
+            entity: entity_spec.id_data_type.type_info
+            for entity, entity_spec in self._model_spec.entities.items()
+        }
+
     def await_updated(self):
         """
         Checks whether schema saved in database is up-to-date and awaits its update
@@ -429,13 +490,13 @@ def await_updated(self):
         if schema_doc is None:
             raise ValueError("Unable to establish schema")
 
-        db_schema, config_schema, _updates, _deleted_entites = self.get_schema_status()
+        db_schema, config_schema, *_ = self.get_schema_status()
         for delay in RECONNECT_DELAYS:
             if db_schema == config_schema:
                 break
             self.log.info("Schema mismatch, will await update by main worker in %s seconds", delay)
             time.sleep(delay)
-            db_schema, config_schema, _updates, _deleted_entites = self.get_schema_status()
+            db_schema, config_schema, *_ = self.get_schema_status()
 
         if db_schema != config_schema:
             raise ValueError("Unable to establish matching schema")
@@ -582,3 +643,24 @@ def migrate_schema_2_to_3(self, schema: dict) -> dict:
         schema["storage"] = {"snapshot_bucket_size": bucket_size}
 
         return schema
+
+    @staticmethod
+    def migrate_schema_3_to_4(schema: dict) -> dict:
+        """
+        Migrates schema from version 3 to version 4.
+
+        This migration adds eid_data_type specification to the schema.
+        As all entities were previously using string as their entity ID type,
+        this is merely an accounting change.
+        Actual data types are to be changed by following migrations.
+
+        Args:
+            schema: Schema to migrate.
+        Returns:
+            Migrated schema.
+        """
+        schema["_id"] += 1
+        schema["version"] = 4
+        schema["entity_id_types"] = {entity: str(str) for entity in schema["schema"].keys()}
+
+        return schema

From 5ba3dda2d1de5031781edd25536344eb31cce423 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Mon, 18 Nov 2024 13:59:52 +0100
Subject: [PATCH 16/22] Docs: Updated list entities endpoint doc.

---
 docs/api.md | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 3aabdfd1..191690a2 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -198,11 +198,27 @@ v -> some_embedded_dict_field
 
 ## List entities
 
-List latest snapshots of all ids present in database under entity type.
+List latest snapshots of all ids present in database under entity type, 
+filtered by `generic_filter` and `fulltext_filters`.
+Contains only the latest snapshot per entity.
 
-Contains only latest snapshot.
+Uses pagination, default limit is 20, setting to 0 will return all results.
 
-Uses pagination.
+Fulltext filters are interpreted as regular expressions.
+Only string values may be filtered this way. There's no validation that queried attribute
+can be fulltext filtered.
+Only plain and observation attributes with string-based data types can be queried.
+Array and set data types are supported as well as long as they are not multi value
+at the same time.
+If you need to filter EIDs, use attribute `eid`.
+
+Generic filter allows filtering using generic MongoDB query (including `$and`, `$or`,`$lt`, etc.).
+For querying non-JSON-native types, you can use the following magic strings,
+as are defined by the search & replace [`magic`][dp3.database.magic] module.
+
+There are no attribute name checks (may be added in the future).
+
+Generic and fulltext filters are merged - fulltext overrides conflicting keys.
 
 ### Request
 
@@ -212,6 +228,8 @@ Uses pagination.
 
 - skip: how many entities to skip (default: 0)
 - limit: how many entities to return (default: 20)
+- fulltext_filters: dictionary of fulltext filters (default: no filters)
+- generic_filter: dictionary of generic filters (default: no filters)
 
 ### Response
 

From 66e658db0ed848b49bba461e02961057bc804da7 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Mon, 18 Nov 2024 14:32:34 +0100
Subject: [PATCH 17/22] Chore: Bump Ruff version.

---
 .pre-commit-config.yaml                  | 2 +-
 dp3/bin/schema_update.py                 | 2 +-
 dp3/database/schema_cleaner.py           | 4 ++--
 dp3/snapshots/snapshooter.py             | 2 +-
 tests/test_api/test_02_get_attr_value.py | 4 ++--
 5 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 51b97842..ea10d842 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -18,7 +18,7 @@ repos:
       - id: black
 
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.6.1'
+    rev: 'v0.7.4'
     hooks:
       - id: ruff
         args: [ "--fix", "." ]
diff --git a/dp3/bin/schema_update.py b/dp3/bin/schema_update.py
index 020c9515..7a45238f 100644
--- a/dp3/bin/schema_update.py
+++ b/dp3/bin/schema_update.py
@@ -78,7 +78,7 @@ def main(args):
         curr_eid_types = config_schema["entity_id_types"]
 
         log.info("Mismatch in entity ID types:")
-        for entity, entity_updates in sorted(eid_updates.items(), key=lambda x: x[0]):
+        for entity in sorted(eid_updates):
             prev = prev_eid_types.get(entity, "None")
             curr = curr_eid_types.get(entity, "None")
             log.info(
diff --git a/dp3/database/schema_cleaner.py b/dp3/database/schema_cleaner.py
index def02c87..8f76b562 100644
--- a/dp3/database/schema_cleaner.py
+++ b/dp3/database/schema_cleaner.py
@@ -134,7 +134,7 @@ def safe_update_schema(self):
             return self._error_on_conflict()
 
         if eid_updates:
-            for entity in eid_updates.keys():
+            for entity in eid_updates:
                 self.log.warning(
                     "Entity ID type mismatch for %s: %s (DB) -> %s (config)",
                     entity,
@@ -661,6 +661,6 @@ def migrate_schema_3_to_4(schema: dict) -> dict:
         """
         schema["_id"] += 1
         schema["version"] = 4
-        schema["entity_id_types"] = {entity: str(str) for entity in schema["schema"].keys()}
+        schema["entity_id_types"] = {entity: str(str) for entity in schema["schema"]}
 
         return schema
diff --git a/dp3/snapshots/snapshooter.py b/dp3/snapshots/snapshooter.py
index f7aeaa0f..713d4426 100644
--- a/dp3/snapshots/snapshooter.py
+++ b/dp3/snapshots/snapshooter.py
@@ -38,9 +38,9 @@
     DataPointTask,
     Snapshot,
     SnapshotMessageType,
+    parse_eid_tuples_from_cache,
     parse_eids_from_cache,
     task_context,
-    parse_eid_tuples_from_cache,
 )
 from dp3.common.types import EventGroupType
 from dp3.common.utils import get_func_name
diff --git a/tests/test_api/test_02_get_attr_value.py b/tests/test_api/test_02_get_attr_value.py
index a92e1d31..f56d7b4d 100644
--- a/tests/test_api/test_02_get_attr_value.py
+++ b/tests/test_api/test_02_get_attr_value.py
@@ -110,8 +110,8 @@ def test_attr_serialization(self):
                     attr_type=2, current_value=sorted(valid, key=lambda x: str(x))
                 )
                 self.query_expected_value(
-                    lambda: self.get_history_attr_value(path),
-                    lambda received: received == expected,
+                    lambda: self.get_history_attr_value(path),  # noqa: B023
+                    lambda received: received == expected,  # noqa: B023
                     attempts=10,
                     delay_s=0.2,
                 )

From 6fdf16f7460233fd515f344838c4b8a609fcedb6 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Mon, 18 Nov 2024 14:41:59 +0100
Subject: [PATCH 18/22] API tests: Loosen timing requirements.

---
 .github/workflows/tests.yml              | 3 +++
 tests/test_api/test_02_get_attr_value.py | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index bc0fdc3e..ef45f027 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -18,6 +18,9 @@ jobs:
       - name: Pause for stack to start
         run: sleep 10
 
+      - name: Show running containers
+        run: docker ps -a
+
       - name: Integration tests - API
         run: docker run 
           --env CONF_DIR=/dp3/tests/test_config 
diff --git a/tests/test_api/test_02_get_attr_value.py b/tests/test_api/test_02_get_attr_value.py
index f56d7b4d..2080dd24 100644
--- a/tests/test_api/test_02_get_attr_value.py
+++ b/tests/test_api/test_02_get_attr_value.py
@@ -112,6 +112,6 @@ def test_attr_serialization(self):
                 self.query_expected_value(
                     lambda: self.get_history_attr_value(path),  # noqa: B023
                     lambda received: received == expected,  # noqa: B023
-                    attempts=10,
+                    attempts=50,
                     delay_s=0.2,
                 )

From 59f1bd94125b86fec8a56dc205adf77d968ff567 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Mon, 18 Nov 2024 15:11:16 +0100
Subject: [PATCH 19/22] API tests: Reduce debug logging.

---
 tests/test_api/common.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/test_api/common.py b/tests/test_api/common.py
index b874abe5..848383db 100644
--- a/tests/test_api/common.py
+++ b/tests/test_api/common.py
@@ -162,8 +162,11 @@ def query_expected_value(
                 time.sleep(delay_s)
             payload = query()
             if assertion(payload):
+                if i > 0:
+                    print(f"After {i * delay_s:.2f}s: {payload}", file=sys.stderr)
                 break
-            print(payload, file=sys.stderr)
+        else:
+            print(f"After {attempts * delay_s:.2f}s: {payload}", file=sys.stderr)
         if msg:
             self.assertTrue(assertion(payload), msg=msg)
         else:

From 6a9b3c342102229ad11ce165573f2e296a5914db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?D=C3=A1vid=20Benko?= <benko@cesnet.cz>
Date: Tue, 19 Nov 2024 10:41:13 +0100
Subject: [PATCH 20/22] Tests: add missing terminating newlines

---
 tests/test_config/db_entities/del_ttl.yml      | 2 +-
 tests/test_config/db_entities/del_ttl_ip.yml   | 2 +-
 tests/test_config/db_entities/del_weak.yml     | 2 +-
 tests/test_config/db_entities/del_weak_mac.yml | 2 +-
 tests/test_config/garbage_collector.yml        | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/test_config/db_entities/del_ttl.yml b/tests/test_config/db_entities/del_ttl.yml
index e18fc664..0fbe4381 100644
--- a/tests/test_config/db_entities/del_ttl.yml
+++ b/tests/test_config/db_entities/del_ttl.yml
@@ -20,4 +20,4 @@ attribs:
     name: ref
     description: reference to weak entity
     type: plain
-    data_type: link<del_weak>
\ No newline at end of file
+    data_type: link<del_weak>
diff --git a/tests/test_config/db_entities/del_ttl_ip.yml b/tests/test_config/db_entities/del_ttl_ip.yml
index 90935249..d33ae1e2 100644
--- a/tests/test_config/db_entities/del_ttl_ip.yml
+++ b/tests/test_config/db_entities/del_ttl_ip.yml
@@ -20,4 +20,4 @@ attribs:
     name: ref
     description: reference to weak entity
     type: plain
-    data_type: link<del_weak_mac>
\ No newline at end of file
+    data_type: link<del_weak_mac>
diff --git a/tests/test_config/db_entities/del_weak.yml b/tests/test_config/db_entities/del_weak.yml
index 5f088b44..bc8b75c7 100644
--- a/tests/test_config/db_entities/del_weak.yml
+++ b/tests/test_config/db_entities/del_weak.yml
@@ -17,4 +17,4 @@ attribs:
     name: ref
     description: reference to weak entity
     type: plain
-    data_type: link<del_weak>
\ No newline at end of file
+    data_type: link<del_weak>
diff --git a/tests/test_config/db_entities/del_weak_mac.yml b/tests/test_config/db_entities/del_weak_mac.yml
index f255d05b..51cc764f 100644
--- a/tests/test_config/db_entities/del_weak_mac.yml
+++ b/tests/test_config/db_entities/del_weak_mac.yml
@@ -17,4 +17,4 @@ attribs:
     name: ref
     description: reference to weak entity
     type: plain
-    data_type: link<del_weak_mac>
\ No newline at end of file
+    data_type: link<del_weak_mac>
diff --git a/tests/test_config/garbage_collector.yml b/tests/test_config/garbage_collector.yml
index b25d3e02..04e39a37 100644
--- a/tests/test_config/garbage_collector.yml
+++ b/tests/test_config/garbage_collector.yml
@@ -2,4 +2,4 @@
 
 # The rate at which the garbage collector will run - A CRON expression
 collection_rate:
-  second: "*/2"
\ No newline at end of file
+  second: "*/2"

From ed9522a9ef648f938683e40745625fdf788c855e Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 19 Nov 2024 12:00:50 +0100
Subject: [PATCH 21/22] Tests: Dump logs on failure

---
 .github/workflows/tests.yml | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index ef45f027..d439a873 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -18,15 +18,16 @@ jobs:
       - name: Pause for stack to start
         run: sleep 10
 
-      - name: Show running containers
-        run: docker ps -a
-
       - name: Integration tests - API
         run: docker run 
           --env CONF_DIR=/dp3/tests/test_config 
           --network container:dp3_api 
           dp3_interpreter python -m unittest discover -s tests/test_api -v
 
+      - name: Dump Logs on Failure
+        if: failure()
+        run: docker compose logs
+
       - name: Check worker errors
         run: docker compose logs worker | grep "WARNING\|ERROR\|exception" | grep -v "RabbitMQ\|it's\ OK\ now,\ we're\ successfully\ connected" || true
 

From 251b5bc3253b49667f54bbb2a3fcf74f3a87db91 Mon Sep 17 00:00:00 2001
From: Ondrej Sedlacek <xsedla1o@stud.fit.vutbr.cz>
Date: Tue, 19 Nov 2024 13:27:33 +0100
Subject: [PATCH 22/22] Tests: Fixed flaky test.

---
 tests/test_api/test_02_get_attr_value.py | 37 ++++++++++++------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/tests/test_api/test_02_get_attr_value.py b/tests/test_api/test_02_get_attr_value.py
index 2080dd24..903ff803 100644
--- a/tests/test_api/test_02_get_attr_value.py
+++ b/tests/test_api/test_02_get_attr_value.py
@@ -18,9 +18,9 @@ class GetEidAttrValue(common.APITest):
     @classmethod
     def setUpClass(cls) -> None:
         super().setUpClass()
-        sleep(1)
+        sleep(8)
+        # Data must pass to DB before this request is processed
         cls.get_request("control/make_snapshots")
-        sleep(6)
 
     def test_unknown_entity_type(self):
         response = self.get_request(
@@ -84,23 +84,9 @@ def get_history_attr_value(self, path: str):
         return result
 
     def test_attr_serialization(self):
-        for data_type, valid in common.values["valid"].items():
-            value = valid[-1]  # Plain attribute has the latest sent value
-            with self.subTest(data_type=data_type, v=value):
-                path = TESTED_PATH.format(
-                    entity="test_entity_type",
-                    eid="test_entity_id",
-                    attr=f"test_attr_{data_type}",
-                )
-                expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=value)
-                result = self.get_entity_data(path, EntityEidAttrValueOrHistory)
-                if expected != result:
-                    print(f"Expected: {expected}")
-                    print(f"Result:   {result}")
-                self.assertEqual(expected, result)
-
+        # Test first with `query_expected_value` to await snapshot
         for data_type, valid in common.observation_values["valid"].items():
-            with self.subTest(data_type=data_type, v=value):
+            with self.subTest(data_type=data_type, v=valid):
                 path = TESTED_PATH.format(
                     entity="test_entity_type",
                     eid="test_entity_id",
@@ -115,3 +101,18 @@ def test_attr_serialization(self):
                     attempts=50,
                     delay_s=0.2,
                 )
+
+        for data_type, valid in common.values["valid"].items():
+            value = valid[-1]  # Plain attribute has the latest sent value
+            with self.subTest(data_type=data_type, v=value):
+                path = TESTED_PATH.format(
+                    entity="test_entity_type",
+                    eid="test_entity_id",
+                    attr=f"test_attr_{data_type}",
+                )
+                expected = EntityEidAttrValueOrHistory(attr_type=1, current_value=value)
+                result = self.get_entity_data(path, EntityEidAttrValueOrHistory)
+                if expected != result:
+                    print(f"Expected: {expected}")
+                    print(f"Result:   {result}")
+                self.assertEqual(expected, result)