karton-core 5.6.1__py3-none-any.whl → 5.8.0__py3-none-any.whl

karton/core/backend.py

@@ -21,6 +21,7 @@ from urllib3.response import HTTPResponse
 
 from .config import Config
 from .exceptions import InvalidIdentityError
+from .resource import RemoteResource
 from .task import Task, TaskPriority, TaskState
 from .utils import chunks, chunks_iter
 
@@ -33,12 +34,20 @@ KARTON_OUTPUTS_NAMESPACE = "karton.outputs"
 
 KartonBind = namedtuple(
     "KartonBind",
-    ["identity", "info", "version", "persistent", "filters", "service_version"],
+    [
+        "identity",
+        "info",
+        "version",
+        "persistent",
+        "filters",
+        "service_version",
+        "is_async",
+    ],
 )
 
 
 KartonOutputs = namedtuple("KartonOutputs", ["identity", "outputs"])
-logger = logging.getLogger("karton.core.backend")
+logger = logging.getLogger(__name__)
 
 
 class KartonMetrics(enum.Enum):
@@ -103,13 +112,13 @@ class KartonServiceInfo:
         )
 
 
-class KartonBackend:
+class KartonBackendBase:
     def __init__(
         self,
         config: Config,
         identity: Optional[str] = None,
         service_info: Optional[KartonServiceInfo] = None,
-    ) -> None:
+    ):
         self.config = config
 
         if identity is not None:
@@ -117,59 +126,6 @@ class KartonBackend:
         self.identity = identity
 
         self.service_info = service_info
-        self.redis = self.make_redis(
-            config, identity=identity, service_info=service_info
-        )
-
-        endpoint = config.get("s3", "address")
-        access_key = config.get("s3", "access_key")
-        secret_key = config.get("s3", "secret_key")
-        iam_auth = config.getboolean("s3", "iam_auth")
-
-        if not endpoint:
-            raise RuntimeError("Attempting to get S3 client without an endpoint set")
-
-        if access_key and secret_key and iam_auth:
-            logger.warning(
-                "Warning: iam is turned on and both S3 access key and secret key are"
-                " provided"
-            )
-
-        if iam_auth:
-            s3_client = self.iam_auth_s3(endpoint)
-            if s3_client:
-                self.s3 = s3_client
-                return
-
-        if access_key is None or secret_key is None:
-            raise RuntimeError(
-                "Attempting to get S3 client without an access_key/secret_key set"
-            )
-
-        self.s3 = boto3.client(
-            "s3",
-            endpoint_url=endpoint,
-            aws_access_key_id=access_key,
-            aws_secret_access_key=secret_key,
-        )
-
-    def iam_auth_s3(self, endpoint: str):
-        boto_session = get_session()
-        iam_providers = [
-            ContainerProvider(),
-            InstanceMetadataProvider(
-                iam_role_fetcher=InstanceMetadataFetcher(timeout=1000, num_attempts=2)
-            ),
-        ]
-
-        for provider in iam_providers:
-            creds = provider.load()
-            if creds:
-                boto_session._credentials = creds  # type: ignore
-                return boto3.Session(botocore_session=boto_session).client(
-                    "s3",
-                    endpoint_url=endpoint,
-                )
 
     @staticmethod
     def _validate_identity(identity: str):
@@ -179,48 +135,6 @@ class KartonBackend:
                 f"Karton identity should not contain {disallowed_chars}"
             )
 
-    @staticmethod
-    def make_redis(
-        config,
-        identity: Optional[str] = None,
-        service_info: Optional[KartonServiceInfo] = None,
-    ) -> StrictRedis:
-        """
-        Create and test a Redis connection.
-
-        :param config: The karton configuration
-        :param identity: Karton service identity
-        :param service_info: Additional service identity metadata
-        :return: Redis connection
-        """
-        if service_info is not None:
-            client_name: Optional[str] = service_info.make_client_name()
-        else:
-            client_name = identity
-
-        redis_args = {
-            "host": config["redis"]["host"],
-            "port": config.getint("redis", "port", 6379),
-            "db": config.getint("redis", "db", 0),
-            "username": config.get("redis", "username"),
-            "password": config.get("redis", "password"),
-            "client_name": client_name,
-            # set socket_timeout to None if set to 0
-            "socket_timeout": config.getint("redis", "socket_timeout", 30) or None,
-            "decode_responses": True,
-        }
-        try:
-            redis = StrictRedis(**redis_args)
-            redis.ping()
-        except AuthenticationError:
-            # Maybe we've sent a wrong password.
-            # Or maybe the server is not (yet) password protected
-            # To make smooth transition possible, try to login insecurely
-            del redis_args["password"]
-            redis = StrictRedis(**redis_args)
-            redis.ping()
-        return redis
-
     @property
     def default_bucket_name(self) -> str:
         bucket_name = self.config.get("s3", "bucket")
@@ -270,6 +184,7 @@ class KartonBackend:
                 "filters": bind.filters,
                 "persistent": bind.persistent,
                 "service_version": bind.service_version,
+                "is_async": bind.is_async,
             },
             sort_keys=True,
         )
@@ -294,6 +209,7 @@ class KartonBackend:
                 persistent=not identity.endswith(".test"),
                 filters=bind,
                 service_version=None,
+                is_async=False,
             )
         return KartonBind(
             identity=identity,
@@ -302,6 +218,7 @@ class KartonBackend:
             persistent=bind["persistent"],
             filters=bind["filters"],
             service_version=bind.get("service_version"),
+            is_async=bind.get("is_async", False),
         )
 
     @staticmethod
@@ -316,6 +233,126 @@ class KartonBackend:
         output = [json.loads(output_type) for output_type in output_data]
         return KartonOutputs(identity=identity, outputs=output)
 
+    @staticmethod
+    def _log_channel(logger_name: Optional[str], level: Optional[str]) -> str:
+        return ".".join(
+            [KARTON_LOG_CHANNEL, (level or "*").lower(), logger_name or "*"]
+        )
+
+
+class KartonBackend(KartonBackendBase):
+    def __init__(
+        self,
+        config: Config,
+        identity: Optional[str] = None,
+        service_info: Optional[KartonServiceInfo] = None,
+    ) -> None:
+        super().__init__(config, identity, service_info)
+        self.redis = self.make_redis(
+            config, identity=identity, service_info=service_info
+        )
+
+        endpoint = config.get("s3", "address")
+        access_key = config.get("s3", "access_key")
+        secret_key = config.get("s3", "secret_key")
+        iam_auth = config.getboolean("s3", "iam_auth")
+
+        if not endpoint:
+            raise RuntimeError("Attempting to get S3 client without an endpoint set")
+
+        if access_key and secret_key and iam_auth:
+            logger.warning(
+                "Warning: iam is turned on and both S3 access key and secret key are"
+                " provided"
+            )
+
+        if iam_auth:
+            s3_client = self.iam_auth_s3(endpoint)
+            if s3_client:
+                self.s3 = s3_client
+                return
+
+        if access_key is None or secret_key is None:
+            raise RuntimeError(
+                "Attempting to get S3 client without an access_key/secret_key set"
+            )
+
+        self.s3 = boto3.client(
+            "s3",
+            endpoint_url=endpoint,
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+        )
+
+    def iam_auth_s3(self, endpoint: str):
+        boto_session = get_session()
+        iam_providers = [
+            ContainerProvider(),
+            InstanceMetadataProvider(
+                iam_role_fetcher=InstanceMetadataFetcher(timeout=1000, num_attempts=2)
+            ),
+        ]
+
+        for provider in iam_providers:
+            creds = provider.load()
+            if creds:
+                boto_session._credentials = creds  # type: ignore
+                return boto3.Session(botocore_session=boto_session).client(
+                    "s3",
+                    endpoint_url=endpoint,
+                )
+
+    @staticmethod
+    def make_redis(
+        config,
+        identity: Optional[str] = None,
+        service_info: Optional[KartonServiceInfo] = None,
+    ) -> StrictRedis:
+        """
+        Create and test a Redis connection.
+
+        :param config: The karton configuration
+        :param identity: Karton service identity
+        :param service_info: Additional service identity metadata
+        :return: Redis connection
+        """
+        if service_info is not None:
+            client_name: Optional[str] = service_info.make_client_name()
+        else:
+            client_name = identity
+
+        redis_args = {
+            "host": config["redis"]["host"],
+            "port": config.getint("redis", "port", 6379),
+            "db": config.getint("redis", "db", 0),
+            "username": config.get("redis", "username"),
+            "password": config.get("redis", "password"),
+            "client_name": client_name,
+            # set socket_timeout to None if set to 0
+            "socket_timeout": config.getint("redis", "socket_timeout", 30) or None,
+            "decode_responses": True,
+        }
+        try:
+            redis = StrictRedis(**redis_args)
+            redis.ping()
+        except AuthenticationError:
+            # Maybe we've sent a wrong password.
+            # Or maybe the server is not (yet) password protected
+            # To make smooth transition possible, try to login insecurely
+            del redis_args["password"]
+            redis = StrictRedis(**redis_args)
+            redis.ping()
+        return redis
+
+    def unserialize_resource(self, resource_spec: Dict[str, Any]) -> RemoteResource:
+        """
+        Unserializes resource into a RemoteResource object bound with current backend
+
+        :param resource_spec: Resource specification
+        :return: RemoteResource object
+        """
+        return RemoteResource.from_dict(resource_spec, backend=self)
+
     def get_bind(self, identity: str) -> KartonBind:
         """
         Get bind object for given identity
@@ -426,7 +463,9 @@ class KartonBackend:
         task_data = self.redis.get(f"{KARTON_TASK_NAMESPACE}:{task_uid}")
         if not task_data:
             return None
-        return Task.unserialize(task_data, backend=self)
+        return Task.unserialize(
+            task_data, resource_unserializer=self.unserialize_resource
+        )
 
     def get_tasks(
         self,
@@ -449,7 +488,11 @@ class KartonBackend:
             chunk_size,
         )
         return [
-            Task.unserialize(task_data, backend=self, parse_resources=parse_resources)
+            Task.unserialize(
+                task_data,
+                parse_resources=parse_resources,
+                resource_unserializer=self.unserialize_resource,
+            )
             for chunk in keys
             for task_data in self.redis.mget(chunk)
             if task_data is not None
@@ -464,7 +507,9 @@ class KartonBackend:
         for chunk in chunks_iter(task_keys, chunk_size):
             yield from (
                 Task.unserialize(
-                    task_data, backend=self, parse_resources=parse_resources
+                    task_data,
+                    parse_resources=parse_resources,
+                    resource_unserializer=self.unserialize_resource,
                 )
                 for task_data in self.redis.mget(chunk)
                 if task_data is not None
@@ -550,7 +595,9 @@ class KartonBackend:
                 lambda task: task.root_uid == root_uid,
                 (
                     Task.unserialize(
-                        task_data, backend=self, parse_resources=parse_resources
+                        task_data,
+                        parse_resources=parse_resources,
+                        resource_unserializer=self.unserialize_resource,
                     )
                     for task_data in self.redis.mget(chunk)
                     if task_data is not None
@@ -798,12 +845,6 @@ class KartonBackend:
         p.execute()
         return new_task
 
-    @staticmethod
-    def _log_channel(logger_name: Optional[str], level: Optional[str]) -> str:
-        return ".".join(
-            [KARTON_LOG_CHANNEL, (level or "*").lower(), logger_name or "*"]
-        )
-
     def produce_log(
         self,
         log_record: Dict[str, Any],
@@ -1040,9 +1081,11 @@ class KartonBackend:
         """
         deletion_chunks = chunks(
             [
-                {"Key": uid, "VersionId": version_id}
-                if version_id != "null" or explicit_version_null
-                else {"Key": uid}
+                (
+                    {"Key": uid, "VersionId": version_id}
+                    if version_id != "null" or explicit_version_null
+                    else {"Key": uid}
+                )
                 for uid, versions in object_versions.items()
                 for version_id in versions
             ],
@@ -1070,17 +1113,20 @@ class KartonBackend:
                 raise e
         return False
 
-    def log_identity_output(self, identity: str, headers: Dict[str, Any]) -> None:
+    def log_identity_output(
+        self, identity: str, headers: Dict[str, Any], task_tracking_ttl: int
+    ) -> None:
         """
         Store the type of task outputted for given producer to
         be used in tracking karton service connections.
 
         :param identity: producer identity
         :param headers: outputted headers
+        :param task_tracking_ttl: expire time (in seconds)
         """
 
         self.redis.sadd(f"{KARTON_OUTPUTS_NAMESPACE}:{identity}", json.dumps(headers))
-        self.redis.expire(f"{KARTON_OUTPUTS_NAMESPACE}:{identity}", 60 * 60 * 24 * 30)
+        self.redis.expire(f"{KARTON_OUTPUTS_NAMESPACE}:{identity}", task_tracking_ttl)
 
     def get_outputs(self) -> List[KartonOutputs]:
         """

karton/core/base.py

@@ -9,33 +9,20 @@ from typing import Optional, Union, cast
 from .__version__ import __version__
 from .backend import KartonBackend, KartonServiceInfo
 from .config import Config
-from .logger import KartonLogHandler
-from .task import Task
+from .logger import KartonLogHandler, TaskContextFilter
+from .task import Task, get_current_task, set_current_task
 from .utils import HardShutdownInterrupt, StrictClassMethod, graceful_killer
 
 
-class KartonBase(abc.ABC):
-    """
-    Base class for all Karton services
-
-    You can set an informative version information by setting the ``version`` class
-    attribute.
-    """
-
-    #: Karton service identity
-    identity: str = ""
-    #: Karton service version
-    version: Optional[str] = None
-    #: Include extended service information for non-consumer services
-    with_service_info: bool = False
+class ConfigMixin:
+    identity: Optional[str]
+    version: Optional[str]
 
-    def __init__(
-        self,
-        config: Optional[Config] = None,
-        identity: Optional[str] = None,
-        backend: Optional[KartonBackend] = None,
-    ) -> None:
+    def __init__(self, config: Optional[Config] = None, identity: Optional[str] = None):
         self.config = config or Config()
+        self.enable_publish_log = self.config.getboolean(
+            "logging", "enable_publish", True
+        )
 
         # If not passed via constructor - get it from class
         if identity is not None:
@@ -47,25 +34,81 @@ class KartonBase(abc.ABC):
 
         self.debug = self.config.getboolean("karton", "debug", False)
 
-        if self.debug:
+        if self.debug and self.identity:
             self.identity += "-" + os.urandom(4).hex() + "-dev"
 
-        self.service_info = None
-        if self.identity is not None and self.with_service_info:
-            self.service_info = KartonServiceInfo(
-                identity=self.identity,
-                karton_version=__version__,
-                service_version=self.version,
-            )
+    @classmethod
+    def args_description(cls) -> str:
+        """Return short description for argument parser."""
+        if not cls.__doc__:
+            return ""
+        return textwrap.dedent(cls.__doc__).strip().splitlines()[0]
 
-        self.backend = backend or KartonBackend(
-            self.config, identity=self.identity, service_info=self.service_info
+    @classmethod
+    def args_parser(cls) -> argparse.ArgumentParser:
+        """
+        Return ArgumentParser for main() class method.
+
+        This method should be overridden and call super methods
+        if you want to add more arguments.
+        """
+        parser = argparse.ArgumentParser(description=cls.args_description())
+        parser.add_argument(
+            "--version", action="version", version=cast(str, cls.version)
+        )
+        parser.add_argument("--config-file", help="Alternative configuration path")
+        parser.add_argument(
+            "--identity", help="Alternative identity for Karton service"
         )
+        parser.add_argument("--log-level", help="Logging level of Karton logger")
+        parser.add_argument(
+            "--debug", help="Enable debugging mode", action="store_true", default=None
+        )
+        return parser
 
-        self._log_handler = KartonLogHandler(
-            backend=self.backend, channel=self.identity
+    @classmethod
+    def config_from_args(cls, config: Config, args: argparse.Namespace) -> None:
+        """
+        Updates configuration with settings from arguments
+
+        This method should be overridden and call super methods
+        if you want to add more arguments.
+        """
+        config.load_from_dict(
+            {
+                "karton": {
+                    "identity": args.identity,
+                    "debug": args.debug,
+                },
+                "logging": {"level": args.log_level},
+            }
         )
-        self.current_task: Optional[Task] = None
+
+    @classmethod
+    def karton_from_args(cls, args: Optional[argparse.Namespace] = None):
+        """
+        Returns Karton instance configured using configuration files
+        and provided arguments
+
+        Used by :py:meth:`KartonServiceBase.main` method
+        """
+        if args is None:
+            parser = cls.args_parser()
+            args = parser.parse_args()
+        config = Config(path=args.config_file)
+        cls.config_from_args(config, args)
+        return cls(config=config)
+
+
+class LoggingMixin:
+    config: Config
+    identity: Optional[str]
+    debug: bool
+    enable_publish_log: bool
+
+    def __init__(self, log_handler: logging.Handler, log_format: str) -> None:
+        self._log_handler = log_handler
+        self._log_format = log_format
 
     def setup_logger(self, level: Optional[Union[str, int]] = None) -> None:
         """
@@ -91,6 +134,7 @@ class KartonBase(abc.ABC):
         self._log_handler.setFormatter(logging.Formatter())
 
         logger = logging.getLogger(self.identity)
+        logger.addFilter(TaskContextFilter())
 
         if logger.handlers:
             # If logger already have handlers set: clear them
@@ -103,16 +147,14 @@ class KartonBase(abc.ABC):
 
         logger.setLevel(log_level)
         stream_handler = logging.StreamHandler()
-        stream_handler.setFormatter(
-            logging.Formatter("[%(asctime)s][%(levelname)s] %(message)s")
-        )
+        stream_handler.setFormatter(logging.Formatter(self._log_format))
         logger.addHandler(stream_handler)
 
-        if not self.debug:
+        if not self.debug and self.enable_publish_log:
             logger.addHandler(self._log_handler)
 
     @property
-    def log_handler(self) -> KartonLogHandler:
+    def log_handler(self) -> logging.Handler:
         """
         Return KartonLogHandler bound to this Karton service.
 
@@ -138,67 +180,54 @@ class KartonBase(abc.ABC):
         """
         return logging.getLogger(self.identity)
 
-    @classmethod
-    def args_description(cls) -> str:
-        """Return short description for argument parser."""
-        if not cls.__doc__:
-            return ""
-        return textwrap.dedent(cls.__doc__).strip().splitlines()[0]
 
-    @classmethod
-    def args_parser(cls) -> argparse.ArgumentParser:
-        """
-        Return ArgumentParser for main() class method.
+class KartonBase(abc.ABC, ConfigMixin, LoggingMixin):
+    """
+    Base class for all Karton services
 
-        This method should be overridden and call super methods
-        if you want to add more arguments.
-        """
-        parser = argparse.ArgumentParser(description=cls.args_description())
-        parser.add_argument(
-            "--version", action="version", version=cast(str, cls.version)
-        )
-        parser.add_argument("--config-file", help="Alternative configuration path")
-        parser.add_argument(
-            "--identity", help="Alternative identity for Karton service"
-        )
-        parser.add_argument("--log-level", help="Logging level of Karton logger")
-        parser.add_argument(
-            "--debug", help="Enable debugging mode", action="store_true", default=None
-        )
-        return parser
+    You can set an informative version information by setting the ``version`` class
+    attribute.
+    """
 
-    @classmethod
-    def config_from_args(cls, config: Config, args: argparse.Namespace) -> None:
-        """
-        Updates configuration with settings from arguments
+    #: Karton service identity
+    identity: str = ""
+    #: Karton service version
+    version: Optional[str] = None
+    #: Include extended service information for non-consumer services
+    with_service_info: bool = False
 
-        This method should be overridden and call super methods
-        if you want to add more arguments.
-        """
-        config.load_from_dict(
-            {
-                "karton": {
-                    "identity": args.identity,
-                    "debug": args.debug,
-                },
-                "logging": {"level": args.log_level},
-            }
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        identity: Optional[str] = None,
+        backend: Optional[KartonBackend] = None,
+    ) -> None:
+        ConfigMixin.__init__(self, config, identity)
+
+        self.service_info = None
+        if self.identity is not None and self.with_service_info:
+            self.service_info = KartonServiceInfo(
+                identity=self.identity,
+                karton_version=__version__,
+                service_version=self.version,
+            )
+
+        self.backend = backend or KartonBackend(
+            self.config, identity=self.identity, service_info=self.service_info
         )
 
-    @classmethod
-    def karton_from_args(cls, args: Optional[argparse.Namespace] = None):
-        """
-        Returns Karton instance configured using configuration files
-        and provided arguments
+        log_handler = KartonLogHandler(backend=self.backend, channel=self.identity)
+        LoggingMixin.__init__(
+            self, log_handler, log_format="[%(asctime)s][%(levelname)s] %(message)s"
+        )
 
-        Used by :py:meth:`KartonServiceBase.main` method
-        """
-        if args is None:
-            parser = cls.args_parser()
-            args = parser.parse_args()
-        config = Config(path=args.config_file)
-        cls.config_from_args(config, args)
-        return cls(config=config)
+    @property
+    def current_task(self) -> Optional[Task]:
+        return get_current_task()
+
+    @current_task.setter
+    def current_task(self, task: Optional[Task]):
+        set_current_task(task)
 
 
 class KartonServiceBase(KartonBase):