ominfra 0.0.0.dev102__py3-none-any.whl → 0.0.0.dev104__py3-none-any.whl

ominfra/clouds/aws/logs.py

@@ -62,7 +62,7 @@ class AwsPutLogEventsResponse(AwsDataclass):
 ##
 
 
-class AwsLogMessagePoster:
+class AwsLogMessageBuilder:
     """
     TODO:
      - max_items
@@ -88,7 +88,7 @@ class AwsLogMessagePoster:
             log_group_name: str,
             log_stream_name: str,
             region_name: str,
-            credentials: AwsSigner.Credentials,
+            credentials: ta.Optional[AwsSigner.Credentials],
 
             url: ta.Optional[str] = None,
             service_name: str = DEFAULT_SERVICE_NAME,
@@ -110,11 +110,16 @@ class AwsLogMessagePoster:
             headers = {**headers, **extra_headers}
         self._headers = {k: [v] for k, v in headers.items()}
 
-        self._signer = V4AwsSigner(
-            credentials,
-            region_name,
-            service_name,
-        )
+        signer: ta.Optional[V4AwsSigner]
+        if credentials is not None:
+            signer = V4AwsSigner(
+                credentials,
+                region_name,
+                service_name,
+            )
+        else:
+            signer = None
+        self._signer = signer
 
     #
 
@@ -158,13 +163,14 @@ class AwsLogMessagePoster:
             payload=body,
         )
 
-        sig_headers = self._signer.sign(
-            sig_req,
-            sign_payload=False,
-        )
-        sig_req = dc.replace(sig_req, headers={**sig_req.headers, **sig_headers})
+        if (signer := self._signer) is not None:
+            sig_headers = signer.sign(
+                sig_req,
+                sign_payload=False,
+            )
+            sig_req = dc.replace(sig_req, headers={**sig_req.headers, **sig_headers})
 
-        post = AwsLogMessagePoster.Post(
+        post = AwsLogMessageBuilder.Post(
             url=self._url,
             headers={k: check_single(v) for k, v in sig_req.headers.items()},
             data=sig_req.payload,

ominfra/clouds/gcp/auth.py

@@ -0,0 +1,48 @@
+import json
+import time
+import typing as ta
+
+from omlish import check
+from omlish import http
+from omlish.http import jwt
+
+
+DEFAULT_JWT_SCOPE = 'https://www.googleapis.com/auth/cloud-platform'
+
+
+def generate_gcp_jwt(
+        creds_dct: ta.Mapping[str, ta.Any],
+        *,
+        issued_at: int | None = None,
+        lifetime_s: int = 3600,
+        scope: str = DEFAULT_JWT_SCOPE,
+) -> str:
+    return jwt.generate_jwt(
+        issuer=creds_dct['client_email'],
+        subject=creds_dct['client_email'],
+        audience=creds_dct['token_uri'],
+        issued_at=(issued_at := int(issued_at if issued_at is not None else time.time())),
+        expires_at=issued_at + lifetime_s,
+        scope=scope,
+        key=creds_dct['private_key'],
+        algorithm='RS256',
+    )
+
+
+def get_gcp_access_token(
+        creds_dct: ta.Mapping[str, ta.Any],
+        *,
+        client: http.HttpClient | None = None,
+) -> str:
+    signed_jwt = generate_gcp_jwt(creds_dct)
+    resp = http.request(
+        creds_dct['token_uri'],
+        'POST',
+        data=jwt.build_get_token_body(signed_jwt).encode('utf-8'),
+        headers={
+            http.consts.HEADER_CONTENT_TYPE: http.consts.CONTENT_TYPE_FORM_URLENCODED,
+        },
+        client=client,
+    )
+    resp_dct = json.loads(check.not_none(resp.data).decode('utf-8'))
+    return resp_dct['access_token']

ominfra/deploy/_executor.py

@@ -82,7 +82,7 @@ if sys.version_info < (3, 8):
 ########################################
 
 
-# ../../../../omlish/lite/check.py
+# ../../../../omlish/lite/cached.py
 T = ta.TypeVar('T')
 
 
@@ -112,7 +112,7 @@ class HostConfig:
 # ../../../../omlish/lite/cached.py
 
 
-class cached_nullary:  # noqa
+class _cached_nullary:  # noqa
     def __init__(self, fn):
         super().__init__()
         self._fn = fn
@@ -129,6 +129,10 @@ class cached_nullary:  # noqa
         return bound
 
 
+def cached_nullary(fn: ta.Callable[..., T]) -> ta.Callable[..., T]:
+    return _cached_nullary(fn)
+
+
 ########################################
 # ../../../../omlish/lite/check.py
 
@@ -661,7 +665,7 @@ class DataclassObjMarshaler(ObjMarshaler):
         return {k: m.marshal(getattr(o, k)) for k, m in self.fs.items()}
 
     def unmarshal(self, o: ta.Any) -> ta.Any:
-        return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if self.nonstrict or k in self.fs})
+        return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if not self.nonstrict or k in self.fs})
 
 
 @dc.dataclass(frozen=True)
@@ -721,7 +725,10 @@ class UuidObjMarshaler(ObjMarshaler):
         return uuid.UUID(o)
 
 
-_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = {
+##
+
+
+_DEFAULT_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = {
     **{t: NopObjMarshaler() for t in (type(None),)},
     **{t: CastObjMarshaler(t) for t in (int, float, str, bool)},
     **{t: Base64ObjMarshaler(t) for t in (bytes, bytearray)},
@@ -750,20 +757,19 @@ _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES: ta.Dict[ta.Any, type] = {
 }
 
 
-def register_opj_marshaler(ty: ta.Any, m: ObjMarshaler) -> None:
-    if ty in _OBJ_MARSHALERS:
-        raise KeyError(ty)
-    _OBJ_MARSHALERS[ty] = m
-
-
-def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
+def _make_obj_marshaler(
+        ty: ta.Any,
+        rec: ta.Callable[[ta.Any], ObjMarshaler],
+        *,
+        nonstrict_dataclasses: bool = False,
+) -> ObjMarshaler:
     if isinstance(ty, type):
         if abc.ABC in ty.__bases__:
             impls = [  # type: ignore
                 PolymorphicObjMarshaler.Impl(
                     ity,
                     ity.__qualname__,
-                    get_obj_marshaler(ity),
+                    rec(ity),
                 )
                 for ity in deep_subclasses(ty)
                 if abc.ABC not in ity.__bases__
@@ -779,7 +785,8 @@ def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
         if dc.is_dataclass(ty):
             return DataclassObjMarshaler(
                 ty,
-                {f.name: get_obj_marshaler(f.type) for f in dc.fields(ty)},
+                {f.name: rec(f.type) for f in dc.fields(ty)},
+                nonstrict=nonstrict_dataclasses,
             )
 
     if is_generic_alias(ty):
@@ -789,7 +796,7 @@ def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
             pass
         else:
             k, v = ta.get_args(ty)
-            return MappingObjMarshaler(mt, get_obj_marshaler(k), get_obj_marshaler(v))
+            return MappingObjMarshaler(mt, rec(k), rec(v))
 
         try:
             st = _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES[ta.get_origin(ty)]
@@ -797,33 +804,71 @@ def _make_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
             pass
         else:
             [e] = ta.get_args(ty)
-            return IterableObjMarshaler(st, get_obj_marshaler(e))
+            return IterableObjMarshaler(st, rec(e))
 
         if is_union_alias(ty):
-            return OptionalObjMarshaler(get_obj_marshaler(get_optional_alias_arg(ty)))
+            return OptionalObjMarshaler(rec(get_optional_alias_arg(ty)))
 
     raise TypeError(ty)
 
 
-def get_obj_marshaler(ty: ta.Any) -> ObjMarshaler:
-    try:
-        return _OBJ_MARSHALERS[ty]
-    except KeyError:
-        pass
+##
 
-    p = ProxyObjMarshaler()
-    _OBJ_MARSHALERS[ty] = p
-    try:
-        m = _make_obj_marshaler(ty)
-    except Exception:
-        del _OBJ_MARSHALERS[ty]
-        raise
-    else:
-        p.m = m
+
+_OBJ_MARSHALERS_LOCK = threading.RLock()
+
+_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = dict(_DEFAULT_OBJ_MARSHALERS)
+
+_OBJ_MARSHALER_PROXIES: ta.Dict[ta.Any, ProxyObjMarshaler] = {}
+
+
+def register_opj_marshaler(ty: ta.Any, m: ObjMarshaler) -> None:
+    with _OBJ_MARSHALERS_LOCK:
+        if ty in _OBJ_MARSHALERS:
+            raise KeyError(ty)
         _OBJ_MARSHALERS[ty] = m
+
+
+def get_obj_marshaler(
+        ty: ta.Any,
+        *,
+        no_cache: bool = False,
+        **kwargs: ta.Any,
+) -> ObjMarshaler:
+    with _OBJ_MARSHALERS_LOCK:
+        if not no_cache:
+            try:
+                return _OBJ_MARSHALERS[ty]
+            except KeyError:
+                pass
+
+        try:
+            return _OBJ_MARSHALER_PROXIES[ty]
+        except KeyError:
+            pass
+
+        rec = functools.partial(
+            get_obj_marshaler,
+            no_cache=no_cache,
+            **kwargs,
+        )
+
+        p = ProxyObjMarshaler()
+        _OBJ_MARSHALER_PROXIES[ty] = p
+        try:
+            m = _make_obj_marshaler(ty, rec, **kwargs)
+        finally:
+            del _OBJ_MARSHALER_PROXIES[ty]
+        p.m = m
+
+        if not no_cache:
+            _OBJ_MARSHALERS[ty] = m
         return m
 
 
+##
+
+
 def marshal_obj(o: ta.Any, ty: ta.Any = None) -> ta.Any:
     return get_obj_marshaler(ty if ty is not None else type(o)).marshal(o)
 
@@ -956,6 +1001,24 @@ def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
     return out.decode().strip() if out is not None else None
 
 
+##
+
+
+def subprocess_close(
+        proc: subprocess.Popen,
+        timeout: ta.Optional[float] = None,
+) -> None:
+    # TODO: terminate, sleep, kill
+    if proc.stdout:
+        proc.stdout.close()
+    if proc.stderr:
+        proc.stderr.close()
+    if proc.stdin:
+        proc.stdin.close()
+
+    proc.wait(timeout)
+
+
 ########################################
 # ../base.py
 

ominfra/deploy/poly/_main.py

@@ -34,8 +34,10 @@ if sys.version_info < (3, 8):
 ########################################
 
 
-# ../base.py
+# ../../../../omlish/lite/cached.py
 T = ta.TypeVar('T')
+
+# ../base.py
 ConcernT = ta.TypeVar('ConcernT')
 ConfigT = ta.TypeVar('ConfigT')
 SiteConcernConfigT = ta.TypeVar('SiteConcernConfigT', bound='SiteConcernConfig')
@@ -84,7 +86,7 @@ class DeployConfig:
 # ../../../../omlish/lite/cached.py
 
 
-class cached_nullary:  # noqa
+class _cached_nullary:  # noqa
     def __init__(self, fn):
         super().__init__()
         self._fn = fn
@@ -101,6 +103,10 @@ class cached_nullary:  # noqa
         return bound
 
 
+def cached_nullary(fn: ta.Callable[..., T]) -> ta.Callable[..., T]:
+    return _cached_nullary(fn)
+
+
 ########################################
 # ../../../../omlish/lite/json.py
 
@@ -809,6 +815,24 @@ def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
     return out.decode().strip() if out is not None else None
 
 
+##
+
+
+def subprocess_close(
+        proc: subprocess.Popen,
+        timeout: ta.Optional[float] = None,
+) -> None:
+    # TODO: terminate, sleep, kill
+    if proc.stdout:
+        proc.stdout.close()
+    if proc.stderr:
+        proc.stderr.close()
+    if proc.stdin:
+        proc.stdin.close()
+
+    proc.wait(timeout)
+
+
 ########################################
 # ../runtime.py
 

ominfra/journald/genmessages.py

@@ -3,6 +3,7 @@
 import argparse
 import datetime
 import json
+import os
 import re
 import sys
 import time
@@ -34,6 +35,11 @@ def _main() -> None:
     else:
         start = 0
 
+    stdout_fd = sys.stdout.fileno()
+    out_fd = os.dup(stdout_fd)
+    null_fd = os.open('/dev/null', os.O_WRONLY)
+    os.dup2(null_fd, stdout_fd)
+
     for i in range(start, args.n):
         if args.sleep_s:
             if not args.sleep_n or (i and i % args.sleep_n == 0):
@@ -46,8 +52,14 @@ def _main() -> None:
             '__CURSOR': f'cursor:{i}',
             '_SOURCE_REALTIME_TIMESTAMP': str(int(ts_us)),
         }
-        print(json.dumps(dct, indent=None, separators=(',', ':')))
-        sys.stdout.flush()
+
+        buf = json.dumps(dct, indent=None, separators=(',', ':')).encode()
+
+        try:
+            os.write(out_fd, buf)
+            os.write(out_fd, b'\n')
+        except BrokenPipeError:
+            break
 
 
 if __name__ == '__main__':

ominfra/journald/tailer.py

@@ -347,9 +347,10 @@ import typing as ta
 from omlish.lite.cached import cached_nullary
 from omlish.lite.check import check_not_none
 from omlish.lite.logs import log
+from omlish.lite.subprocesses import subprocess_close
 from omlish.lite.subprocesses import subprocess_shell_wrap_exec
 
-from ..threadworker import ThreadWorker
+from ..threadworkers import ThreadWorker
 from .messages import JournalctlMessage  # noqa
 from .messages import JournalctlMessageBuilder
 
@@ -385,7 +386,7 @@ class JournalctlTailerWorker(ThreadWorker):
         self._read_size = read_size
         self._sleep_s = sleep_s
 
-        self._mb = JournalctlMessageBuilder()
+        self._builder = JournalctlMessageBuilder()
 
         self._proc: ta.Optional[subprocess.Popen] = None
 
@@ -409,45 +410,50 @@ class JournalctlTailerWorker(ThreadWorker):
 
         return cmd
 
+    def _read_loop(self, stdout: ta.IO) -> None:
+        while stdout.readable():
+            self._heartbeat()
+
+            buf = stdout.read(self._read_size)
+            if not buf:
+                log.debug('Journalctl empty read')
+                break
+
+            log.debug('Journalctl read buffer: %r', buf)
+            msgs = self._builder.feed(buf)
+            if msgs:
+                while True:
+                    try:
+                        self._output.put(msgs, timeout=1.)
+                    except queue.Full:
+                        self._heartbeat()
+                    else:
+                        break
+
     def _run(self) -> None:
         with subprocess.Popen(
             self._full_cmd(),
             stdout=subprocess.PIPE,
         ) as self._proc:
-            stdout = check_not_none(self._proc.stdout)
+            try:
+                stdout = check_not_none(self._proc.stdout)
+
+                fd = stdout.fileno()
+                fl = fcntl.fcntl(fd, fcntl.F_GETFL)
+                fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
 
-            fd = stdout.fileno()
-            fl = fcntl.fcntl(fd, fcntl.F_GETFL)
-            fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
+                while True:
+                    self._heartbeat()
 
-            while True:
-                if not self._heartbeat():
-                    return
+                    self._read_loop(stdout)
 
-                while stdout.readable():
-                    if not self._heartbeat():
+                    log.debug('Journalctl not readable')
+
+                    if self._proc.poll() is not None:
+                        log.critical('Journalctl process terminated')
                         return
 
-                    buf = stdout.read(self._read_size)
-                    if not buf:
-                        log.debug('Journalctl empty read')
-                        break
+                    time.sleep(self._sleep_s)
 
-                    log.debug('Journalctl read buffer: %r', buf)
-                    msgs = self._mb.feed(buf)
-                    if msgs:
-                        while True:
-                            try:
-                                self._output.put(msgs, timeout=1.)
-                            except queue.Full:
-                                if not self._heartbeat():
-                                    return
-                            else:
-                                break
-
-                if self._proc.poll() is not None:
-                    log.critical('Journalctl process terminated')
-                    return
-
-                log.debug('Journalctl readable')
-                time.sleep(self._sleep_s)
+            finally:
+                subprocess_close(self._proc)