qena-shared-lib 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

qena_shared_lib/remotelogging/__init__.py

@@ -0,0 +1,15 @@
+from . import logstash
+from ._base import (
+    BaseRemoteLogSender,
+    LogLevel,
+    RemoteLogRecord,
+    SenderResponse,
+)
+
+__all__ = [
+    "BaseRemoteLogSender",
+    "LogLevel",
+    "logstash",
+    "RemoteLogRecord",
+    "SenderResponse",
+]

qena_shared_lib/{logstash → remotelogging}/_base.py

@@ -17,9 +17,9 @@ from ..logging import LoggerProvider
 from ..utils import AsyncEventLoopMixin
 
 __all__ = [
-    "BaseLogstashSender",
+    "BaseRemoteLogSender",
     "LogLevel",
-    "LogstashLogRecord",
+    "RemoteLogRecord",
 ]
 
 
@@ -30,7 +30,7 @@ class LogLevel(Enum):
     ERROR = 3
 
 
-class LogstashLogRecord:
+class RemoteLogRecord:
     def __init__(
         self,
         message: str,
@@ -139,8 +139,9 @@ class LogstashLogRecord:
 
     def __repr__(self) -> str:
         return (
-            "LogstashLogRecord (\n\tlevel : `%s`,\n\tmessage : `%s`,\n\ttags : %s,\n\tlabel : %s,\n\terror_type : `%s`,\n\terror_message: `%s`\n)%s"
+            "%s (\n\tlevel : `%s`,\n\tmessage : `%s`,\n\ttags : %s,\n\tlabel : %s,\n\terror_type : `%s`,\n\terror_message: `%s`\n)%s"
             % (
+                self.__class__.__name__,
                 self._log_level.name,
                 self._message,
                 self._tags or [],
@@ -161,31 +162,6 @@ class LogstashLogRecord:
     def log_retries(self, log_retries: int) -> None:
         self._log_retries = log_retries
 
-    def to_dict(self) -> dict:
-        log: dict[str, str | list | dict[str, str]] = {
-            "message": self._message,
-            "service.name": self._service_name,
-            "log.level": self._log_level.name.lower(),
-            "log.logger": self._log_logger,
-        }
-
-        if self._tags is not None:
-            log["tags"] = self._tags
-
-        if self._labels is not None:
-            log["labels"] = self._labels
-
-        if self._error_type is not None:
-            log["error.type"] = self._error_type
-
-        if self._error_message is not None:
-            log["error.message"] = self._error_message
-
-        if self._error_stack_trace is not None:
-            log["error.stack_trace"] = self._error_stack_trace
-
-        return log
-
 
 @dataclass
 class SenderResponse:
@@ -198,20 +174,20 @@ class EndOfLogMarker:
     pass
 
 
-class BaseLogstashSender(AsyncEventLoopMixin):
-    LOGSTASH_LOGS = Counter(
-        name="successful_logstash_logs",
-        documentation="Successfully sent logstash log count",
+class BaseRemoteLogSender(AsyncEventLoopMixin):
+    REMOTE_LOGS = Counter(
+        name="successful_remote_logs",
+        documentation="Successfully sent remote log count",
         labelnames=["log_level"],
     )
-    FAILED_LOGSTASH_LOGS = Counter(
-        name="failed_logstash_logs",
-        documentation="Failed logstash log count",
+    FAILED_REMOTE_LOGS = Counter(
+        name="failed_remote_logs",
+        documentation="Failed remote log count",
         labelnames=["log_level", "exception"],
     )
-    LOGSTASH_SENDER_STATE = PrometheusEnum(
-        name="logstash_sender_state",
-        documentation="Logstash sender state",
+    REMOTE_SENDER_STATE = PrometheusEnum(
+        name="remote_sender_state",
+        documentation="Remote sender state",
         states=["running", "stopped"],
     )
 
@@ -222,25 +198,27 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         log_queue_size: int = 1000,
         failed_log_queue_size: int = 1000,
     ) -> None:
-        self._sender = f"qena_shared_lib.logstash.{self.__class__.__name__}"
+        self._sender = (
+            f"qena_shared_lib.remotelogging.{self.__class__.__name__}"
+        )
         self._service_name = service_name
         self._max_log_retry = max_log_retry
         self._started = False
         self._closed = False
-        self._log_queue: Queue[LogstashLogRecord | EndOfLogMarker] = Queue(
+        self._log_queue: Queue[RemoteLogRecord | EndOfLogMarker] = Queue(
             log_queue_size
         )
-        self._dead_letter_log_queue: Queue[
-            LogstashLogRecord | EndOfLogMarker
-        ] = Queue(failed_log_queue_size)
+        self._dead_letter_log_queue: Queue[RemoteLogRecord | EndOfLogMarker] = (
+            Queue(failed_log_queue_size)
+        )
         self._level = LogLevel.INFO
         self._logger = LoggerProvider.default().get_logger(
-            f"logstash.{self.__class__.__name__.lower()}"
+            f"remotelogging.{self.__class__.__name__.lower()}"
         )
 
     async def start(self) -> None:
         if self._started:
-            raise RuntimeError("logstash sender already started")
+            raise RuntimeError("remote sender already started")
 
         self._started = True
         self._closed = False
@@ -254,10 +232,10 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             self._on_log_flusher_closed
         )
         self._logger.info(
-            "logstash logger `%s` started accepting logs",
+            "remote logger `%s` started accepting logs",
             self.__class__.__name__,
         )
-        self.LOGSTASH_SENDER_STATE.state("running")
+        self.REMOTE_SENDER_STATE.state("running")
 
     def _hook_on_start(self) -> None:
         pass
@@ -265,7 +243,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
     async def _hook_on_start_async(self) -> None:
         pass
 
-    def _on_log_flusher_closed(self, task: Task) -> None:
+    def _on_log_flusher_closed(self, task: Task[None]) -> None:
         if task.cancelled():
             self._close_future.set_result(None)
 
@@ -283,13 +261,13 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             self._hook_on_stop_async(),
         ).add_done_callback(self._on_close_hook_done)
 
-    def _on_close_hook_done(self, task_or_future: Task | Future) -> None:
-        if task_or_future.cancelled():
+    def _on_close_hook_done(self, future: Future[tuple[None, None]]) -> None:
+        if future.cancelled():
             self._close_future.set_result(None)
 
             return
 
-        exception = task_or_future.exception()
+        exception = future.exception()
 
         if exception is not None:
             self._close_future.set_exception(exception)
@@ -298,12 +276,12 @@ class BaseLogstashSender(AsyncEventLoopMixin):
 
         self._close_future.set_result(None)
         self._logger.debug(
-            "logstash http logger closed, will no longer accept logs"
+            "remote http logger closed, will no longer accept logs"
         )
 
-    def stop(self) -> Future:
+    def stop(self) -> Future[None]:
         if self._closed:
-            raise RuntimeError("logstash sender already closed")
+            raise RuntimeError("remote sender already closed")
 
         self._closed = True
         self._started = False
@@ -318,10 +296,10 @@ class BaseLogstashSender(AsyncEventLoopMixin):
 
         return self._close_future
 
-    def _on_close_future_done(self, future: Future) -> None:
+    def _on_close_future_done(self, future: Future[None]) -> None:
         del future
 
-        self.LOGSTASH_SENDER_STATE.state("stopped")
+        self.REMOTE_SENDER_STATE.state("stopped")
 
     def _hook_on_stop(self) -> None:
         pass
@@ -340,7 +318,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             if not self._dead_letter_log_queue.empty():
                 log = await self._dead_letter_log_queue.get()
 
-                if isinstance(log, LogstashLogRecord):
+                if isinstance(log, RemoteLogRecord):
                     if log.log_retries >= self._max_log_retry:
                         self._logger.exception(
                             "failed to send log too many times, falling back to stdout or stderr. \n%r",
@@ -368,9 +346,9 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             except Exception as e:
                 self._put_to_dead_letter_log_queue(log)
                 self._logger.exception(
-                    "error occurred while sending log to logstash"
+                    "error occurred while sending log to remote logging facility"
                 )
-                self.FAILED_LOGSTASH_LOGS.labels(
+                self.FAILED_REMOTE_LOGS.labels(
                     log_level=log.log_level.name, exception=e.__class__.__name__
                 ).inc()
 
@@ -393,15 +371,17 @@ class BaseLogstashSender(AsyncEventLoopMixin):
                     sender_response.reason or "No reason",
                 )
             else:
-                self.LOGSTASH_LOGS.labels(log_level=log.log_level.name).inc()
-                self._logger.debug("log sent to logstash.\n%r", log)
+                self.REMOTE_LOGS.labels(log_level=log.log_level.name).inc()
+                self._logger.debug(
+                    "log sent to remote logging facility.\n%r", log
+                )
 
-    async def _send(self, log: LogstashLogRecord) -> SenderResponse:
+    async def _send(self, log: RemoteLogRecord) -> SenderResponse:
         del log
 
         raise NotImplementedError()
 
-    def _put_to_dead_letter_log_queue(self, log: LogstashLogRecord) -> None:
+    def _put_to_dead_letter_log_queue(self, log: RemoteLogRecord) -> None:
         if self._closed:
             self._logger.error(
                 "%s logger closed, falling back to stdout or stderr.\n%r",
@@ -524,7 +504,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         exception: BaseException | None = None,
     ) -> None:
         if self._closed:
-            self._logger.warning("Logstash http logger is already close")
+            self._logger.warning("Remote logger is already close")
 
             return
 
@@ -551,8 +531,8 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ) -> LogstashLogRecord:
-        log = LogstashLogRecord(
+    ) -> RemoteLogRecord:
+        log = RemoteLogRecord(
             message=message,
             service_name=self._service_name,
             log_level=level,

qena_shared_lib/remotelogging/logstash/__init__.py

@@ -0,0 +1,9 @@
+from ._base import BaseLogstashSender
+from ._http_sender import HTTPSender
+from ._tcp_sender import TCPSender
+
+__all__ = [
+    "BaseLogstashSender",
+    "HTTPSender",
+    "TCPSender",
+]

qena_shared_lib/remotelogging/logstash/_base.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+from .._base import BaseRemoteLogSender, RemoteLogRecord
+
+
+class BaseLogstashSender(BaseRemoteLogSender):
+    def remote_log_record_to_ecs(self, log: RemoteLogRecord) -> dict[str, Any]:
+        log_dict: dict[str, Any] = {
+            "message": log.message,
+            "service.name": log.service_name,
+            "log.level": log.log_level.name.lower(),
+            "log.logger": log.log_logger,
+        }
+
+        if log.tags is not None:
+            log_dict["tags"] = log.tags
+
+        if log.labels is not None:
+            log_dict["labels"] = log.labels
+
+        error_type, error_message, error_stack_trace = log.error
+
+        if error_type is not None:
+            log_dict["error.type"] = error_type
+
+        if error_message is not None:
+            log_dict["error.message"] = error_message
+
+        if error_stack_trace is not None:
+            log_dict["error.stack_trace"] = error_stack_trace
+
+        return log_dict

qena_shared_lib/{logstash → remotelogging/logstash}/_http_sender.py

@@ -1,7 +1,8 @@
 from httpx import AsyncClient, Timeout
 
-from ..logging import LoggerProvider
-from ._base import BaseLogstashSender, LogstashLogRecord, SenderResponse
+from ...logging import LoggerProvider
+from .._base import RemoteLogRecord, SenderResponse
+from ._base import BaseLogstashSender
 
 __all__ = ["HTTPSender"]
 
@@ -42,10 +43,10 @@ class HTTPSender(BaseLogstashSender):
             "logstash.httpsender"
         )
 
-    async def _send(self, log: LogstashLogRecord) -> SenderResponse:
+    async def _send(self, log: RemoteLogRecord) -> SenderResponse:
         send_log_response = await self._client.post(
             url=self._url,
-            json=log.to_dict(),
+            json=self.remote_log_record_to_ecs(log),
         )
 
         if not send_log_response.is_success:

qena_shared_lib/{logstash → remotelogging/logstash}/_tcp_sender.py

@@ -1,9 +1,11 @@
 from asyncio import StreamWriter, open_connection
+from typing import Any
 
 from pydantic_core import to_json
 
-from ..logging import LoggerProvider
-from ._base import BaseLogstashSender, LogstashLogRecord, SenderResponse
+from ...logging import LoggerProvider
+from .._base import RemoteLogRecord, SenderResponse
+from ._base import BaseLogstashSender
 
 __all__ = ["TCPSender"]
 
@@ -28,8 +30,8 @@ class TCPSender(BaseLogstashSender):
         self._client = AsyncTcpClient(host=host, port=port)
         self._logger = LoggerProvider.default().get_logger("logstash.tcpsender")
 
-    async def _send(self, log: LogstashLogRecord) -> SenderResponse:
-        await self._client.write(log.to_dict())
+    async def _send(self, log: RemoteLogRecord) -> SenderResponse:
+        await self._client.write(self.remote_log_record_to_ecs(log))
 
         return SenderResponse(sent=True)
 
@@ -47,7 +49,7 @@ class AsyncTcpClient:
         self._writer: StreamWriter | None = None
         self._client_closed = False
 
-    async def write(self, json: dict) -> None:
+    async def write(self, json: dict[str, Any]) -> None:
         if self._client_closed:
             raise RuntimeError("async tcp client already closed")
 

qena_shared_lib/scheduler.py

@@ -1,11 +1,11 @@
-from asyncio import Future, Task, iscoroutinefunction, sleep
+from asyncio import Future, Task, gather, iscoroutinefunction, sleep
 from dataclasses import dataclass
 from datetime import datetime
 from functools import partial
 from importlib import import_module
 from inspect import signature
 from os import name as osname
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, TypeVar, cast
 from zoneinfo import ZoneInfo
 
 from cronsim import CronSim
@@ -14,7 +14,7 @@ from punq import Container, Scope
 
 from .dependencies.miscellaneous import validate_annotation
 from .logging import LoggerProvider
-from .logstash import BaseLogstashSender
+from .remotelogging import BaseRemoteLogSender
 from .utils import AsyncEventLoopMixin
 
 __all__ = [
@@ -32,7 +32,7 @@ SCHEDULED_TASK_ATTRIBUTE = "__scheduled_task__"
 
 @dataclass
 class ScheduledTask:
-    task: Callable
+    task: Callable[..., None]
     cron_expression: str
     zone_info: ZoneInfo | None
 
@@ -78,8 +78,8 @@ class Scheduler:
 
     def schedule(
         self, cron_expression: str, timezone: str | None = None
-    ) -> Callable[[Callable], Callable]:
-        def wrapper(task: Callable) -> Callable:
+    ) -> Callable[[Callable[..., None]], Callable[..., None]]:
+        def wrapper(task: Callable[..., Any]) -> Callable[..., None]:
             self.add_task(
                 task=task, cron_expression=cron_expression, timezone=timezone
             )
@@ -89,7 +89,10 @@ class Scheduler:
         return wrapper
 
     def add_task(
-        self, task: Callable, cron_expression: str, timezone: str | None
+        self,
+        task: Callable[..., Any],
+        cron_expression: str,
+        timezone: str | None,
     ) -> None:
         self._scheduled_tasks.append(
             ScheduledTask(
@@ -121,8 +124,8 @@ class ScheduledTaskMeta:
 
 def schedule(
     cron_expression: str, *, timezone: str | None = None
-) -> Callable[[Callable], Callable]:
-    def wrapper(task: Callable) -> Callable:
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    def wrapper(task: Callable[..., Any]) -> Callable[..., Any]:
         setattr(
             task,
             SCHEDULED_TASK_ATTRIBUTE,
@@ -145,7 +148,7 @@ class SchedulerBase:
                 f"{self.__class__.__name__} not a scheduler, possibly no annotated with either `Scheduler`"
             )
 
-        return scheduler
+        return cast(Scheduler, scheduler)
 
     def register_scheduled_tasks(self) -> Scheduler:
         scheduler = self.get_scheduler()
@@ -186,14 +189,16 @@ class ScheduleManager(AsyncEventLoopMixin):
 
     def __init__(
         self,
-        logstash: BaseLogstashSender,
+        remote_logger: BaseRemoteLogSender,
         container: Container | None = None,
     ) -> None:
         self._container = container or Container()
-        self._logstash = logstash
+        self._remote_logger = remote_logger
         self._scheduled_tasks: list[ScheduledTask] = []
         self._next_run_in: int | None = None
-        self._scheduler_task: Task | None = None
+        self._scheduler_task: Task[None] | None = None
+        self._scheduled_tasks_or_futures: list[Task[Any] | Future[Any]] = []
+        self._stopped = False
         self._logger = LoggerProvider.default().get_logger("schedule_manager")
 
     def include_scheduler(
@@ -260,20 +265,26 @@ class ScheduleManager(AsyncEventLoopMixin):
         ]:
             self._scheduled_tasks.extend(scheduler.scheduled_tasks)
 
-    def stop(self) -> None:
+    async def stop(self) -> None:
+        self._stopped = True
+
+        _ = await gather(
+            *self._scheduled_tasks_or_futures, return_exceptions=True
+        )
+
         if self._scheduler_task is not None and not self._scheduler_task.done():
             self._scheduler_task.cancel()
 
         self.SCHEDULE_MANAGER_STATE.state("stopped")
 
-    def _on_scheduler_done(self, task: Task) -> None:
+    def _on_scheduler_done(self, task: Task[None]) -> None:
         if task.cancelled():
             return
 
         exception = task.exception()
 
         if exception is not None:
-            self._logstash.error(
+            self._remote_logger.error(
                 message="error occured in schedule manager",
                 exception=exception,
             )
@@ -309,7 +320,7 @@ class ScheduleManager(AsyncEventLoopMixin):
         return True
 
     async def _run_scheduler(self) -> None:
-        while True:
+        while not self._stopped:
             self._calculate_next_schedule()
             self._logger.debug(
                 "next tasks will be executed after `%d` seconds",
@@ -331,15 +342,23 @@ class ScheduleManager(AsyncEventLoopMixin):
                     continue
 
                 args = self._resolve_dependencies(scheduled_task)
+                scheduled_task_or_future: Task[Any] | Future[Any] | None = None
 
                 if iscoroutinefunction(scheduled_task.task):
-                    self.loop.create_task(
+                    scheduled_task_or_future = self.loop.create_task(
                         scheduled_task.task(**args)
-                    ).add_done_callback(self._on_task_done)
+                    )
                 else:
-                    self.loop.run_in_executor(
-                        executor=None, func=scheduled_task.task
-                    ).add_done_callback(partial(self._on_task_done, **args))
+                    scheduled_task_or_future = self.loop.run_in_executor(
+                        executor=None, func=partial(scheduled_task.task, **args)
+                    )
+
+                assert scheduled_task_or_future is not None
+
+                scheduled_task_or_future.add_done_callback(self._on_task_done)
+                self._scheduled_tasks_or_futures.append(
+                    scheduled_task_or_future
+                )
 
                 scheduled_task.ran = True
 
@@ -356,14 +375,20 @@ class ScheduleManager(AsyncEventLoopMixin):
 
         return args
 
-    def _on_task_done(self, task_or_future: Task | Future) -> None:
+    def _on_task_done(self, task_or_future: Task[Any] | Future[Any]) -> None:
+        if (
+            not self._stopped
+            and task_or_future in self._scheduled_tasks_or_futures
+        ):
+            self._scheduled_tasks_or_futures.remove(task_or_future)
+
         if task_or_future.cancelled():
             return
 
         exception = task_or_future.exception()
 
         if exception is not None:
-            self._logstash.error(
+            self._remote_logger.error(
                 message="error occured while executing task",
                 exception=exception,
             )

qena_shared_lib/security.py

@@ -1,7 +1,7 @@
 from asyncio import Future
 from enum import Enum
 from os import environ
-from typing import AbstractSet, Annotated, Any
+from typing import AbstractSet, Annotated, Any, cast
 
 from fastapi import Depends, Header
 from jwt import JWT, AbstractJWKBase, jwk_from_dict
@@ -129,7 +129,7 @@ async def extract_user_info(
             extract_exc_info=True,
         ) from e
 
-    return user_info
+    return cast(UserInfo, user_info)
 
 
 class PermissionMatch(Enum):

qena_shared_lib/utils.py

@@ -1,5 +1,4 @@
 from asyncio import AbstractEventLoop, get_running_loop
-from functools import lru_cache
 
 from pydantic import TypeAdapter
 
@@ -7,10 +6,17 @@ __all__ = ["AsyncEventLoopMixin", "TypeAdapterCache"]
 
 
 class AsyncEventLoopMixin:
+    _LOOP: AbstractEventLoop | None = None
+
     @property
-    @lru_cache
     def loop(self) -> AbstractEventLoop:
-        return get_running_loop()
+        if self._LOOP is None:
+            self._LOOP = get_running_loop()
+
+        return self._LOOP
+
+    def init(self) -> None:
+        self._LOOP = get_running_loop()
 
 
 class TypeAdapterCache: