UncountablePythonSDK 0.0.83__py3-none-any.whl → 0.0.132__py3-none-any.whl

uncountable/integration/executors/executors.py

@@ -9,6 +9,7 @@ from uncountable.types import (
     entity_t,
     field_values_t,
     identifier_t,
+    integration_server_t,
     job_definition_t,
     transition_entity_phase_t,
 )
@@ -33,7 +34,10 @@ def resolve_executor(
 
 
 def _create_run_entity(
-    *, client: Client, logging_settings: job_definition_t.JobLoggingSettings
+    *,
+    client: Client,
+    logging_settings: job_definition_t.JobLoggingSettings,
+    job_uuid: str,
 ) -> entity_t.Entity:
     run_entity = client.create_entity(
         entity_type=entity_t.EntityType.ASYNC_JOB,
@@ -49,6 +53,10 @@ def _create_run_entity(
                 field_ref_name=async_jobs_t.ASYNC_JOB_STATUS_FIELD_REF_NAME,
                 value=async_jobs_t.AsyncJobStatus.IN_PROGRESS,
             ),
+            field_values_t.FieldRefNameValue(
+                field_ref_name=integration_server_t.INTEGRATION_SERVER_RUN_UUID_FIELD_REF_NAME,
+                value=job_uuid,
+            ),
         ],
     ).entity
     client.transition_entity_phase(
@@ -93,7 +101,9 @@ def execute_job(
                 and job_definition.logging_settings.enabled
             ):
                 run_entity = _create_run_entity(
-                    client=args.client, logging_settings=job_definition.logging_settings
+                    client=args.client,
+                    logging_settings=job_definition.logging_settings,
+                    job_uuid=args.job_uuid,
                 )
             result = job.run_outer(args=args)
         except Exception as e:

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(
@@ -130,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]
@@ -184,16 +219,33 @@ class GenericUploadJob(Job[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:

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,14 +1,43 @@
 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 import base_t, webhook_job_t
-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(kw_only=True)
@@ -19,16 +48,14 @@ class JobArguments:
     batch_processor: AsyncBatchProcessor
     logger: JobLogger
     payload: base_t.JsonValue
+    job_uuid: str
 
 
 # only for compatibility:
 CronJobArguments = JobArguments
 
 
-PT = typing.TypeVar("PT")
-
-
-class Job(ABC, typing.Generic[PT]):
+class Job[PT](ABC):
     _unc_job_registered: bool = False
 
     @property
@@ -62,7 +89,83 @@ class CronJob(Job):
 WPT = typing.TypeVar("WPT")
 
 
-class WebhookJob(Job[webhook_job_t.WebhookEventPayload], typing.Generic[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
@@ -71,6 +174,60 @@ class WebhookJob(Job[webhook_job_t.WebhookEventPayload], typing.Generic[WPT]):
     @abstractmethod
     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(
@@ -85,3 +242,31 @@ class WebhookJob(Job[webhook_job_t.WebhookEventPayload], typing.Generic[WPT]):
 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

@@ -4,6 +4,8 @@ from .types import (
     CommandEnqueueJob,
     CommandEnqueueJobResponse,
     CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
     CommandServerBadResponse,
     CommandServerException,
     CommandServerTimeout,
@@ -16,6 +18,8 @@ __all__: list[str] = [
     "send_job_queue_message",
     "CommandEnqueueJob",
     "CommandEnqueueJobResponse",
+    "CommandRetryJob",
+    "CommandRetryJobResponse",
     "CommandTask",
     "CommandQueue",
     "CommandServerTimeout",

uncountable/integration/queue_runner/command_server/command_client.py

@@ -10,6 +10,12 @@ from uncountable.integration.queue_runner.command_server.protocol.command_server
     CheckHealthResult,
     EnqueueJobRequest,
     EnqueueJobResult,
+    ListQueuedJobsRequest,
+    ListQueuedJobsResult,
+    RetryJobRequest,
+    RetryJobResult,
+    VaccuumQueuedJobsRequest,
+    VaccuumQueuedJobsResult,
 )
 from uncountable.integration.queue_runner.command_server.types import (
     CommandServerBadResponse,
@@ -57,6 +63,26 @@ def send_job_queue_message(
         return response.queued_job_uuid
 
 
+def send_retry_job_message(
+    *,
+    job_uuid: str,
+    host: str = "localhost",
+    port: int,
+) -> str:
+    with command_server_connection(host=host, port=port) as stub:
+        request = RetryJobRequest(uuid=job_uuid)
+
+        try:
+            response = stub.RetryJob(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+            assert isinstance(response, RetryJobResult)
+            if not response.successfully_queued:
+                raise CommandServerBadResponse("queue operation was not successful")
+
+            return response.queued_job_uuid
+        except grpc.RpcError as e:
+            raise ValueError(e.details())  # type: ignore
+
+
 def check_health(*, host: str = _LOCAL_RPC_HOST, port: int) -> bool:
     with command_server_connection(host=host, port=port) as stub:
         request = CheckHealthRequest()
@@ -66,3 +92,42 @@ def check_health(*, host: str = _LOCAL_RPC_HOST, port: int) -> bool:
         assert isinstance(response, CheckHealthResult)
 
         return response.success
+
+
+def send_list_queued_jobs_message(
+    *,
+    offset: int,
+    limit: int,
+    host: str = "localhost",
+    port: int,
+) -> list[ListQueuedJobsResult.ListQueuedJobsResultItem]:
+    with command_server_connection(host=host, port=port) as stub:
+        request = ListQueuedJobsRequest(
+            offset=offset,
+            limit=limit,
+        )
+
+        try:
+            response = stub.ListQueuedJobs(
+                request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS
+            )
+        except grpc.RpcError as e:
+            raise ValueError(e.details())  # type: ignore
+
+        assert isinstance(response, ListQueuedJobsResult)
+        return list(response.queued_jobs)
+
+
+def send_vaccuum_queued_jobs_message(*, host: str = "localhost", port: int) -> None:
+    with command_server_connection(host=host, port=port) as stub:
+        request = VaccuumQueuedJobsRequest()
+
+        try:
+            response = stub.VaccuumQueuedJobs(
+                request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS
+            )
+        except grpc.RpcError as e:
+            raise ValueError(e.details())  # type: ignore
+
+        assert isinstance(response, VaccuumQueuedJobsResult)
+        return None