ominfra 0.0.0.dev154__py3-none-any.whl → 0.0.0.dev156__py3-none-any.whl

ominfra/manage/targets/connection.py

@@ -0,0 +1,150 @@
+# ruff: noqa: UP006 UP007
+import abc
+import contextlib
+import dataclasses as dc
+import typing as ta
+
+from omlish.lite.check import check
+
+from ..bootstrap import MainBootstrap
+from ..commands.base import CommandExecutor
+from ..commands.local import LocalCommandExecutor
+from ..remote.connection import InProcessRemoteExecutionConnector
+from ..remote.connection import PyremoteRemoteExecutionConnector
+from ..remote.spawning import RemoteSpawning
+from .targets import DockerManageTarget
+from .targets import InProcessManageTarget
+from .targets import LocalManageTarget
+from .targets import ManageTarget
+from .targets import SshManageTarget
+from .targets import SubprocessManageTarget
+
+
+##
+
+
+class ManageTargetConnector(abc.ABC):
+    @abc.abstractmethod
+    def connect(self, tgt: ManageTarget) -> ta.AsyncContextManager[CommandExecutor]:
+        raise NotImplementedError
+
+
+##
+
+
+ManageTargetConnectorMap = ta.NewType('ManageTargetConnectorMap', ta.Mapping[ta.Type[ManageTarget], ManageTargetConnector])  # noqa
+
+
+@dc.dataclass(frozen=True)
+class TypeSwitchedManageTargetConnector(ManageTargetConnector):
+    connectors: ManageTargetConnectorMap
+
+    def get_connector(self, ty: ta.Type[ManageTarget]) -> ManageTargetConnector:
+        for k, v in self.connectors.items():
+            if issubclass(ty, k):
+                return v
+        raise KeyError(ty)
+
+    def connect(self, tgt: ManageTarget) -> ta.AsyncContextManager[CommandExecutor]:
+        return self.get_connector(type(tgt)).connect(tgt)
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class LocalManageTargetConnector(ManageTargetConnector):
+    _local_executor: LocalCommandExecutor
+    _in_process_connector: InProcessRemoteExecutionConnector
+    _pyremote_connector: PyremoteRemoteExecutionConnector
+    _bootstrap: MainBootstrap
+
+    @contextlib.asynccontextmanager
+    async def connect(self, tgt: ManageTarget) -> ta.AsyncGenerator[CommandExecutor, None]:
+        lmt = check.isinstance(tgt, LocalManageTarget)
+
+        if isinstance(lmt, InProcessManageTarget):
+            imt = check.isinstance(lmt, InProcessManageTarget)
+
+            if imt.mode == InProcessManageTarget.Mode.DIRECT:
+                yield self._local_executor
+
+            elif imt.mode == InProcessManageTarget.Mode.FAKE_REMOTE:
+                async with self._in_process_connector.connect() as rce:
+                    yield rce
+
+            else:
+                raise TypeError(imt.mode)
+
+        elif isinstance(lmt, SubprocessManageTarget):
+            async with self._pyremote_connector.connect(
+                    RemoteSpawning.Target(
+                        python=lmt.python,
+                    ),
+                    self._bootstrap,
+            ) as rce:
+                yield rce
+
+        else:
+            raise TypeError(lmt)
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class DockerManageTargetConnector(ManageTargetConnector):
+    _pyremote_connector: PyremoteRemoteExecutionConnector
+    _bootstrap: MainBootstrap
+
+    @contextlib.asynccontextmanager
+    async def connect(self, tgt: ManageTarget) -> ta.AsyncGenerator[CommandExecutor, None]:
+        dmt = check.isinstance(tgt, DockerManageTarget)
+
+        sh_parts: ta.List[str] = ['docker']
+        if dmt.image is not None:
+            sh_parts.extend(['run', '-i', dmt.image])
+        elif dmt.container_id is not None:
+            sh_parts.extend(['exec', dmt.container_id])
+        else:
+            raise ValueError(dmt)
+
+        async with self._pyremote_connector.connect(
+                RemoteSpawning.Target(
+                    shell=' '.join(sh_parts),
+                    python=dmt.python,
+                ),
+                self._bootstrap,
+        ) as rce:
+            yield rce
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class SshManageTargetConnector(ManageTargetConnector):
+    _pyremote_connector: PyremoteRemoteExecutionConnector
+    _bootstrap: MainBootstrap
+
+    @contextlib.asynccontextmanager
+    async def connect(self, tgt: ManageTarget) -> ta.AsyncGenerator[CommandExecutor, None]:
+        smt = check.isinstance(tgt, SshManageTarget)
+
+        sh_parts: ta.List[str] = ['ssh']
+        if smt.key_file is not None:
+            sh_parts.extend(['-i', smt.key_file])
+        addr = check.not_none(smt.host)
+        if smt.username is not None:
+            addr = f'{smt.username}@{addr}'
+        sh_parts.append(addr)
+
+        async with self._pyremote_connector.connect(
+                RemoteSpawning.Target(
+                    shell=' '.join(sh_parts),
+                    shell_quote=True,
+                    python=smt.python,
+                ),
+                self._bootstrap,
+        ) as rce:
+            yield rce

ominfra/manage/targets/inject.py

@@ -0,0 +1,42 @@
+# ruff: noqa: UP006 UP007
+import typing as ta
+
+from omlish.lite.inject import Injector
+from omlish.lite.inject import InjectorBindingOrBindings
+from omlish.lite.inject import InjectorBindings
+from omlish.lite.inject import inj
+
+from .connection import DockerManageTargetConnector
+from .connection import LocalManageTargetConnector
+from .connection import ManageTargetConnector
+from .connection import ManageTargetConnectorMap
+from .connection import SshManageTargetConnector
+from .connection import TypeSwitchedManageTargetConnector
+from .targets import DockerManageTarget
+from .targets import LocalManageTarget
+from .targets import SshManageTarget
+
+
+def bind_targets() -> InjectorBindings:
+    lst: ta.List[InjectorBindingOrBindings] = [
+        inj.bind(LocalManageTargetConnector, singleton=True),
+        inj.bind(DockerManageTargetConnector, singleton=True),
+        inj.bind(SshManageTargetConnector, singleton=True),
+
+        inj.bind(TypeSwitchedManageTargetConnector, singleton=True),
+        inj.bind(ManageTargetConnector, to_key=TypeSwitchedManageTargetConnector),
+    ]
+
+    #
+
+    def provide_manage_target_connector_map(injector: Injector) -> ManageTargetConnectorMap:
+        return ManageTargetConnectorMap({
+            LocalManageTarget: injector[LocalManageTargetConnector],
+            DockerManageTarget: injector[DockerManageTargetConnector],
+            SshManageTarget: injector[SshManageTargetConnector],
+        })
+    lst.append(inj.bind(provide_manage_target_connector_map, singleton=True))
+
+    #
+
+    return inj.as_bindings(*lst)

ominfra/manage/targets/targets.py

@@ -0,0 +1,87 @@
+# ruff: noqa: UP006 UP007
+"""
+It's desugaring. Subprocess and locals are only leafs. Retain an origin?
+** TWO LAYERS ** - ManageTarget is user-facing, ConnectorTarget is bound, internal
+"""
+import abc
+import dataclasses as dc
+import enum
+import typing as ta
+
+from omlish.lite.check import check
+
+
+##
+
+
+class ManageTarget(abc.ABC):  # noqa
+    def __init_subclass__(cls, **kwargs: ta.Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        check.state(cls.__name__.endswith('ManageTarget'))
+
+
+#
+
+
+@dc.dataclass(frozen=True)
+class PythonRemoteManageTarget:
+    DEFAULT_PYTHON: ta.ClassVar[str] = 'python3'
+    python: str = DEFAULT_PYTHON
+
+
+#
+
+
+class RemoteManageTarget(ManageTarget, abc.ABC):
+    pass
+
+
+class PhysicallyRemoteManageTarget(RemoteManageTarget, abc.ABC):
+    pass
+
+
+class LocalManageTarget(ManageTarget, abc.ABC):
+    pass
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class SshManageTarget(PhysicallyRemoteManageTarget, PythonRemoteManageTarget):
+    host: ta.Optional[str] = None
+    username: ta.Optional[str] = None
+    key_file: ta.Optional[str] = None
+
+    def __post_init__(self) -> None:
+        check.non_empty_str(self.host)
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class DockerManageTarget(RemoteManageTarget, PythonRemoteManageTarget):  # noqa
+    image: ta.Optional[str] = None
+    container_id: ta.Optional[str] = None
+
+    def __post_init__(self) -> None:
+        check.arg(bool(self.image) ^ bool(self.container_id))
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class InProcessManageTarget(LocalManageTarget):
+    class Mode(enum.Enum):
+        DIRECT = enum.auto()
+        FAKE_REMOTE = enum.auto()
+
+    mode: Mode = Mode.DIRECT
+
+
+@dc.dataclass(frozen=True)
+class SubprocessManageTarget(LocalManageTarget, PythonRemoteManageTarget):
+    pass

ominfra/scripts/journald2aws.py

@@ -961,8 +961,6 @@ def async_cached_nullary(fn):  # ta.Callable[..., T]) -> ta.Callable[..., T]:
 """
 TODO:
  - def maybe(v: lang.Maybe[T])
- - patch / override lite.check ?
-  - checker interface?
 """
 
 
@@ -2270,6 +2268,20 @@ def attr_setting(obj, attr, val, *, default=None):  # noqa
             setattr(obj, attr, orig)
 
 
+##
+
+
+class aclosing(contextlib.AbstractAsyncContextManager):  # noqa
+    def __init__(self, thing):
+        self.thing = thing
+
+    async def __aenter__(self):
+        return self.thing
+
+    async def __aexit__(self, *exc_info):
+        await self.thing.aclose()
+
+
 ########################################
 # ../../../../../omlish/lite/logs.py
 """
@@ -2545,7 +2557,8 @@ def configure_standard_logging(
 """
 TODO:
  - pickle stdlib objs? have to pin to 3.8 pickle protocol, will be cross-version
- - nonstrict toggle
+ - namedtuple
+ - literals
 """
 
 
@@ -2835,14 +2848,18 @@ class ObjMarshalerManager:
     ) -> ObjMarshaler:
         if isinstance(ty, type):
             if abc.ABC in ty.__bases__:
-                return PolymorphicObjMarshaler.of([  # type: ignore
+                impls = [ity for ity in deep_subclasses(ty) if abc.ABC not in ity.__bases__]  # type: ignore
+                if all(ity.__qualname__.endswith(ty.__name__) for ity in impls):
+                    ins = {ity: snake_case(ity.__qualname__[:-len(ty.__name__)]) for ity in impls}
+                else:
+                    ins = {ity: ity.__qualname__ for ity in impls}
+                return PolymorphicObjMarshaler.of([
                     PolymorphicObjMarshaler.Impl(
                         ity,
-                        ity.__qualname__,
+                        itn,
                         rec(ity),
                     )
-                    for ity in deep_subclasses(ty)
-                    if abc.ABC not in ity.__bases__
+                    for ity, itn in ins.items()
                 ])
 
             if issubclass(ty, enum.Enum):
@@ -3576,168 +3593,222 @@ SUBPROCESS_CHANNEL_OPTION_VALUES: ta.Mapping[SubprocessChannelOption, int] = {
 _SUBPROCESS_SHELL_WRAP_EXECS = False
 
 
-def subprocess_shell_wrap_exec(*args: str) -> ta.Tuple[str, ...]:
-    return ('sh', '-c', ' '.join(map(shlex.quote, args)))
+def subprocess_shell_wrap_exec(*cmd: str) -> ta.Tuple[str, ...]:
+    return ('sh', '-c', ' '.join(map(shlex.quote, cmd)))
 
 
-def subprocess_maybe_shell_wrap_exec(*args: str) -> ta.Tuple[str, ...]:
+def subprocess_maybe_shell_wrap_exec(*cmd: str) -> ta.Tuple[str, ...]:
     if _SUBPROCESS_SHELL_WRAP_EXECS or is_debugger_attached():
-        return subprocess_shell_wrap_exec(*args)
+        return subprocess_shell_wrap_exec(*cmd)
     else:
-        return args
-
-
-def prepare_subprocess_invocation(
-        *args: str,
-        env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
-        extra_env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
-        quiet: bool = False,
-        shell: bool = False,
-        **kwargs: ta.Any,
-) -> ta.Tuple[ta.Tuple[ta.Any, ...], ta.Dict[str, ta.Any]]:
-    log.debug('prepare_subprocess_invocation: args=%r', args)
-    if extra_env:
-        log.debug('prepare_subprocess_invocation: extra_env=%r', extra_env)
-
-    if extra_env:
-        env = {**(env if env is not None else os.environ), **extra_env}
-
-    if quiet and 'stderr' not in kwargs:
-        if not log.isEnabledFor(logging.DEBUG):
-            kwargs['stderr'] = subprocess.DEVNULL
-
-    if not shell:
-        args = subprocess_maybe_shell_wrap_exec(*args)
-
-    return args, dict(
-        env=env,
-        shell=shell,
-        **kwargs,
-    )
+        return cmd
 
 
 ##
 
 
-@contextlib.contextmanager
-def subprocess_common_context(*args: ta.Any, **kwargs: ta.Any) -> ta.Iterator[None]:
-    start_time = time.time()
-    try:
-        log.debug('subprocess_common_context.try: args=%r', args)
-        yield
-
-    except Exception as exc:  # noqa
-        log.debug('subprocess_common_context.except: exc=%r', exc)
-        raise
+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()
 
-    finally:
-        end_time = time.time()
-        elapsed_s = end_time - start_time
-        log.debug('subprocess_common_context.finally: elapsed_s=%f args=%r', elapsed_s, args)
+    proc.wait(timeout)
 
 
 ##
 
 
-def subprocess_check_call(
-        *args: str,
-        stdout: ta.Any = sys.stderr,
-        **kwargs: ta.Any,
-) -> None:
-    args, kwargs = prepare_subprocess_invocation(*args, stdout=stdout, **kwargs)
-    with subprocess_common_context(*args, **kwargs):
-        return subprocess.check_call(args, **kwargs)  # type: ignore
+class AbstractSubprocesses(abc.ABC):  # noqa
+    DEFAULT_LOGGER: ta.ClassVar[ta.Optional[logging.Logger]] = log
 
+    def __init__(
+            self,
+            *,
+            log: ta.Optional[logging.Logger] = None,
+            try_exceptions: ta.Optional[ta.Tuple[ta.Type[Exception], ...]] = None,
+    ) -> None:
+        super().__init__()
 
-def subprocess_check_output(
-        *args: str,
-        **kwargs: ta.Any,
-) -> bytes:
-    args, kwargs = prepare_subprocess_invocation(*args, **kwargs)
-    with subprocess_common_context(*args, **kwargs):
-        return subprocess.check_output(args, **kwargs)
+        self._log = log if log is not None else self.DEFAULT_LOGGER
+        self._try_exceptions = try_exceptions if try_exceptions is not None else self.DEFAULT_TRY_EXCEPTIONS
 
+    #
 
-def subprocess_check_output_str(*args: str, **kwargs: ta.Any) -> str:
-    return subprocess_check_output(*args, **kwargs).decode().strip()
+    def prepare_args(
+            self,
+            *cmd: str,
+            env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
+            extra_env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
+            quiet: bool = False,
+            shell: bool = False,
+            **kwargs: ta.Any,
+    ) -> ta.Tuple[ta.Tuple[ta.Any, ...], ta.Dict[str, ta.Any]]:
+        if self._log:
+            self._log.debug('Subprocesses.prepare_args: cmd=%r', cmd)
+            if extra_env:
+                self._log.debug('Subprocesses.prepare_args: extra_env=%r', extra_env)
 
+        if extra_env:
+            env = {**(env if env is not None else os.environ), **extra_env}
 
-##
+        if quiet and 'stderr' not in kwargs:
+            if self._log and not self._log.isEnabledFor(logging.DEBUG):
+                kwargs['stderr'] = subprocess.DEVNULL
 
+        if not shell:
+            cmd = subprocess_maybe_shell_wrap_exec(*cmd)
 
-DEFAULT_SUBPROCESS_TRY_EXCEPTIONS: ta.Tuple[ta.Type[Exception], ...] = (
-    FileNotFoundError,
-    subprocess.CalledProcessError,
-)
+        return cmd, dict(
+            env=env,
+            shell=shell,
+            **kwargs,
+        )
 
+    @contextlib.contextmanager
+    def wrap_call(self, *cmd: ta.Any, **kwargs: ta.Any) -> ta.Iterator[None]:
+        start_time = time.time()
+        try:
+            if self._log:
+                self._log.debug('Subprocesses.wrap_call.try: cmd=%r', cmd)
+            yield
 
-def _subprocess_try_run(
-        fn: ta.Callable[..., T],
-        *args: ta.Any,
-        try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
-        **kwargs: ta.Any,
-) -> ta.Union[T, Exception]:
-    try:
-        return fn(*args, **kwargs)
-    except try_exceptions as e:  # noqa
-        if log.isEnabledFor(logging.DEBUG):
-            log.exception('command failed')
-        return e
-
-
-def subprocess_try_call(
-        *args: str,
-        try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
-        **kwargs: ta.Any,
-) -> bool:
-    if isinstance(_subprocess_try_run(
-            subprocess_check_call,
-            *args,
-            try_exceptions=try_exceptions,
-            **kwargs,
-    ), Exception):
-        return False
-    else:
-        return True
+        except Exception as exc:  # noqa
+            if self._log:
+                self._log.debug('Subprocesses.wrap_call.except: exc=%r', exc)
+            raise
 
+        finally:
+            end_time = time.time()
+            elapsed_s = end_time - start_time
+            if self._log:
+                self._log.debug('sSubprocesses.wrap_call.finally: elapsed_s=%f cmd=%r', elapsed_s, cmd)
 
-def subprocess_try_output(
-        *args: str,
-        try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
-        **kwargs: ta.Any,
-) -> ta.Optional[bytes]:
-    if isinstance(ret := _subprocess_try_run(
-            subprocess_check_output,
-            *args,
-            try_exceptions=try_exceptions,
-            **kwargs,
-    ), Exception):
-        return None
-    else:
-        return ret
+    @contextlib.contextmanager
+    def prepare_and_wrap(
+            self,
+            *cmd: ta.Any,
+            **kwargs: ta.Any,
+    ) -> ta.Iterator[ta.Tuple[
+        ta.Tuple[ta.Any, ...],
+        ta.Dict[str, ta.Any],
+    ]]:
+        cmd, kwargs = self.prepare_args(*cmd, **kwargs)
+        with self.wrap_call(*cmd, **kwargs):
+            yield cmd, kwargs
 
+    #
 
-def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
-    out = subprocess_try_output(*args, **kwargs)
-    return out.decode().strip() if out is not None else None
+    DEFAULT_TRY_EXCEPTIONS: ta.Tuple[ta.Type[Exception], ...] = (
+        FileNotFoundError,
+        subprocess.CalledProcessError,
+    )
+
+    def try_fn(
+            self,
+            fn: ta.Callable[..., T],
+            *cmd: str,
+            try_exceptions: ta.Optional[ta.Tuple[ta.Type[Exception], ...]] = None,
+            **kwargs: ta.Any,
+    ) -> ta.Union[T, Exception]:
+        if try_exceptions is None:
+            try_exceptions = self._try_exceptions
+
+        try:
+            return fn(*cmd, **kwargs)
+
+        except try_exceptions as e:  # noqa
+            if self._log and self._log.isEnabledFor(logging.DEBUG):
+                self._log.exception('command failed')
+            return e
+
+    async def async_try_fn(
+            self,
+            fn: ta.Callable[..., ta.Awaitable[T]],
+            *cmd: ta.Any,
+            try_exceptions: ta.Optional[ta.Tuple[ta.Type[Exception], ...]] = None,
+            **kwargs: ta.Any,
+    ) -> ta.Union[T, Exception]:
+        if try_exceptions is None:
+            try_exceptions = self._try_exceptions
+
+        try:
+            return await fn(*cmd, **kwargs)
+
+        except try_exceptions as e:  # noqa
+            if self._log and self._log.isEnabledFor(logging.DEBUG):
+                self._log.exception('command failed')
+            return e
 
 
 ##
 
 
-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()
+class Subprocesses(AbstractSubprocesses):
+    def check_call(
+            self,
+            *cmd: str,
+            stdout: ta.Any = sys.stderr,
+            **kwargs: ta.Any,
+    ) -> None:
+        with self.prepare_and_wrap(*cmd, stdout=stdout, **kwargs) as (cmd, kwargs):  # noqa
+            subprocess.check_call(cmd, **kwargs)
 
-    proc.wait(timeout)
+    def check_output(
+            self,
+            *cmd: str,
+            **kwargs: ta.Any,
+    ) -> bytes:
+        with self.prepare_and_wrap(*cmd, **kwargs) as (cmd, kwargs):  # noqa
+            return subprocess.check_output(cmd, **kwargs)
+
+    def check_output_str(
+            self,
+            *cmd: str,
+            **kwargs: ta.Any,
+    ) -> str:
+        return self.check_output(*cmd, **kwargs).decode().strip()
+
+    #
+
+    def try_call(
+            self,
+            *cmd: str,
+            **kwargs: ta.Any,
+    ) -> bool:
+        if isinstance(self.try_fn(self.check_call, *cmd, **kwargs), Exception):
+            return False
+        else:
+            return True
+
+    def try_output(
+            self,
+            *cmd: str,
+            **kwargs: ta.Any,
+    ) -> ta.Optional[bytes]:
+        if isinstance(ret := self.try_fn(self.check_output, *cmd, **kwargs), Exception):
+            return None
+        else:
+            return ret
+
+    def try_output_str(
+            self,
+            *cmd: str,
+            **kwargs: ta.Any,
+    ) -> ta.Optional[str]:
+        if (ret := self.try_output(*cmd, **kwargs)) is None:
+            return None
+        else:
+            return ret.decode().strip()
+
+
+subprocesses = Subprocesses()
 
 
 ########################################