07fb811788
Some checks failed
Build / Test on ubuntu-latest (push) Successful in 3m6s
Lint / lint (push) Successful in 22s
Build / typecheck (ubuntu-latest, 3.9) (push) Successful in 1m49s
Lint / format (push) Successful in 42s
Build / Test on macOS-latest (push) Has been cancelled
Build / Test on windows-latest (push) Has been cancelled
CodeQL / Analyze (python) (push) Failing after 1m31s
389 lines
14 KiB
Python
389 lines
14 KiB
Python
import functools
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from socket import gethostname
|
|
from typing import (
|
|
Any,
|
|
Dict,
|
|
List,
|
|
Optional,
|
|
)
|
|
from uuid import uuid4
|
|
|
|
import iso8601
|
|
from aw_core.dirs import get_data_dir
|
|
from aw_core.log import get_log_file_path
|
|
from aw_core.models import Event
|
|
from aw_query import query2
|
|
from aw_query.exceptions import QueryException
|
|
from aw_transform import heartbeat_merge
|
|
|
|
from .__about__ import __version__
|
|
from .exceptions import NotFound
|
|
from .profile import profile_from_env
|
|
from .settings import Settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_device_id() -> str:
|
|
path = Path(get_data_dir("aw-server")) / "device_id"
|
|
if path.exists():
|
|
with open(path) as f:
|
|
return f.read()
|
|
else:
|
|
uuid = str(uuid4())
|
|
with open(path, "w") as f:
|
|
f.write(uuid)
|
|
return uuid
|
|
|
|
|
|
def check_bucket_exists(f):
|
|
@functools.wraps(f)
|
|
def g(self, bucket_id, *args, **kwargs):
|
|
if bucket_id not in self.db.buckets():
|
|
raise NotFound("NoSuchBucket", f"There's no bucket named {bucket_id}")
|
|
return f(self, bucket_id, *args, **kwargs)
|
|
|
|
return g
|
|
|
|
|
|
class ServerAPI:
|
|
def __init__(self, db, testing) -> None:
|
|
self.db = db
|
|
self.settings = Settings(testing)
|
|
self.testing = testing
|
|
self.profile = profile_from_env(testing=testing)
|
|
self.last_event = {} # type: dict
|
|
|
|
def get_info(self) -> Dict[str, Any]:
|
|
"""Get server info"""
|
|
payload = {
|
|
"hostname": gethostname(),
|
|
"version": __version__,
|
|
"testing": self.testing,
|
|
"device_id": get_device_id(),
|
|
"profile": self.profile,
|
|
}
|
|
return payload
|
|
|
|
def get_buckets(self) -> Dict[str, Dict]:
|
|
"""Get dict {bucket_name: Bucket} of all buckets"""
|
|
logger.debug("Received get request for buckets")
|
|
buckets = self.db.buckets()
|
|
for b in buckets:
|
|
# TODO: Move this code to aw-core?
|
|
last_events = self.db[b].get(limit=1)
|
|
if len(last_events) > 0:
|
|
last_event = last_events[0]
|
|
last_updated = last_event.timestamp + last_event.duration
|
|
buckets[b]["last_updated"] = last_updated.isoformat()
|
|
return buckets
|
|
|
|
@check_bucket_exists
|
|
def get_bucket_metadata(self, bucket_id: str) -> Dict[str, Any]:
|
|
"""Get metadata about bucket."""
|
|
bucket = self.db[bucket_id]
|
|
return bucket.metadata()
|
|
|
|
@check_bucket_exists
|
|
def export_bucket(self, bucket_id: str) -> Dict[str, Any]:
|
|
"""Export a bucket to a dataformat consistent across versions, including all events in it."""
|
|
bucket = self.get_bucket_metadata(bucket_id)
|
|
bucket["events"] = self.get_events(bucket_id, limit=-1)
|
|
# Scrub event IDs
|
|
for event in bucket["events"]:
|
|
del event["id"]
|
|
return bucket
|
|
|
|
def export_all(self) -> Dict[str, Any]:
|
|
"""Exports all buckets and their events to a format consistent across versions"""
|
|
buckets = self.get_buckets()
|
|
exported_buckets = {}
|
|
for bid in buckets.keys():
|
|
exported_buckets[bid] = self.export_bucket(bid)
|
|
return exported_buckets
|
|
|
|
def import_bucket(self, bucket_data: Any):
|
|
bucket_id = bucket_data["id"]
|
|
logger.info(f"Importing bucket {bucket_id}")
|
|
|
|
# TODO: Check that bucket doesn't already exist
|
|
self.db.create_bucket(
|
|
bucket_id,
|
|
type=bucket_data["type"],
|
|
client=bucket_data["client"],
|
|
hostname=bucket_data["hostname"],
|
|
created=(
|
|
bucket_data["created"]
|
|
if isinstance(bucket_data["created"], datetime)
|
|
else iso8601.parse_date(bucket_data["created"])
|
|
),
|
|
)
|
|
|
|
# scrub IDs from events
|
|
# (otherwise causes weird bugs with no events seemingly imported when importing events exported from aw-server-rust, which contains IDs)
|
|
for event in bucket_data["events"]:
|
|
if "id" in event:
|
|
del event["id"]
|
|
|
|
self.create_events(
|
|
bucket_id,
|
|
[Event(**e) if isinstance(e, dict) else e for e in bucket_data["events"]],
|
|
)
|
|
|
|
def import_all(self, buckets: Dict[str, Any]):
|
|
for bid, bucket in buckets.items():
|
|
self.import_bucket(bucket)
|
|
|
|
def create_bucket(
|
|
self,
|
|
bucket_id: str,
|
|
event_type: str,
|
|
client: str,
|
|
hostname: str,
|
|
created: Optional[datetime] = None,
|
|
data: Optional[Dict[str, Any]] = None,
|
|
) -> bool:
|
|
"""
|
|
Create a bucket.
|
|
|
|
If hostname is "!local", the hostname and device_id will be set from the server info.
|
|
This is useful for watchers which are known/assumed to run locally but might not know their hostname (like aw-watcher-web).
|
|
|
|
Returns True if successful, otherwise false if a bucket with the given ID already existed.
|
|
"""
|
|
if created is None:
|
|
created = datetime.now()
|
|
if bucket_id in self.db.buckets():
|
|
return False
|
|
if hostname == "!local":
|
|
info = self.get_info()
|
|
if data is None:
|
|
data = {}
|
|
hostname = info["hostname"]
|
|
data["device_id"] = info["device_id"]
|
|
self.db.create_bucket(
|
|
bucket_id,
|
|
type=event_type,
|
|
client=client,
|
|
hostname=hostname,
|
|
created=created,
|
|
data=data,
|
|
)
|
|
return True
|
|
|
|
@check_bucket_exists
|
|
def update_bucket(
|
|
self,
|
|
bucket_id: str,
|
|
event_type: Optional[str] = None,
|
|
client: Optional[str] = None,
|
|
hostname: Optional[str] = None,
|
|
data: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
"""Update bucket metadata"""
|
|
self.db.update_bucket(
|
|
bucket_id,
|
|
type=event_type,
|
|
client=client,
|
|
hostname=hostname,
|
|
data=data,
|
|
)
|
|
return None
|
|
|
|
@check_bucket_exists
|
|
def delete_bucket(self, bucket_id: str) -> None:
|
|
"""Delete a bucket"""
|
|
self.db.delete_bucket(bucket_id)
|
|
logger.debug(f"Deleted bucket '{bucket_id}'")
|
|
return None
|
|
|
|
@check_bucket_exists
|
|
def get_event(
|
|
self,
|
|
bucket_id: str,
|
|
event_id: int,
|
|
) -> Optional[Event]:
|
|
"""Get a single event from a bucket"""
|
|
logger.debug(
|
|
f"Received get request for event {event_id} in bucket '{bucket_id}'"
|
|
)
|
|
event = self.db[bucket_id].get_by_id(event_id)
|
|
return event.to_json_dict() if event else None
|
|
|
|
@check_bucket_exists
|
|
def get_events(
|
|
self,
|
|
bucket_id: str,
|
|
limit: int = -1,
|
|
start: Optional[datetime] = None,
|
|
end: Optional[datetime] = None,
|
|
) -> List[Event]:
|
|
"""Get events from a bucket"""
|
|
logger.debug(f"Received get request for events in bucket '{bucket_id}'")
|
|
if limit is None: # Let limit = None also mean "no limit"
|
|
limit = -1
|
|
events = [
|
|
event.to_json_dict() for event in self.db[bucket_id].get(limit, start, end)
|
|
]
|
|
return events
|
|
|
|
@check_bucket_exists
|
|
def create_events(self, bucket_id: str, events: List[Event]) -> List[Event]:
|
|
"""Create events for a bucket. Can handle both single events and multiple ones.
|
|
|
|
Always returns a list of inserted events (matching aw-server-rust behavior).
|
|
For single events, the returned event includes the server-assigned ID.
|
|
For bulk inserts, returns empty list (events may not have IDs without a response-SQL roundtrip).
|
|
"""
|
|
if len(events) == 1:
|
|
# Pass as single Event so Bucket.insert uses insert_one (returns Event with ID)
|
|
inserted = self.db[bucket_id].insert(events[0])
|
|
return [inserted]
|
|
else:
|
|
self.db[bucket_id].insert(events)
|
|
return []
|
|
|
|
@check_bucket_exists
|
|
def get_eventcount(
|
|
self,
|
|
bucket_id: str,
|
|
start: Optional[datetime] = None,
|
|
end: Optional[datetime] = None,
|
|
) -> int:
|
|
"""Get eventcount from a bucket"""
|
|
logger.debug(f"Received get request for eventcount in bucket '{bucket_id}'")
|
|
return self.db[bucket_id].get_eventcount(start, end)
|
|
|
|
@check_bucket_exists
|
|
def delete_event(self, bucket_id: str, event_id) -> bool:
|
|
"""Delete a single event from a bucket"""
|
|
return self.db[bucket_id].delete(event_id)
|
|
|
|
@check_bucket_exists
|
|
def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event:
|
|
"""
|
|
Heartbeats are useful when implementing watchers that simply keep
|
|
track of a state, how long it's in that state and when it changes.
|
|
A single heartbeat always has a duration of zero.
|
|
|
|
If the heartbeat was identical to the last (apart from timestamp), then the last event has its duration updated.
|
|
If the heartbeat differed, then a new event is created.
|
|
|
|
Such as:
|
|
- Active application and window title
|
|
- Example: aw-watcher-window
|
|
- Currently open document/browser tab/playing song
|
|
- Example: wakatime
|
|
- Example: aw-watcher-web
|
|
- Example: aw-watcher-spotify
|
|
- Is the user active/inactive?
|
|
Send an event on some interval indicating if the user is active or not.
|
|
- Example: aw-watcher-afk
|
|
|
|
Inspired by: https://wakatime.com/developers#heartbeats
|
|
"""
|
|
logger.debug(
|
|
"Received heartbeat in bucket '{}'\n\ttimestamp: {}, duration: {}, pulsetime: {}\n\tdata: {}".format(
|
|
bucket_id,
|
|
heartbeat.timestamp,
|
|
heartbeat.duration,
|
|
pulsetime,
|
|
heartbeat.data,
|
|
)
|
|
)
|
|
|
|
# The endtime here is set such that in the event that the heartbeat is older than an
|
|
# existing event we should try to merge it with the last event before the heartbeat instead.
|
|
# FIXME: This (the endtime=heartbeat.timestamp) gets rid of the "heartbeat was older than last event"
|
|
# warning and also causes a already existing "newer" event to be overwritten in the
|
|
# replace_last call below. This is problematic.
|
|
# Solution: This could be solved if we were able to replace arbitrary events.
|
|
# That way we could double check that the event has been applied
|
|
# and if it hasn't we simply replace it with the updated counterpart.
|
|
|
|
last_event = None
|
|
if bucket_id not in self.last_event:
|
|
last_events = self.db[bucket_id].get(limit=1)
|
|
if len(last_events) > 0:
|
|
last_event = last_events[0]
|
|
else:
|
|
last_event = self.last_event[bucket_id]
|
|
|
|
if last_event:
|
|
if last_event.data == heartbeat.data:
|
|
merged = heartbeat_merge(last_event, heartbeat, pulsetime)
|
|
if merged is not None:
|
|
# Heartbeat was merged into last_event
|
|
logger.debug(
|
|
"Received valid heartbeat, merging. (bucket: {})".format(
|
|
bucket_id
|
|
)
|
|
)
|
|
self.last_event[bucket_id] = merged
|
|
self.db[bucket_id].replace_last(merged)
|
|
return merged
|
|
else:
|
|
logger.info(
|
|
"Received heartbeat after pulse window, inserting as new event. (bucket: {})".format(
|
|
bucket_id
|
|
)
|
|
)
|
|
else:
|
|
logger.debug(
|
|
"Received heartbeat with differing data, inserting as new event. (bucket: {})".format(
|
|
bucket_id
|
|
)
|
|
)
|
|
else:
|
|
logger.info(
|
|
"Received heartbeat, but bucket was previously empty, inserting as new event. (bucket: {})".format(
|
|
bucket_id
|
|
)
|
|
)
|
|
|
|
self.db[bucket_id].insert(heartbeat)
|
|
self.last_event[bucket_id] = heartbeat
|
|
return heartbeat
|
|
|
|
def query2(self, name, query, timeperiods, cache):
|
|
result = []
|
|
for timeperiod in timeperiods:
|
|
period = timeperiod.split("/")[
|
|
:2
|
|
] # iso8601 timeperiods are separated by a slash
|
|
if len(period) != 2:
|
|
raise QueryException(
|
|
f"Invalid timeperiod '{timeperiod}': expected two ISO8601 "
|
|
"datetimes separated by a slash (start/end)"
|
|
)
|
|
try:
|
|
starttime = iso8601.parse_date(period[0])
|
|
endtime = iso8601.parse_date(period[1])
|
|
except iso8601.ParseError as e:
|
|
raise QueryException(f"Invalid timeperiod '{timeperiod}': {e}")
|
|
query = "".join(query)
|
|
result.append(query2.query(name, query, starttime, endtime, self.db))
|
|
return result
|
|
|
|
# TODO: Right now the log format on disk has to be JSON, this is hard to read by humans...
|
|
def get_log(self):
|
|
"""Get the server log in json format"""
|
|
payload = []
|
|
with open(get_log_file_path()) as log_file:
|
|
for line in log_file.readlines()[::-1]:
|
|
payload.append(json.loads(line))
|
|
return payload, 200
|
|
|
|
def get_setting(self, key):
|
|
"""Get a setting"""
|
|
return self.settings.get(key, None)
|
|
|
|
def set_setting(self, key, value):
|
|
"""Set a setting"""
|
|
self.settings[key] = value
|
|
return value
|