omserv 0.0.0.dev7__py3-none-any.whl

omserv/node/sql.py

@@ -0,0 +1,131 @@
+"""
+TODO:
+ - move to omlish.sql, pg-specific trigger compiler - https://docs.sqlalchemy.org/en/20/core/compiler.html
+"""
+import textwrap
+import typing as ta
+
+import sqlalchemy as sa
+import sqlalchemy.ext.compiler
+
+
+##
+
+
+class IdMixin:
+    _id = sa.Column(
+        sa.Integer,
+        nullable=False,
+        primary_key=True,
+        autoincrement=True,
+    )
+
+
+##
+
+
+class utcnow(sa.sql.expression.FunctionElement):  # noqa
+    inherit_cache = True
+    type = sa.TIMESTAMP()
+
+
+@sa.ext.compiler.compiles(utcnow)
+def _compile_utcnow(
+        element: utcnow,
+        compiler: sa.sql.compiler.SQLCompiler,
+        **kw: ta.Any,
+) -> str:
+    return "timezone('utc', now())"
+
+
+##
+
+
+class TimestampsMixin:
+    created_at = sa.Column(
+        sa.TIMESTAMP(timezone=True),
+        server_default=utcnow(),
+        nullable=False,
+    )
+
+    updated_at = sa.Column(
+        sa.TIMESTAMP(timezone=True),
+        server_default=utcnow(),
+        server_onupdate=sa.schema.FetchedValue(for_update=True),
+        nullable=False,
+    )
+
+
+##
+
+
+CREATE_UPDATED_AT_FUNCTION_STATEMENT = textwrap.dedent("""
+    create or replace function set_updated_at_timestamp()
+    returns trigger as $$
+    begin
+        new.updated_at = now() at time zone 'utc';
+        return new;
+    end;
+    $$ language 'plpgsql';
+""")
+
+
+##
+
+
+def get_update_at_trigger_name(table_name: str) -> str:
+    return f'trigger__updated_at_{table_name}'
+
+
+#
+
+
+class CreateUpdateAtTrigger(sa.schema.DDLElement):
+    inherit_cache = False
+
+    def __init__(self, table_name: str) -> None:
+        super().__init__()
+        self.table_name = table_name
+
+
+@sa.ext.compiler.compiles(CreateUpdateAtTrigger)
+def _compile_create_update_at_trigger(
+        element: CreateUpdateAtTrigger,
+        compiler: sa.sql.compiler.SQLCompiler,
+        **kw: ta.Any,
+):
+    return textwrap.dedent(f"""
+        create or replace trigger {get_update_at_trigger_name(element.table_name)}
+            before update
+            on {element.table_name}
+            for each row
+            execute procedure set_updated_at_timestamp()
+    """)
+
+
+#
+
+
+class DropUpdateAtTrigger(sa.schema.DDLElement):
+    inherit_cache = False
+
+    def __init__(self, table_name: str) -> None:
+        super().__init__()
+        self.table_name = table_name
+
+
+@sa.ext.compiler.compiles(DropUpdateAtTrigger)
+def _compile_drop_update_at_trigger(
+        element: DropUpdateAtTrigger,
+        compiler: sa.sql.compiler.SQLCompiler,
+        **kw: ta.Any,
+) -> str:
+    return f'drop trigger if exists {get_update_at_trigger_name(element.table_name)} on {element.table_name}'
+
+
+#
+
+
+def install_updated_at_trigger(metadata: sa.MetaData, table_name: str) -> None:
+    sa.event.listen(metadata, 'after_create', CreateUpdateAtTrigger(table_name))
+    sa.event.listen(metadata, 'before_drop', DropUpdateAtTrigger(table_name))

omserv/secrets.py

@@ -0,0 +1,12 @@
+import os.path
+import typing as ta
+
+import yaml
+
+
+SECRETS_PATH = os.getenv('SECRETS_PATH', os.path.expanduser('~/Dropbox/.dotfiles/secrets.yml'))
+
+
+def load_secrets() -> dict[str, ta.Any]:
+    with open(SECRETS_PATH) as f:
+        return yaml.safe_load(f)

omserv/server/__init__.py

@@ -0,0 +1,18 @@
+"""
+Based on https://github.com/pgjones/hypercorn
+
+TODO:
+ - !!! error handling jfc
+ - add ssl back lol
+ - events as dc's
+ - injectify
+ - lifecycle / otp-ify
+ - configify
+
+Lookit:
+ - https://github.com/davidbrochart/anycorn
+ - https://github.com/encode/starlette
+ - https://github.com/tiangolo/fastapi
+"""
+from .config import Config  # noqa
+from .workers import serve  # noqa

omserv/server/config.py

@@ -0,0 +1,51 @@
+import dataclasses as dc
+import typing as ta
+
+
+BYTES = 1
+OCTETS = 1
+SECONDS = 1.0
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class Config:
+    bind: ta.Sequence[str] = ('127.0.0.1:8000',)
+
+    umask: int | None = None
+    user: int | None = None
+    group: int | None = None
+
+    workers: int = 0
+
+    max_app_queue_size: int = 10
+
+    startup_timeout = 60 * SECONDS
+    shutdown_timeout = 60 * SECONDS
+
+    server_names: ta.Sequence[str] = ()
+
+    max_requests: int | None = None
+    max_requests_jitter: int = 0
+
+    backlog: int = 100
+
+    graceful_timeout: float = 3 * SECONDS
+
+    keep_alive_timeout: float = 5 * SECONDS
+    keep_alive_max_requests: int = 1000
+
+    read_timeout: int | None = None
+
+    h11_max_incomplete_size: int = 16 * 1024 * BYTES
+    h11_pass_raw_headers: bool = False
+
+    h2_max_concurrent_streams: int = 100
+    h2_max_header_list_size: int = 2 ** 16
+    h2_max_inbound_frame_size: int = 2 ** 14 * OCTETS
+
+    include_date_header: bool = True
+    include_server_header: bool = True
+    alt_svc_headers: list[str] = dc.field(default_factory=list)
+
+    websocket_max_message_size: int = 16 * 1024 * 1024 * BYTES
+    websocket_ping_interval: float | None = None

omserv/server/debug.py

@@ -0,0 +1,14 @@
+import logging
+import sys
+
+from omlish.diag import pydevd as pdu
+
+
+log = logging.getLogger(__name__)
+
+
+def handle_error_debug(e: BaseException) -> None:
+    exc_info = sys.exc_info()
+    log.warning('Launching debugger')
+    pdu.debug_unhandled_exception(exc_info)
+    log.warning('Done debugging')

omserv/server/events.py

@@ -0,0 +1,83 @@
+import dataclasses as dc
+
+from omlish import lang
+
+
+##
+
+
+class ServerEvent(lang.Abstract):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class RawData(ServerEvent):
+    data: bytes
+    address: tuple[str, int] | None = None
+
+
+@dc.dataclass(frozen=True)
+class Closed(ServerEvent):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class Updated(ServerEvent):
+    idle: bool
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class ProtocolEvent:
+    stream_id: int
+
+
+@dc.dataclass(frozen=True)
+class Request(ProtocolEvent):
+    headers: list[tuple[bytes, bytes]]
+    http_version: str
+    method: str
+    raw_path: bytes
+
+
+@dc.dataclass(frozen=True)
+class Body(ProtocolEvent):
+    data: bytes
+
+
+@dc.dataclass(frozen=True)
+class EndBody(ProtocolEvent):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class Data(ProtocolEvent):
+    data: bytes
+
+
+@dc.dataclass(frozen=True)
+class EndData(ProtocolEvent):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class Response(ProtocolEvent):
+    headers: list[tuple[bytes, bytes]]
+    status_code: int
+
+
+@dc.dataclass(frozen=True)
+class InformationalResponse(ProtocolEvent):
+    headers: list[tuple[bytes, bytes]]
+    status_code: int
+
+    def __post_init__(self) -> None:
+        if self.status_code >= 200 or self.status_code < 100:
+            raise ValueError(f'Status code must be 1XX not {self.status_code}')
+
+
+@dc.dataclass(frozen=True)
+class StreamClosed(ProtocolEvent):
+    pass

omserv/server/headers.py

@@ -0,0 +1,36 @@
+import time
+import wsgiref.handlers
+
+from .config import Config
+
+
+def _now() -> float:
+    return time.time()
+
+
+def response_headers(config: Config, protocol: str) -> list[tuple[bytes, bytes]]:
+    headers = []
+    if config.include_date_header:
+        headers.append((b'date', wsgiref.handlers.format_date_time(_now()).encode('ascii')))
+    if config.include_server_header:
+        headers.append((b'server', f'omlicorn-{protocol}'.encode('ascii')))
+
+    for alt_svc_header in config.alt_svc_headers:
+        headers.append((b'alt-svc', alt_svc_header.encode()))
+
+    return headers
+
+
+def filter_pseudo_headers(headers: list[tuple[bytes, bytes]]) -> list[tuple[bytes, bytes]]:
+    filtered_headers: list[tuple[bytes, bytes]] = [(b'host', b'')]  # Placeholder
+    authority = None
+    host = b''
+    for name, value in headers:
+        if name == b':authority':  # h2 & h3 libraries validate this is present
+            authority = value
+        elif name == b'host':
+            host = value
+        elif name[0] != b':'[0]:
+            filtered_headers.append((name, value))
+    filtered_headers[0] = (b'host', authority if authority is not None else host)
+    return filtered_headers

omserv/server/lifespans.py

@@ -0,0 +1,132 @@
+import logging
+import typing as ta
+
+import anyio
+import anyio.abc
+import anyio.from_thread
+import anyio.to_thread
+
+from .config import Config
+from .debug import handle_error_debug
+from .types import AppWrapper
+from .types import AsgiReceiveEvent
+from .types import AsgiSendEvent
+from .types import LifespanScope
+from .types import UnexpectedMessageError
+
+
+log = logging.getLogger(__name__)
+
+
+class LifespanTimeoutError(Exception):
+    def __init__(self, stage: str) -> None:
+        super().__init__(
+            f'Timeout whilst awaiting {stage}. Your application may not support the Asgi Lifespan '
+            f'protocol correctly, alternatively the {stage}_timeout configuration is incorrect.',
+        )
+
+
+class LifespanFailureError(Exception):
+    def __init__(self, stage: str, message: str) -> None:
+        super().__init__(f"Lifespan failure in {stage}. '{message}'")
+
+
+class Lifespan:
+    def __init__(self, app: AppWrapper, config: Config) -> None:
+        super().__init__()
+        self.app = app
+        self.config = config
+        self.startup = anyio.Event()
+        self.shutdown = anyio.Event()
+        self.app_send_channel, self.app_receive_channel = anyio.create_memory_object_stream[ta.Any](
+            config.max_app_queue_size,
+        )
+        self.supported = True
+
+    async def handle_lifespan(
+            self,
+            *,
+            task_status: anyio.abc.TaskStatus[ta.Any] = anyio.TASK_STATUS_IGNORED,
+    ) -> None:
+        task_status.started()
+        scope: LifespanScope = {
+            'type': 'lifespan',
+            'asgi': {'spec_version': '2.0', 'version': '3.0'},
+        }
+
+        try:
+            await self.app(
+                scope,
+                self.asgi_receive,
+                self.asgi_send,
+                anyio.to_thread.run_sync,
+                anyio.from_thread.run,
+            )
+        except (LifespanFailureError, anyio.get_cancelled_exc_class()):
+            raise
+        except (BaseExceptionGroup, Exception) as error:
+            handle_error_debug(error)
+
+            if isinstance(error, BaseExceptionGroup):
+                failure_error = error.subgroup(LifespanFailureError)
+                if failure_error is not None:
+                    # Lifespan failures should crash the server
+                    raise failure_error  # noqa
+                reraise_error = error.subgroup((LifespanFailureError, anyio.get_cancelled_exc_class()))
+                if reraise_error is not None:
+                    raise reraise_error  # noqa
+
+            self.supported = False
+            if not self.startup.is_set():
+                log.warning('Asgi Framework Lifespan error, continuing without Lifespan support')
+            elif not self.shutdown.is_set():
+                log.exception('Asgi Framework Lifespan error, shutdown without Lifespan support')
+            else:
+                log.exception('Asgi Framework Lifespan errored after shutdown.')
+
+        finally:
+            self.startup.set()
+            self.shutdown.set()
+            await self.app_send_channel.aclose()
+            await self.app_receive_channel.aclose()
+
+    async def wait_for_startup(self) -> None:
+        if not self.supported:
+            return
+
+        await self.app_send_channel.send({'type': 'lifespan.startup'})
+        try:
+            with anyio.fail_after(self.config.startup_timeout):
+                await self.startup.wait()
+        except TimeoutError as error:
+            raise LifespanTimeoutError('startup') from error
+
+    async def wait_for_shutdown(self) -> None:
+        if not self.supported:
+            return
+
+        await self.app_send_channel.send({'type': 'lifespan.shutdown'})
+        try:
+            with anyio.fail_after(self.config.shutdown_timeout):
+                await self.shutdown.wait()
+        except TimeoutError as error:
+            raise LifespanTimeoutError('startup') from error
+
+    async def asgi_receive(self) -> AsgiReceiveEvent:
+        return await self.app_receive_channel.receive()
+
+    async def asgi_send(self, message: AsgiSendEvent) -> None:
+        if message['type'] == 'lifespan.startup.complete':
+            self.startup.set()
+
+        elif message['type'] == 'lifespan.shutdown.complete':
+            self.shutdown.set()
+
+        elif message['type'] == 'lifespan.startup.failed':
+            raise LifespanFailureError('startup', message.get('message', ''))
+
+        elif message['type'] == 'lifespan.shutdown.failed':
+            raise LifespanFailureError('shutdown', message.get('message', ''))
+
+        else:
+            raise UnexpectedMessageError(message['type'])

omserv/server/multiprocess.py

@@ -0,0 +1,157 @@
+import functools
+import multiprocessing
+import multiprocessing.connection
+import multiprocessing.context
+import multiprocessing.synchronize
+import pickle
+import platform
+import signal
+import time
+import typing as ta
+
+import anyio
+
+from .config import Config
+from .sockets import Sockets
+from .sockets import create_sockets
+from .types import AsgiFramework
+from .types import wrap_app
+from .workers import serve
+
+
+async def check_multiprocess_shutdown_event(
+        shutdown_event: multiprocessing.synchronize.Event,
+        sleep: ta.Callable[[float], ta.Awaitable[ta.Any]],
+) -> None:
+    while True:
+        if shutdown_event.is_set():
+            return
+        await sleep(0.1)
+
+
+def _multiprocess_serve(
+        app: AsgiFramework,
+        config: Config,
+        sockets: Sockets | None = None,
+        shutdown_event: multiprocessing.synchronize.Event | None = None,
+) -> None:
+    if sockets is not None:
+        for sock in sockets.insecure_sockets:
+            sock.listen(config.backlog)
+
+    shutdown_trigger = None
+    if shutdown_event is not None:
+        shutdown_trigger = functools.partial(check_multiprocess_shutdown_event, shutdown_event, anyio.sleep)
+
+    anyio.run(
+        functools.partial(
+            serve,
+            wrap_app(app),
+            config,
+            sockets=sockets,
+            shutdown_trigger=shutdown_trigger,
+        ),
+        # backend='trio',
+    )
+
+
+def serve_multiprocess(
+        app: AsgiFramework,
+        config: Config,
+) -> int:
+    sockets = create_sockets(config)
+
+    exitcode = 0
+    ctx = multiprocessing.get_context('spawn')
+
+    active = True
+    shutdown_event = ctx.Event()
+
+    def shutdown(*args: ta.Any) -> None:
+        nonlocal active, shutdown_event
+        shutdown_event.set()
+        active = False
+
+    processes: list[multiprocessing.Process] = []
+    while active:
+        # Ignore SIGINT before creating the processes, so that they inherit the signal handling. This means that the
+        # shutdown function controls the shutdown.
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+        _populate(
+            processes,
+            app,
+            config,
+            _multiprocess_serve,
+            sockets,
+            shutdown_event,
+            ctx,
+        )
+
+        for signal_name in ('SIGINT', 'SIGTERM', 'SIGBREAK'):
+            if hasattr(signal, signal_name):
+                signal.signal(getattr(signal, signal_name), shutdown)
+
+        multiprocessing.connection.wait(process.sentinel for process in processes)
+
+        exitcode = _join_exited(processes)
+        if exitcode != 0:
+            shutdown_event.set()
+            active = False
+
+    for process in processes:
+        process.terminate()
+
+    exitcode = _join_exited(processes) if exitcode != 0 else exitcode
+
+    for sock in sockets.insecure_sockets:
+        sock.close()
+
+    return exitcode
+
+
+def _populate(
+        processes: list[multiprocessing.Process],
+        app: AsgiFramework,
+        config: Config,
+        worker_func: ta.Callable,
+        sockets: Sockets,
+        shutdown_event: multiprocessing.synchronize.Event,
+        ctx: multiprocessing.context.BaseContext,
+) -> None:
+    num_workers = config.workers or 1
+    if num_workers < 0:
+        num_workers = multiprocessing.cpu_count()
+    for _ in range(num_workers - len(processes)):
+        process = ctx.Process(  # type: ignore
+            target=worker_func,
+            kwargs={
+                'app': app,
+                'config': config,
+                'shutdown_event': shutdown_event,
+                'sockets': sockets,
+            },
+        )
+        process.daemon = True
+        try:
+            process.start()
+        except pickle.PicklingError as error:
+            raise RuntimeError(
+                'Cannot pickle the config, see https://docs.python.org/3/library/pickle.html#pickle-picklable',
+                # noqa: E501
+            ) from error
+        processes.append(process)
+        if platform.system() == 'Windows':
+            time.sleep(0.1)
+
+
+def _join_exited(processes: list[multiprocessing.Process]) -> int:
+    exitcode = 0
+    for index in reversed(range(len(processes))):
+        worker = processes[index]
+        if worker.exitcode is not None:
+            worker.join()
+            exitcode = worker.exitcode if exitcode == 0 else exitcode
+            del processes[index]
+
+    return exitcode

omserv/server/protocols/__init__.py

@@ -0,0 +1 @@
+from .protocols import ProtocolWrapper  # noqa