commit cc2ac65d6b9577035d6c02d7bc3c63333042c3f1 · zenfyr.dev/xpost

+35 -7

atproto/identity.py

···

       1
       1
       -
       from typing import Any

     

       1
       1
       +
       from pathlib import Path

     

       2
       2
       +
       from typing import Any, override

     

       2
       3
        
       

     

       3
       4
        
       import dns.resolver

     

       4
       5
        
       import requests

     

       5
       6
        
       

     

       6
       7
        
       import env

     

       7
       7
       -
       from util.cache import TTLCache

     

       8
       8
       -
       from util.util import LOGGER, normalize_service_url

     

       8
       8
       +
       from util.cache import Cacheable, TTLCache

     

       9
       9
       +
       from util.util import LOGGER, normalize_service_url, shutdown_hook

     

       9
       10
        
       

     

       10
       11
        
       

     

       11
       11
       -
       class DidDocument:

     

       12
       12
       +
       class DidDocument():

     

       12
       13
        
           def __init__(self, raw_doc: dict[str, Any]) -> None:

     

       13
       14
        
               self.raw: dict[str, Any] = raw_doc

     

       14
       15
        
               self.atproto_pds: str | None = None

     
···

       35
       36
        
               return None

     

       36
       37
        
       

     

       37
       38
        
       

     

       38
       38
       -
       class DidResolver:

     

       39
       39
       +
       class DidResolver(Cacheable):

     

       39
       40
        
           def __init__(self, plc_host: str) -> None:

     

       40
       41
        
               self.plc_host: str = plc_host

     

       41
       42
        
               self.__cache: TTLCache[str, DidDocument] = TTLCache(ttl_seconds=12 * 60 * 60)

     
···

       79
       80
        
                       return from_web

     

       80
       81
        
               raise Exception(f"Failed to resolve {did}!")

     

       81
       82
        
       

     

       83
       83
       +
           @override

     

       84
       84
       +
           def dump_cache(self, path: Path):

     

       85
       85
       +
               self.__cache.dump_cache(path)

     

       82
       86
        
       

     

       83
       83
       -
       class HandleResolver:

     

       87
       87
       +
           @override

     

       88
       88
       +
           def load_cache(self, path: Path):

     

       89
       89
       +
               self.__cache.load_cache(path)

     

       90
       90
       +
       

     

       91
       91
       +
       class HandleResolver(Cacheable):

     

       84
       92
        
           def __init__(self) -> None:

     

       85
       85
       -
               self.__cache: TTLCache[str, str] = TTLCache()

     

       93
       93
       +
               self.__cache: TTLCache[str, str] = TTLCache(ttl_seconds=12 * 60 * 60)

     

       86
       94
        
       

     

       87
       95
        
           def try_resolve_dns(self, handle: str) -> str | None:

     

       88
       96
        
               try:

     
···

       131
       139
        
       

     

       132
       140
        
               raise Exception(f"Failed to resolve handle {handle}!")

     

       133
       141
        
       

     

       142
       142
       +
           @override

     

       143
       143
       +
           def dump_cache(self, path: Path):

     

       144
       144
       +
               self.__cache.dump_cache(path)

     

       145
       145
       +
       

     

       146
       146
       +
           @override

     

       147
       147
       +
           def load_cache(self, path: Path):

     

       148
       148
       +
               self.__cache.load_cache(path)

     

       149
       149
       +
       

     

       134
       150
        
       

     

       135
       151
        
       handle_resolver = HandleResolver()

     

       136
       152
        
       did_resolver = DidResolver(env.PLC_HOST)

     

       153
       153
       +
       

     

       154
       154
       +
       did_cache = Path(env.CACHE_DIR).joinpath('did.cache')

     

       155
       155
       +
       handle_cache = Path(env.CACHE_DIR).joinpath('handle.cache')

     

       156
       156
       +
       

     

       157
       157
       +
       did_resolver.load_cache(did_cache)

     

       158
       158
       +
       handle_resolver.load_cache(handle_cache)

     

       159
       159
       +
       

     

       160
       160
       +
       def cache_dump():

     

       161
       161
       +
           did_resolver.dump_cache(did_cache)

     

       162
       162
       +
           handle_resolver.dump_cache(handle_cache)

     

       163
       163
       +
       

     

       164
       164
       +
       shutdown_hook.append(cache_dump)

+2 -1

env.py

···

       2
       2
        
       

     

       3
       3
        
       DEV = bool(os.environ.get("DEV")) or False

     

       4
       4
        
       DATA_DIR = os.environ.get("DATA_DIR") or "./data"

     

       5
       5
       +
       CACHE_DIR = os.environ.get("CACHE_DIR") or "./data/cache"

     

       5
       6
        
       MIGRATIONS_DIR = os.environ.get("MIGRATIONS_DIR") or "./migrations"

     

       6
       6
       -
       PLC_HOST = os.environ.get("PLC_HOST") or "https://plc.wtf"

     

       7
       7
       +
       PLC_HOST = os.environ.get("PLC_HOST") or "https://plc.directory"

+4 -1

main.py

···

       10
       10
        
       from database.migrations import DatabaseMigrator

     

       11
       11
        
       from registry import create_input_service, create_output_service

     

       12
       12
        
       from registry_bootstrap import bootstrap

     

       13
       13
       -
       from util.util import LOGGER, read_env

     

       13
       13
       +
       from util.util import LOGGER, read_env, shutdown_hook

     

       14
       14
        
       

     

       15
       15
        
       

     

       16
       16
        
       def main() -> None:

     
···

       87
       87
        
           task_queue.put(None)

     

       88
       88
        
           thread.join()

     

       89
       89
        
           db_pool.close()

     

       90
       90
       +
       

     

       91
       91
       +
           for shook in shutdown_hook:

     

       92
       92
       +
               shook()

     

       90
       93
        
       

     

       91
       94
        
       

     

       92
       95
        
       if __name__ == "__main__":

+26 -2

util/cache.py

···

       1
       1
       +
       from abc import ABC, abstractmethod

     

       2
       2
       +
       from pathlib import Path

     

       1
       3
        
       import time

     

       2
       2
       -
       from typing import Generic, TypeVar

     

       4
       4
       +
       from typing import Generic, TypeVar, override

     

       5
       5
       +
       import pickle

     

       3
       6
        
       

     

       4
       7
        
       K = TypeVar("K")

     

       5
       8
        
       V = TypeVar("V")

     

       6
       9
        
       

     

       7
       7
       -
       class TTLCache(Generic[K, V]):

     

       10
       10
       +
       class Cacheable(ABC):

     

       11
       11
       +
           @abstractmethod

     

       12
       12
       +
           def dump_cache(self, path: Path):

     

       13
       13
       +
               pass

     

       14
       14
       +
       

     

       15
       15
       +
           @abstractmethod

     

       16
       16
       +
           def load_cache(self, path: Path):

     

       17
       17
       +
               pass

     

       18
       18
       +
       

     

       19
       19
       +
       class TTLCache(Generic[K, V], Cacheable):

     

       8
       20
        
           def __init__(self, ttl_seconds: int = 3600) -> None:

     

       9
       21
        
               self.ttl: int = ttl_seconds

     

       10
       22
        
               self.__cache: dict[K, tuple[V, float]] = {}

     
···

       23
       35
        
       

     

       24
       36
        
           def clear(self) -> None:

     

       25
       37
        
               self.__cache.clear()

     

       38
       38
       +
       

     

       39
       39
       +
           @override

     

       40
       40
       +
           def dump_cache(self, path: Path) -> None:

     

       41
       41
       +
               path.parent.mkdir(parents=True, exist_ok=True)

     

       42
       42
       +
               with open(path, 'wb') as f:

     

       43
       43
       +
                   pickle.dump(self.__cache, f)

     

       44
       44
       +
       

     

       45
       45
       +
           @override

     

       46
       46
       +
           def load_cache(self, path: Path):

     

       47
       47
       +
               if path.exists():

     

       48
       48
       +
                   with open(path, 'rb') as f:

     

       49
       49
       +
                       self.__cache = pickle.load(f)

+3 -1

util/util.py

···

       1
       1
        
       import logging

     

       2
       2
        
       import sys

     

       3
       3
        
       import os

     

       4
       4
       -
       from typing import Any

     

       4
       4
       +
       from typing import Any, Callable

     

       5
       5
        
       

     

       6
       6
        
       import env

     

       7
       7
       +
       

     

       8
       8
       +
       shutdown_hook: list[Callable[[], None]] = []

     

       7
       9
        
       

     

       8
       10
        
       logging.basicConfig(stream=sys.stderr, level=logging.DEBUG if env.DEV else logging.INFO)

     

       9
       11
        
       LOGGER = logging.getLogger("XPost")