hat-event 0.9.27__cp310.cp311.cp312.cp313-abi3-win_amd64.whl

hat/event/backends/lmdb/latestdb.py

@@ -0,0 +1,116 @@
+import itertools
+import typing
+
+import lmdb
+
+from hat.event.backends.lmdb import common
+from hat.event.backends.lmdb import environment
+from hat.event.backends.lmdb.conditions import Conditions
+
+
+class Changes(typing.NamedTuple):
+    data: dict[common.EventTypeRef, common.Event]
+    types: dict[common.EventTypeRef, common.EventId]
+
+
+class AddResult(typing.NamedTuple):
+    added_ref: common.EventRef | None
+    removed_ref: tuple[common.EventId, common.EventRef] | None
+
+
+def ext_create(env: environment.Environment,
+               txn: lmdb.Transaction,
+               conditions: Conditions,
+               subscription: common.Subscription
+               ) -> 'LatestDb':
+    db = LatestDb()
+    db._env = env
+    db._subscription = subscription
+    db._changes = Changes({}, {})
+
+    db._event_type_refs = {
+        event_type: ref
+        for ref, event_type in env.ext_read(txn, common.DbType.LATEST_TYPE)}
+
+    db._events = {
+        event.type: event
+        for ref, event in env.ext_read(txn, common.DbType.LATEST_DATA)
+        if conditions.matches(event) and
+        subscription.matches(event.type) and
+        db._event_type_refs.get(event.type) == ref}
+
+    db._next_event_type_refs = itertools.count(
+        max(db._event_type_refs.values(), default=0) + 1)
+
+    return db
+
+
+class LatestDb:
+
+    def add(self,
+            event: common.Event
+            ) -> AddResult:
+        if not self._subscription.matches(event.type):
+            return AddResult(added_ref=None,
+                             removed_ref=None)
+
+        previous_event = self._events.get(event.type)
+        if previous_event and previous_event > event:
+            return AddResult(added_ref=None,
+                             removed_ref=None)
+
+        event_type_ref = self._event_type_refs.get(event.type)
+        if event_type_ref is None:
+            event_type_ref = next(self._next_event_type_refs)
+            self._changes.types[event_type_ref] = event.type
+            self._event_type_refs[event.type] = event_type_ref
+
+        self._changes.data[event_type_ref] = event
+        self._events[event.type] = event
+
+        added_ref = common.LatestEventRef(event_type_ref)
+        removed_ref = ((previous_event.id, added_ref)
+                       if previous_event else None)
+        return AddResult(added_ref=added_ref,
+                         removed_ref=removed_ref)
+
+    def query(self,
+              params: common.QueryLatestParams
+              ) -> common.QueryResult:
+        event_types = (set(params.event_types)
+                       if params.event_types is not None
+                       else None)
+
+        if event_types is None or ('*', ) in event_types:
+            events = self._events.values()
+
+        elif any('*' in event_type or '?' in event_type
+                 for event_type in event_types):
+            subscription = common.create_subscription(event_types)
+            events = (event for event in self._events.values()
+                      if subscription.matches(event.type))
+
+        elif len(event_types) < len(self._events):
+            events = (self._events.get(event_type)
+                      for event_type in event_types)
+            events = (event for event in events if event)
+
+        else:
+            events = (event for event in self._events.values()
+                      if event.type in event_types)
+
+        return common.QueryResult(events=list(events),
+                                  more_follows=False)
+
+    def create_changes(self) -> Changes:
+        changes, self._changes = self._changes, Changes({}, {})
+        return changes
+
+    def ext_write(self,
+                  txn: lmdb.Transaction,
+                  changes: Changes):
+        self._env.ext_write(txn, common.DbType.LATEST_DATA,
+                            changes.data.items())
+
+        self._env.ext_write(txn, common.DbType.LATEST_TYPE,
+                            changes.types.items())

hat/event/backends/lmdb/manager/__main__.py

@@ -0,0 +1,8 @@
+import sys
+
+from hat.event.backends.lmdb.manager.main import main
+
+
+if __name__ == '__main__':
+    sys.argv[0] = 'hat-event-lmdb-manager'
+    sys.exit(main())

hat/event/backends/lmdb/manager/common.py

@@ -0,0 +1,45 @@
+from hat.event.backends.lmdb.common import *  # NOQA
+
+from hat import json
+
+from hat.event.backends.lmdb.common import (EventId,
+                                            Timestamp,
+                                            Event,
+                                            EventPayloadJson,
+                                            EventPayloadBinary)
+
+
+def event_id_to_json(event_id: EventId) -> json.Data:
+    return {'server': event_id.server,
+            'session': event_id.session,
+            'instance': event_id.instance}
+
+
+def timestamp_to_json(timestamp: Timestamp) -> json.Data:
+    return {'s': timestamp.s,
+            'us': timestamp.us}
+
+
+def event_to_json(event: Event) -> json.Data:
+    return {'id': event_id_to_json(event.id),
+            'type': list(event.type),
+            'timestamp': timestamp_to_json(event.timestamp),
+            'source_timestamp': (timestamp_to_json(event.source_timestamp)
+                                 if event.source_timestamp else None),
+            'payload': _event_payload_to_json(event.payload)}
+
+
+def _event_payload_to_json(payload):
+    if payload is None:
+        return None
+
+    if isinstance(payload, EventPayloadBinary):
+        return {'type': 'BINARY',
+                'subtype': payload.type,
+                'data': bytes(payload.data).hex()}
+
+    if isinstance(payload, EventPayloadJson):
+        return {'type': 'JSON',
+                'data': payload.data}
+
+    raise ValueError('unsupported payload type')

hat/event/backends/lmdb/manager/copy.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+import argparse
+
+from hat.event.backends.lmdb.manager import common
+
+
+def create_argument_parser(subparsers) -> argparse.ArgumentParser:
+    parser = subparsers.add_parser('copy')
+    parser.add_argument('--skip-latest', action='store_true')
+    parser.add_argument('--skip-timeseries', action='store_true')
+    parser.add_argument('src_path', type=Path)
+    parser.add_argument('dst_path', type=Path)
+    return parser
+
+
+def copy(args):
+    if args.dst_path.exists():
+        raise Exception('destination db already exists')
+
+    with common.ext_create_env(args.src_path, readonly=True) as src_env:
+        with common.ext_create_env(args.dst_path) as dst_env:
+            src_dbs = {db_type: common.ext_open_db(src_env, db_type, False)
+                       for db_type in common.DbType}
+            dst_dbs = {db_type: common.ext_open_db(dst_env, db_type, True)
+                       for db_type in common.DbType}
+
+            with src_env.begin(buffers=True) as src_txn:
+                with dst_env.begin(write=True) as dst_txn:
+                    _copy_system(src_dbs, src_txn, dst_dbs, dst_txn)
+                    _copy_ref(src_dbs, src_txn, dst_dbs, dst_txn,
+                              args.skip_latest, args.skip_timeseries)
+
+                    if not args.skip_latest:
+                        _copy_latest(src_dbs, src_txn, dst_dbs, dst_txn)
+
+                    if not args.skip_timeseries:
+                        _copy_timeseries(src_dbs, src_txn, dst_dbs, dst_txn)
+
+
+def _copy_system(src_dbs, src_txn, dst_dbs, dst_txn):
+    for db_type in [common.DbType.SYSTEM_SETTINGS,
+                    common.DbType.SYSTEM_LAST_EVENT_ID,
+                    common.DbType.SYSTEM_LAST_TIMESTAMP]:
+        _copy_db(src_dbs[db_type], src_txn, dst_dbs[db_type], dst_txn)
+
+
+def _copy_ref(src_dbs, src_txn, dst_dbs, dst_txn, skip_latest,
+              skip_timeseries):
+    db_type = common.DbType.REF
+    db_def = common.db_defs[db_type]
+
+    for key, value in src_txn.cursor(db=src_dbs[db_type]):
+        if not skip_latest and not skip_timeseries:
+            dst_txn.put(key, value, db=dst_dbs[db_type])
+            continue
+
+        event_refs = db_def.decode_value(value)
+
+        if skip_latest:
+            event_refs = (
+                event_ref for event_ref in event_refs
+                if not isinstance(event_ref, common.LatestEventRef))
+
+        if skip_timeseries:
+            event_refs = (
+                event_ref for event_ref in event_refs
+                if not isinstance(event_ref, common.TimeseriesEventRef))
+
+        event_refs = set(event_refs)
+        if not event_refs:
+            continue
+
+        value = db_def.encode_value(event_refs)
+        dst_txn.put(key, value, db=dst_dbs[db_type])
+
+
+def _copy_latest(src_dbs, src_txn, dst_dbs, dst_txn):
+    for db_type in [common.DbType.LATEST_DATA,
+                    common.DbType.LATEST_TYPE]:
+        _copy_db(src_dbs[db_type], src_txn, dst_dbs[db_type], dst_txn)
+
+
+def _copy_timeseries(src_dbs, src_txn, dst_dbs, dst_txn):
+    for db_type in [common.DbType.TIMESERIES_DATA,
+                    common.DbType.TIMESERIES_PARTITION,
+                    common.DbType.TIMESERIES_COUNT]:
+        _copy_db(src_dbs[db_type], src_txn, dst_dbs[db_type], dst_txn)
+
+
+def _copy_db(src_db, src_txn, dst_db, dst_txn):
+    for key, value in src_txn.cursor(db=src_db):
+        dst_txn.put(key, value, db=dst_db)

hat/event/backends/lmdb/manager/main.py

@@ -0,0 +1,34 @@
+import argparse
+import sys
+
+from hat.event.backends.lmdb.manager import copy
+from hat.event.backends.lmdb.manager import query
+
+
+def create_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest='action', required=True)
+
+    query.create_argument_parser(subparsers)
+    copy.create_argument_parser(subparsers)
+
+    return parser
+
+
+def main():
+    parser = create_argument_parser()
+    args = parser.parse_args()
+
+    if args.action == 'query':
+        query.query(args)
+
+    elif args.action == 'copy':
+        copy.copy(args)
+
+    else:
+        raise ValueError('unsupported action')
+
+
+if __name__ == '__main__':
+    sys.argv[0] = 'hat-event-lmdb-manager'
+    sys.exit(main())

hat/event/backends/lmdb/manager/query.py

@@ -0,0 +1,215 @@
+from pathlib import Path
+import argparse
+
+from hat import json
+
+from hat.event.backends.lmdb.manager import common
+
+
+def create_argument_parser(subparsers) -> argparse.ArgumentParser:
+    parser = subparsers.add_parser('query')
+    parser.add_argument('--db', metavar='PATH', type=Path, required=True)
+    subsubparsers = parser.add_subparsers(dest='subaction', required=True)
+
+    subsubparsers.add_parser('settings')
+
+    last_event_id_parser = subsubparsers.add_parser('last_event_id')
+    last_event_id_parser.add_argument('--server-id', type=int, default=None)
+
+    last_timestamps_parser = subsubparsers.add_parser('last_timestamp')
+    last_timestamps_parser.add_argument('--server-id', type=int, default=None)
+
+    ref_parser = subsubparsers.add_parser('ref')
+    ref_parser.add_argument('--server-id', type=int, default=None)
+
+    latest_parser = subsubparsers.add_parser('latest')
+    latest_parser.add_argument('--event-types', type=str, default=None,
+                               nargs='*')
+
+    partition_parser = subsubparsers.add_parser('partition')
+    partition_parser.add_argument('--partition-id', type=int, default=None)
+
+    timeseries_parser = subsubparsers.add_parser('timeseries')
+    timeseries_parser.add_argument('--partition-id', type=int, default=None)
+    timeseries_parser.add_argument('--server-id', type=int, default=None)
+    timeseries_parser.add_argument('--event-types', type=str, default=None,
+                                   nargs='*')
+
+    return parser
+
+
+def query(args):
+    with common.ext_create_env(args.db, readonly=True) as env:
+        dbs = {db_type: common.ext_open_db(env, db_type, False)
+               for db_type in common.DbType}
+
+        with env.begin(buffers=True) as txn:
+            if args.subaction == 'settings':
+                for result in _query_settings(dbs, txn):
+                    _print_result(result)
+
+            elif args.subaction == 'last_event_id':
+                for result in _query_last_event_id(dbs, txn, args.server_id):
+                    _print_result(result)
+
+            elif args.subaction == 'last_timestamp':
+                for result in _query_last_timestamp(dbs, txn, args.server_id):
+                    _print_result(result)
+
+            elif args.subaction == 'ref':
+                for result in _query_ref(dbs, txn):
+                    _print_result(result)
+
+            elif args.subaction == 'latest':
+                event_types = ([tuple(i.split('/')) for i in args.event_types]
+                               if args.event_types is not None else None)
+                for result in _query_ref(dbs, txn, event_types):
+                    _print_result(result)
+
+            elif args.subaction == 'partition':
+                for result in _query_partition(dbs, txn, args.partition_id):
+                    _print_result(result)
+
+            elif args.subaction == 'timeseries':
+                event_types = ([tuple(i.split('/')) for i in args.event_types]
+                               if args.event_types is not None else None)
+                for result in _query_timeseries(dbs, txn, args.partition_id,
+                                                args.server_id, event_types):
+                    _print_result(result)
+
+            else:
+                raise ValueError('unsupported subaction')
+
+
+def _print_result(result):
+    print(json.encode(result))
+
+
+def _query_settings(dbs, txn):
+    db_type = common.DbType.SYSTEM_SETTINGS
+    db = dbs[db_type]
+    db_def = common.db_defs[db_type]
+
+    for key, value in txn.cursor(db=db):
+        settings_id = db_def.decode_key(key)
+        data = db_def.decode_value(value)
+
+        yield {'settings_id': settings_id.name,
+               'data': data}
+
+
+def _query_last_event_id(dbs, txn, server_id):
+    db_type = common.DbType.SYSTEM_LAST_EVENT_ID
+    db = dbs[db_type]
+    db_def = common.db_defs[db_type]
+
+    for key, value in txn.cursor(db=db):
+        decoded_server_id = db_def.decode_key(key)
+        if server_id is not None and server_id != decoded_server_id:
+            continue
+
+        event_id = db_def.decode_value(value)
+        yield common.event_id_to_json(event_id)
+
+
+def _query_last_timestamp(dbs, txn, server_id):
+    db_type = common.DbType.SYSTEM_LAST_TIMESTAMP
+    db = dbs[db_type]
+    db_def = common.db_defs[db_type]
+
+    for key, value in txn.cursor(db=db):
+        decoded_server_id = db_def.decode_key(key)
+        if server_id is not None and server_id != decoded_server_id:
+            continue
+
+        timestamp = db_def.decode_value(value)
+        yield common.timestamp_to_json(timestamp)
+
+
+def _query_ref(dbs, txn, server_id):
+    ref_db_type = common.DbType.REF
+    ref_db = dbs[ref_db_type]
+    ref_db_def = common.db_defs[ref_db_type]
+
+    latest_type_db_type = common.DbType.LATEST_TYPE
+    latest_type_db = dbs[latest_type_db_type]
+    latest_type_db_def = common.db_defs[latest_type_db_type]
+
+    for key, value in txn.cursor(db=ref_db):
+        event_id = ref_db_def.decode_key(key)
+        if server_id is not None and event_id.server != server_id:
+            continue
+
+        event_refs = ref_db_def.decode_value(value)
+
+        for event_ref in event_refs:
+            if isinstance(event_ref, common.LatestEventRef):
+                latest_type_key = latest_type_db_def.encoded_key(event_ref.key)
+                latest_type_value = txn.get(latest_type_key, db=latest_type_db)
+
+                event_type = latest_type_db_def.decode_value(latest_type_value)
+                yield {'event_id': common.event_id_to_json(event_id),
+                       'ref_type': 'latest',
+                       'event_type': list(event_type)}
+
+            if isinstance(event_ref, common.TimeseriesEventRef):
+                partition_id, timestamp, _ = event_ref.key
+                yield {'event_id': common.event_id_to_json(event_id),
+                       'ref_type': 'timeseries',
+                       'partition_id': partition_id,
+                       'timestamp': common.timestamp_to_json(timestamp)}
+
+
+def _query_latest(dbs, txn, event_types):
+    db_type = common.DbType.LATEST_DATA
+    db = dbs[db_type]
+    db_def = common.db_defs[db_type]
+
+    subscription = common.create_subscription([('*', )] if event_types is None
+                                              else event_types)
+
+    for _, value in txn.cursor(db=db):
+        event = db_def.decode_value(value)
+        if not subscription.matches(event.type):
+            continue
+
+        yield common.event_to_json(value)
+
+
+def _query_partition(dbs, txn, partition_id):
+    db_type = common.DbType.TIMESERIES_PARTITION
+    db = dbs[db_type]
+    db_def = common.db_defs[db_type]
+
+    for key, value in txn.cursor(db=db):
+        decoded_partition_id = db_def.decode_key(key)
+        if partition_id is not None and partition_id != decoded_partition_id:
+            continue
+
+        data = db_def.decode_value(value)
+        yield {'partition_id': decoded_partition_id,
+               'data': data}
+
+
+def _query_timeseries(dbs, txn, partition_id, server_id, event_types):
+    db_type = common.DbType.TIMESERIES_DATA
+    db = dbs[db_type]
+    db_def = common.db_defs[db_type]
+
+    subscription = common.create_subscription([('*', )] if event_types is None
+                                              else event_types)
+
+    for key, value in txn.cursor(db=db):
+        decoded_partition_id, _, event_id = db_def.decode(key)
+
+        if partition_id is not None and partition_id != decoded_partition_id:
+            continue
+
+        if server_id is not None and server_id != event_id.server:
+            continue
+
+        event = db_def.decode_value(value)
+        if not subscription.matches(event.type):
+            continue
+
+        yield common.event_to_json(value)