karton-core 5.7.0__py3-none-any.whl → 5.9.0__py3-none-any.whl

karton/core/backend.py

@@ -17,10 +17,12 @@ from botocore.credentials import (
 from botocore.session import get_session
 from redis import AuthenticationError, StrictRedis
 from redis.client import Pipeline
+from redis.connection import parse_url as parse_redis_url
 from urllib3.response import HTTPResponse
 
 from .config import Config
 from .exceptions import InvalidIdentityError
+from .resource import LocalResource, RemoteResource
 from .task import Task, TaskPriority, TaskState
 from .utils import chunks, chunks_iter
 
@@ -33,12 +35,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 +113,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 +127,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,54 +136,47 @@ class KartonBackend:
                 f"Karton identity should not contain {disallowed_chars}"
             )
 
+    @property
+    def default_bucket_name(self) -> str:
+        bucket_name = self.config.get("s3", "bucket")
+        if not bucket_name:
+            raise RuntimeError("S3 default bucket is not defined in configuration")
+        return bucket_name
+
     @staticmethod
-    def make_redis(
-        config,
+    def get_redis_configuration(
+        config: 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
-        """
+    ) -> Dict[str, Any]:
         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")
-        if not bucket_name:
-            raise RuntimeError("S3 default bucket is not defined in configuration")
-        return bucket_name
+        redis_url = config.get("redis", "url")
+        if redis_url is not None:
+            redis_conf = parse_redis_url(redis_url)
+        else:
+            redis_conf = {
+                "host": config["redis"]["host"],
+                "port": config.getint("redis", "port", 6379),
+                "db": config.getint("redis", "db", 0),
+                "ssl": config.getboolean("redis", "ssl", False),
+            }
+
+        if username := config.get("redis", "username"):
+            redis_conf["username"] = username
+            password = config.get("redis", "password")
+            if password is None:
+                raise RuntimeError("You must set both username and password, or none")
+            redis_conf["password"] = password
+        # Don't set if set to 0
+        if socket_timeout := config.get("redis", "socket_timeout", 30):
+            redis_conf["socket_timeout"] = socket_timeout
+        redis_conf["client_name"] = client_name
+        redis_conf["decode_responses"] = True
+        return redis_conf
 
     @staticmethod
     def get_queue_name(identity: str, priority: TaskPriority) -> str:
@@ -270,6 +220,7 @@ class KartonBackend:
                 "filters": bind.filters,
                 "persistent": bind.persistent,
                 "service_version": bind.service_version,
+                "is_async": bind.is_async,
             },
             sort_keys=True,
         )
@@ -294,6 +245,7 @@ class KartonBackend:
                 persistent=not identity.endswith(".test"),
                 filters=bind,
                 service_version=None,
+                is_async=False,
             )
         return KartonBind(
             identity=identity,
@@ -302,6 +254,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 +269,115 @@ 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,
+                )
+
+    @classmethod
+    def make_redis(
+        cls,
+        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
+        """
+        redis_args = cls.get_redis_configuration(
+            config, identity=identity, service_info=service_info
+        )
+        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["username"]
+            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 +488,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 +513,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 +532,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 +620,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
@@ -584,10 +656,29 @@ class KartonBackend:
             task_keys, chunk_size=chunk_size, parse_resources=parse_resources
         )
 
+    def declare_task(self, task: Task) -> None:
+        """
+        Declares a new task to send it to the queue.
+
+        Task producers should use this method for new tasks.
+
+        :param task: Task to declare
+        """
+        # Ensure all local resources have good buckets
+        for resource in task.iterate_resources():
+            if isinstance(resource, LocalResource) and not resource.bucket:
+                resource.bucket = self.default_bucket_name
+
+        # Register new task
+        self.register_task(task)
+
     def register_task(self, task: Task, pipe: Optional[Pipeline] = None) -> None:
         """
         Register or update task in Redis.
 
+        This method is used internally to alter task data. If you want to declare new
+        task in Redis, use declare_task.
+
         :param task: Task object
         :param pipe: Optional pipeline object if operation is a part of pipeline
         """
@@ -798,12 +889,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],