snapshot of https://github.com/ActivityWatch/aw-core.git (master, 2026-09-16; blobs > 131072B omitted)

aw_datastore/datastore.py

@@ -0,0 +1,189 @@
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import (
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Union,
+)
+
+from aw_core.models import Event
+
+from .storages import AbstractStorage
+
+logger = logging.getLogger(__name__)
+
+
+class Datastore:
+    def __init__(
+        self,
+        storage_strategy: Callable[..., AbstractStorage],
+        testing=False,
+        **kwargs,
+    ) -> None:
+        self.logger = logger.getChild("Datastore")
+        self.bucket_instances: Dict[str, Bucket] = dict()
+
+        self.storage_strategy = storage_strategy(testing=testing, **kwargs)
+
+    def __repr__(self):
+        return f"<Datastore object using {self.storage_strategy.__class__.__name__}>"
+
+    def __getitem__(self, bucket_id: str) -> "Bucket":
+        # If this bucket doesn't have a initialized object, create it
+        if bucket_id not in self.bucket_instances:
+            # If the bucket exists in the database, create an object representation of it
+            if bucket_id in self.buckets():
+                bucket = Bucket(self, bucket_id)
+                self.bucket_instances[bucket_id] = bucket
+            else:
+                self.logger.error(
+                    f"Cannot create a Bucket object for {bucket_id} because it doesn't exist in the database"
+                )
+                raise KeyError
+
+        return self.bucket_instances[bucket_id]
+
+    def create_bucket(
+        self,
+        bucket_id: str,
+        type: str,
+        client: str,
+        hostname: str,
+        created: Optional[datetime] = None,
+        name: Optional[str] = None,
+        data: Optional[dict] = None,
+    ) -> "Bucket":
+        created = created or datetime.now(timezone.utc)
+        self.logger.info(f"Creating bucket '{bucket_id}'")
+        self.storage_strategy.create_bucket(
+            bucket_id, type, client, hostname, created.isoformat(), name=name, data=data
+        )
+        return self[bucket_id]
+
+    def update_bucket(self, bucket_id: str, **kwargs):
+        self.logger.info(f"Updating bucket '{bucket_id}'")
+        return self.storage_strategy.update_bucket(bucket_id, **kwargs)
+
+    def delete_bucket(self, bucket_id: str):
+        self.logger.info(f"Deleting bucket '{bucket_id}'")
+        if bucket_id in self.bucket_instances:
+            del self.bucket_instances[bucket_id]
+        return self.storage_strategy.delete_bucket(bucket_id)
+
+    def buckets(self):
+        return self.storage_strategy.buckets()
+
+
+class Bucket:
+    def __init__(self, datastore: Datastore, bucket_id: str) -> None:
+        self.logger = logger.getChild("Bucket")
+        self.ds = datastore
+        self.bucket_id = bucket_id
+
+    def metadata(self) -> dict:
+        return self.ds.storage_strategy.get_metadata(self.bucket_id)
+
+    def get(
+        self,
+        limit: int = -1,
+        starttime: Optional[datetime] = None,
+        endtime: Optional[datetime] = None,
+    ) -> List[Event]:
+        """Returns events sorted in descending order by timestamp"""
+        # Resolution is rounded down since not all datastores like microsecond precision
+        if starttime:
+            starttime = starttime.replace(
+                microsecond=1000 * int(starttime.microsecond / 1000)
+            )
+        if endtime:
+            # Rounding up here in order to ensure events aren't missed
+            # second_offset and microseconds modulo required since replace() only takes microseconds up to 999999 (doesn't handle overflow)
+            milliseconds = 1 + int(endtime.microsecond / 1000)
+            second_offset = int(milliseconds / 1000)  # usually 0, rarely 1
+            microseconds = (
+                (1000 * milliseconds) % 1000000
+            )  # will likely just be 1000 * milliseconds, if it overflows it would become zero
+            endtime = endtime.replace(microsecond=microseconds) + timedelta(
+                seconds=second_offset
+            )
+
+        return self.ds.storage_strategy.get_events(
+            self.bucket_id, limit, starttime, endtime
+        )
+
+    def get_by_id(self, event_id) -> Optional[Event]:
+        """Will return the event with the provided ID, or None if not found."""
+        return self.ds.storage_strategy.get_event(self.bucket_id, event_id)
+
+    def get_eventcount(
+        self, starttime: Optional[datetime] = None, endtime: Optional[datetime] = None
+    ) -> int:
+        return self.ds.storage_strategy.get_eventcount(
+            self.bucket_id, starttime, endtime
+        )
+
+    def insert(self, events: Union[Event, List[Event]]) -> Optional[Event]:
+        """
+        Inserts one or several events.
+        If a single event is inserted, return the event with its id assigned.
+        If several events are inserted, returns None. (This is due to there being no efficient way of getting ids out when doing bulk inserts with some datastores such as peewee/SQLite)
+        """
+
+        # NOTE: Should we keep the timestamp checking?
+        warn_older_event = False
+
+        # Get last event for timestamp check after insert
+        if warn_older_event:
+            last_event_list = self.get(1)
+            last_event = None
+            if last_event_list:
+                last_event = last_event_list[0]
+
+        now = datetime.now(tz=timezone.utc)
+
+        inserted: Optional[Event] = None
+
+        # Call insert
+        if isinstance(events, Event):
+            oldest_event: Optional[Event] = events
+            if events.timestamp + events.duration > now:
+                self.logger.warning(
+                    f"Event inserted into bucket {self.bucket_id} reaches into the future. Current UTC time: {str(now)}. Event data: {str(events)}"
+                )
+            inserted = self.ds.storage_strategy.insert_one(self.bucket_id, events)
+            # assert inserted
+        elif isinstance(events, list):
+            if events:
+                oldest_event = sorted(events, key=lambda k: k["timestamp"])[0]
+            else:  # pragma: no cover
+                oldest_event = None
+            for event in events:
+                if event.timestamp + event.duration > now:
+                    self.logger.warning(
+                        f"Event inserted into bucket {self.bucket_id} reaches into the future. Current UTC time: {str(now)}. Event data: {str(event)}"
+                    )
+            self.ds.storage_strategy.insert_many(self.bucket_id, events)
+        else:
+            raise TypeError
+
+        # Warn if timestamp is older than last event
+        if warn_older_event and last_event and oldest_event:
+            if oldest_event.timestamp < last_event.timestamp:  # pragma: no cover
+                self.logger.warning(
+                    f"""Inserting event that has a older timestamp than previous event!
+Previous: {last_event}
+Inserted: {oldest_event}"""
+                )
+
+        return inserted
+
+    def delete(self, event_id):
+        return self.ds.storage_strategy.delete(self.bucket_id, event_id)
+
+    def replace_last(self, event):
+        return self.ds.storage_strategy.replace_last(self.bucket_id, event)
+
+    def replace(self, event_id, event):
+        return self.ds.storage_strategy.replace(self.bucket_id, event_id, event)