UncountablePythonSDK 0.0.52__py3-none-any.whl → 0.0.131__py3-none-any.whl

uncountable/integration/executors/generic_upload_executor.py

@@ -1,7 +1,8 @@
+import datetime
 import io
 import os
 import re
-from datetime import datetime, timezone
+from datetime import UTC
 
 import paramiko
 
@@ -10,12 +11,12 @@ from pkgs.filesystem_utils import (
     FileSystemFileReference,
     FileSystemObject,
     FileSystemS3Config,
+    FileSystemSession,
     FileSystemSFTPConfig,
     FileTransfer,
     S3Session,
     SFTPSession,
 )
-from pkgs.filesystem_utils.filesystem_session import FileSystemSession
 from uncountable.core.file_upload import DataFileUpload, FileUpload
 from uncountable.integration.job import Job, JobArguments
 from uncountable.integration.secret_retrieval import retrieve_secret
@@ -33,6 +34,27 @@ from uncountable.types.job_definition_t import (
 )
 
 
+def _get_extension(filename: str) -> str | None:
+    _, ext = os.path.splitext(filename)
+    return ext.strip().lower()
+
+
+def _run_keyword_detection(data: io.BytesIO, keyword: str) -> bool:
+    try:
+        text = io.TextIOWrapper(data, encoding="utf-8")
+        for line in text:
+            if (
+                keyword in line
+                or re.search(keyword, line, flags=re.IGNORECASE) is not None
+            ):
+                return True
+        return False
+    except re.error:
+        return False
+    except UnicodeError:
+        return False
+
+
 def _filter_files_by_keyword(
     remote_directory: GenericRemoteDirectoryScope,
     files: list[FileObjectData],
@@ -41,7 +63,20 @@ def _filter_files_by_keyword(
     if remote_directory.detection_keyword is None:
         return files
 
-    raise NotImplementedError("keyword detection not implemented yet")
+    filtered_files = []
+
+    for file in files:
+        extension = _get_extension(file.filename)
+
+        if extension not in (".txt", ".csv"):
+            raise NotImplementedError(
+                "keyword detection is only supported for csv, txt files"
+            )
+
+        if _run_keyword_detection(file.file_IO, remote_directory.detection_keyword):
+            filtered_files.append(file)
+
+    return filtered_files
 
 
 def _filter_by_filename(
@@ -68,7 +103,8 @@ def _filter_by_file_extension(
         file
         for file in files
         if file.filename is not None
-        and os.path.splitext(file.filename)[-1] in remote_directory.valid_file_extensions
+        and os.path.splitext(file.filename)[-1]
+        in remote_directory.valid_file_extensions
     ]
 
 
@@ -129,7 +165,7 @@ def _move_files_post_upload(
     appended_text = ""
 
     if remote_directory_scope.prepend_date_on_archive:
-        appended_text = f"-{datetime.now(timezone.utc).timestamp()}"
+        appended_text = f"-{datetime.datetime.now(UTC).timestamp()}"
 
     for file_path in success_file_paths:
         filename = os.path.split(file_path)[-1]
@@ -164,7 +200,7 @@ def _move_files_post_upload(
     filesystem_session.move_files([*success_file_transfers, *failed_file_transfers])
 
 
-class GenericUploadJob(Job):
+class GenericUploadJob(Job[None]):
     def __init__(
         self,
         data_source: GenericUploadDataSource,
@@ -176,19 +212,40 @@ class GenericUploadJob(Job):
         self.upload_strategy = upload_strategy
         self.data_source = data_source
 
+    @property
+    def payload_type(self) -> type[None]:
+        return type(None)
+
     def _construct_filesystem_session(self, args: JobArguments) -> FileSystemSession:
         match self.data_source:
             case GenericUploadDataSourceSFTP():
-                pem_secret = retrieve_secret(
-                    self.data_source.pem_secret, profile_metadata=args.profile_metadata
-                )
-                pem_key = paramiko.RSAKey.from_private_key(io.StringIO(pem_secret))
-                sftp_config = FileSystemSFTPConfig(
-                    ip=self.data_source.host,
-                    username=self.data_source.username,
-                    pem_path=None,
-                    pem_key=pem_key,
-                )
+                if self.data_source.pem_secret is not None:
+                    pem_secret = retrieve_secret(
+                        self.data_source.pem_secret,
+                        profile_metadata=args.profile_metadata,
+                    )
+                    pem_key = paramiko.RSAKey.from_private_key(io.StringIO(pem_secret))
+                    sftp_config = FileSystemSFTPConfig(
+                        ip=self.data_source.host,
+                        username=self.data_source.username,
+                        pem_path=None,
+                        pem_key=pem_key,
+                    )
+                elif self.data_source.password_secret is not None:
+                    password_secret = retrieve_secret(
+                        self.data_source.password_secret,
+                        profile_metadata=args.profile_metadata,
+                    )
+                    sftp_config = FileSystemSFTPConfig(
+                        ip=self.data_source.host,
+                        username=self.data_source.username,
+                        pem_path=None,
+                        password=password_secret,
+                    )
+                else:
+                    raise ValueError(
+                        "Either pem_secret or password_secret must be specified for sftp data source"
+                    )
                 return SFTPSession(sftp_config=sftp_config)
             case GenericUploadDataSourceS3():
                 if self.data_source.access_key_secret is not None:
@@ -200,19 +257,17 @@ class GenericUploadJob(Job):
                     secret_access_key = None
 
                 if self.data_source.endpoint_url is None:
-                    assert (
-                        self.data_source.cloud_provider is not None
-                    ), "either cloud_provider or endpoint_url must be specified"
+                    assert self.data_source.cloud_provider is not None, (
+                        "either cloud_provider or endpoint_url must be specified"
+                    )
                     match self.data_source.cloud_provider:
                         case S3CloudProvider.AWS:
                             endpoint_url = "https://s3.amazonaws.com"
                         case S3CloudProvider.OVH:
-                            assert (
-                                self.data_source.region_name is not None
-                            ), "region_name must be specified for cloud_provider OVH"
-                            endpoint_url = (
-                                f"https://s3.{self.data_source.region_name}.cloud.ovh.net"
+                            assert self.data_source.region_name is not None, (
+                                "region_name must be specified for cloud_provider OVH"
                             )
+                            endpoint_url = f"https://s3.{self.data_source.region_name}.cloud.ovh.net"
                 else:
                     endpoint_url = self.data_source.endpoint_url
 
@@ -227,7 +282,7 @@ class GenericUploadJob(Job):
 
                 return S3Session(s3_config=s3_config)
 
-    def run(self, args: JobArguments) -> JobResult:
+    def run_outer(self, args: JobArguments) -> JobResult:
         client = args.client
         batch_processor = args.batch_processor
         logger = args.logger
@@ -248,7 +303,8 @@ class GenericUploadJob(Job):
                 for file_data in filtered_file_data:
                     files_to_upload.append(
                         DataFileUpload(
-                            data=io.BytesIO(file_data.file_data), name=file_data.filename
+                            data=io.BytesIO(file_data.file_data),
+                            name=file_data.filename,
                         )
                     )
                 if not self.upload_strategy.skip_moving_files:
@@ -256,7 +312,9 @@ class GenericUploadJob(Job):
                         filesystem_session=filesystem_session,
                         remote_directory_scope=remote_directory,
                         success_file_paths=[
-                            file.filepath if file.filepath is not None else file.filename
+                            file.filepath
+                            if file.filepath is not None
+                            else file.filename
                             for file in filtered_file_data
                         ],
                         # IMPROVE: use triggers/webhooks to mark failed files as failed
@@ -267,12 +325,12 @@ class GenericUploadJob(Job):
 
         file_ids = [file.file_id for file in uploaded_files]
 
-        for material_family_key in self.upload_strategy.material_family_keys:
+        for destination in self.upload_strategy.destinations:
             for file_id in file_ids:
                 batch_processor.invoke_uploader(
                     file_id=file_id,
                     uploader_key=self.upload_strategy.uploader_key,
-                    material_family_key=material_family_key,
+                    destination=destination,
                 )
 
         return JobResult(success=True)

uncountable/integration/executors/script_executor.py

@@ -19,7 +19,7 @@ def resolve_script_executor(
     for _, job_class in inspect.getmembers(job_module, inspect.isclass):
         if getattr(job_class, "_unc_job_registered", False):
             found_jobs.append(job_class())
-    assert (
-        len(found_jobs) == 1
-    ), f"expected exactly one job class in {executor.import_path}, found {len(found_jobs)}"
+    assert len(found_jobs) == 1, (
+        f"expected exactly one job class in {executor.import_path}, found {len(found_jobs)}"
+    )
     return found_jobs[0]

uncountable/integration/http_server/__init__.py

@@ -0,0 +1,5 @@
+# CLOSED MODULE
+
+from .types import GenericHttpRequest as GenericHttpRequest
+from .types import GenericHttpResponse as GenericHttpResponse
+from .types import HttpException as HttpException

uncountable/integration/http_server/types.py

@@ -0,0 +1,69 @@
+import base64
+import functools
+import json
+from dataclasses import dataclass
+
+from flask.wrappers import Response
+
+
+class HttpException(Exception):
+    error_code: int
+    message: str
+
+    def __init__(self, *, error_code: int, message: str) -> None:
+        self.error_code = error_code
+        self.message = message
+
+    @staticmethod
+    def payload_failed_signature() -> "HttpException":
+        return HttpException(
+            error_code=401, message="webhook payload did not match signature"
+        )
+
+    @staticmethod
+    def no_signature_passed() -> "HttpException":
+        return HttpException(error_code=400, message="missing signature")
+
+    @staticmethod
+    def body_parse_error() -> "HttpException":
+        return HttpException(error_code=400, message="body parse error")
+
+    @staticmethod
+    def unknown_error() -> "HttpException":
+        return HttpException(error_code=500, message="internal server error")
+
+    @staticmethod
+    def configuration_error(
+        message: str = "internal configuration error",
+    ) -> "HttpException":
+        return HttpException(error_code=500, message=message)
+
+    def __str__(self) -> str:
+        return f"[{self.error_code}]: {self.message}"
+
+    def make_error_response(self) -> Response:
+        return Response(
+            status=self.error_code,
+            response=json.dumps({"error": {"message": str(self)}}),
+        )
+
+
+@dataclass(kw_only=True, frozen=True)
+class GenericHttpRequest:
+    body_base64: str
+    headers: dict[str, str]
+
+    @functools.cached_property
+    def body_bytes(self) -> bytes:
+        return base64.b64decode(self.body_base64)
+
+    @functools.cached_property
+    def body_text(self) -> str:
+        return self.body_bytes.decode()
+
+
+@dataclass(kw_only=True)
+class GenericHttpResponse:
+    response: str
+    status_code: int
+    headers: dict[str, str] | None = None

uncountable/integration/job.py

@@ -1,42 +1,272 @@
+import functools
+import hmac
+import typing
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
+import simplejson
+
+from pkgs.argument_parser import CachedParser
+from pkgs.serialization_util import serialize_for_api
 from uncountable.core.async_batch import AsyncBatchProcessor
 from uncountable.core.client import Client
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.core.file_upload import FileUpload
+from uncountable.core.types import AuthDetailsOAuth
+from uncountable.integration.http_server import (
+    GenericHttpRequest,
+    GenericHttpResponse,
+    HttpException,
+)
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_job_queue_message,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandServerException,
+)
+from uncountable.integration.secret_retrieval.retrieve_secret import retrieve_secret
 from uncountable.integration.telemetry import JobLogger
-from uncountable.types.job_definition_t import JobDefinition, JobResult, ProfileMetadata
+from uncountable.types import (
+    base_t,
+    job_definition_t,
+    queued_job_t,
+    webhook_job_t,
+)
+from uncountable.types.job_definition_t import (
+    HttpJobDefinitionBase,
+    JobDefinition,
+    JobResult,
+    ProfileMetadata,
+)
 
 
-@dataclass
-class JobArgumentsBase:
+@dataclass(kw_only=True)
+class JobArguments:
     job_definition: JobDefinition
     profile_metadata: ProfileMetadata
     client: Client
     batch_processor: AsyncBatchProcessor
     logger: JobLogger
+    payload: base_t.JsonValue
+    job_uuid: str
 
 
-@dataclass
-class CronJobArguments(JobArgumentsBase):
-    # can imagine passing additional data such as in the sftp or webhook cases
-    pass
+# only for compatibility:
+CronJobArguments = JobArguments
 
 
-JobArguments = CronJobArguments
+class Job[PT](ABC):
+    _unc_job_registered: bool = False
 
+    @property
+    @abstractmethod
+    def payload_type(self) -> type[PT]: ...
 
-class Job(ABC):
-    _unc_job_registered: bool = False
+    @abstractmethod
+    def run_outer(self, args: JobArguments) -> JobResult: ...
+
+    @functools.cached_property
+    def _cached_payload_parser(self) -> CachedParser[PT]:
+        return CachedParser(self.payload_type)
+
+    def get_payload(self, payload: base_t.JsonValue) -> PT:
+        return self._cached_payload_parser.parse_storage(payload)
+
+
+class CronJob(Job):
+    @property
+    def payload_type(self) -> type[None]:
+        return type(None)
+
+    def run_outer(self, args: JobArguments) -> JobResult:
+        assert isinstance(args, CronJobArguments)
+        return self.run(args)
 
     @abstractmethod
     def run(self, args: JobArguments) -> JobResult: ...
 
 
-class CronJob(Job):
+WPT = typing.TypeVar("WPT")
+
+
+@dataclass(kw_only=True)
+class WebhookResponse:
+    pass
+
+
+class _RequestValidatorClient(Client):
+    def __init__(self, *, base_url: str, oauth_bearer_token: str):
+        super().__init__(
+            base_url=base_url,
+            auth_details=AuthDetailsOAuth(refresh_token=""),
+            config=None,
+        )
+        self._oauth_bearer_token = oauth_bearer_token
+
+    def _get_oauth_bearer_token(self, *, oauth_details: AuthDetailsOAuth) -> str:
+        return self._oauth_bearer_token
+
+
+class CustomHttpJob(Job[GenericHttpRequest]):
+    @property
+    def payload_type(self) -> type[GenericHttpRequest]:
+        return GenericHttpRequest
+
+    @staticmethod
+    @abstractmethod
+    def validate_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+    ) -> None:
+        """
+        Validate that the request is valid. If the request is invalid, raise an
+        exception.
+        """
+        ...
+
+    @staticmethod
+    def get_validated_oauth_request_user_id(
+        *, profile_metadata: ProfileMetadata, request: GenericHttpRequest
+    ) -> base_t.ObjectId:
+        token = request.headers.get("Authorization", "").replace("Bearer ", "")
+        if token == "":
+            raise HttpException(
+                message="unauthorized; no bearer token in request", error_code=401
+            )
+        return (
+            _RequestValidatorClient(
+                base_url=profile_metadata.base_url,
+                oauth_bearer_token=token,
+            )
+            .get_current_user_info()
+            .user_id
+        )
+
+    @staticmethod
+    @abstractmethod
+    def handle_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+    ) -> GenericHttpResponse:
+        """
+        Handle the request synchronously. Normally this should just enqueue a job
+        and return immediately (see WebhookJob as an example).
+        """
+        ...
+
+    def run_outer(self, args: JobArguments) -> JobResult:
+        args.logger.log_warning(
+            message=f"Unexpected call to run_outer for CustomHttpJob: {args.job_definition.id}"
+        )
+        return JobResult(success=False)
+
+
+class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload]):
+    @property
+    def payload_type(self) -> type[webhook_job_t.WebhookEventPayload]:
+        return webhook_job_t.WebhookEventPayload
+
+    @property
     @abstractmethod
-    def run(self, args: CronJobArguments) -> JobResult: ...
+    def webhook_payload_type(self) -> type[WPT]: ...
+
+    @staticmethod
+    def validate_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: job_definition_t.HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+    ) -> None:
+        assert isinstance(job_definition, job_definition_t.WebhookJobDefinition)
+        signature_key = retrieve_secret(
+            profile_metadata=profile_meta,
+            secret_retrieval=job_definition.signature_key_secret,
+        )
+        passed_signature = request.headers.get("Uncountable-Webhook-Signature")
+        if passed_signature is None:
+            raise HttpException.no_signature_passed()
+
+        request_body_signature = hmac.new(
+            signature_key.encode("utf-8"), msg=request.body_bytes, digestmod="sha256"
+        ).hexdigest()
+
+        if request_body_signature != passed_signature:
+            raise HttpException.payload_failed_signature()
+
+    @staticmethod
+    def handle_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: job_definition_t.HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,  # noqa: ARG004
+    ) -> GenericHttpResponse:
+        try:
+            request_body = simplejson.loads(request.body_text)
+            webhook_payload = typing.cast(base_t.JsonValue, request_body)
+        except (simplejson.JSONDecodeError, ValueError) as e:
+            raise HttpException.body_parse_error() from e
+
+        try:
+            send_job_queue_message(
+                job_ref_name=job_definition.id,
+                payload=queued_job_t.QueuedJobPayload(
+                    invocation_context=queued_job_t.InvocationContextWebhook(
+                        webhook_payload=webhook_payload
+                    )
+                ),
+                port=get_local_admin_server_port(),
+            )
+        except CommandServerException as e:
+            raise HttpException.unknown_error() from e
+
+        return GenericHttpResponse(
+            response=simplejson.dumps(serialize_for_api(WebhookResponse())),
+            status_code=200,
+        )
+
+    def run_outer(self, args: JobArguments) -> JobResult:
+        webhook_body = self.get_payload(args.payload)
+        inner_payload = CachedParser(self.webhook_payload_type).parse_api(
+            webhook_body.data
+        )
+        return self.run(args, inner_payload)
+
+    @abstractmethod
+    def run(self, args: JobArguments, payload: WPT) -> JobResult: ...
 
 
 def register_job(cls: type[Job]) -> type[Job]:
     cls._unc_job_registered = True
     return cls
+
+
+class RunsheetWebhookJob(WebhookJob[webhook_job_t.RunsheetWebhookPayload]):
+    @property
+    def webhook_payload_type(self) -> type:
+        return webhook_job_t.RunsheetWebhookPayload
+
+    @abstractmethod
+    def build_runsheet(
+        self,
+        *,
+        args: JobArguments,
+        payload: webhook_job_t.RunsheetWebhookPayload,
+    ) -> FileUpload: ...
+
+    def run(
+        self, args: JobArguments, payload: webhook_job_t.RunsheetWebhookPayload
+    ) -> JobResult:
+        runsheet = self.build_runsheet(args=args, payload=payload)
+
+        files = args.client.upload_files(file_uploads=[runsheet])
+        args.client.complete_async_upload(
+            async_job_id=payload.async_job_id, file_id=files[0].file_id
+        )
+
+        return JobResult(
+            success=True,
+        )

uncountable/integration/queue_runner/command_server/__init__.py

@@ -0,0 +1,28 @@
+from .command_client import check_health, send_job_queue_message
+from .command_server import serve
+from .types import (
+    CommandEnqueueJob,
+    CommandEnqueueJobResponse,
+    CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
+    CommandServerBadResponse,
+    CommandServerException,
+    CommandServerTimeout,
+    CommandTask,
+)
+
+__all__: list[str] = [
+    "serve",
+    "check_health",
+    "send_job_queue_message",
+    "CommandEnqueueJob",
+    "CommandEnqueueJobResponse",
+    "CommandRetryJob",
+    "CommandRetryJobResponse",
+    "CommandTask",
+    "CommandQueue",
+    "CommandServerTimeout",
+    "CommandServerException",
+    "CommandServerBadResponse",
+]