UncountablePythonSDK 0.0.114__py3-none-any.whl → 0.0.116__py3-none-any.whl

pkgs/type_spec/ui_entry_actions/generate_ui_entry_actions.py

@@ -1,13 +1,14 @@
 import json
+import re
 from dataclasses import dataclass
 from io import StringIO
 from pathlib import Path
 from typing import assert_never
 
 from main.base.types import (
+    base_t,
     ui_entry_actions_t,
 )
-from pkgs.argument_parser import snake_to_camel_case
 from pkgs.serialization_util import serialize_for_api
 from pkgs.type_spec import emit_typescript_util
 from pkgs.type_spec.builder import (
@@ -45,8 +46,8 @@ def ui_entry_variable_to_type_spec_type(
     match variable:
         case ui_entry_actions_t.UiEntryActionVariableString():
             return BaseTypeName.s_string
-        case ui_entry_actions_t.UiEntryActionVariableEntityFields():
-            return "entity.Entity"
+        case ui_entry_actions_t.UiEntryActionVariableSingleEntity():
+            return "ObjectId"
         case _:
             assert_never(variable)
 
@@ -157,6 +158,15 @@ def emit_entry_action_definition(
     )
 
 
+def _validate_input(input: ui_entry_actions_t.UiEntryActionVariable) -> None:
+    if "_" in input.vs_var_name:
+        raise ValueError(f"Expected camelCase for variable {input.vs_var_name}")
+    if not re.fullmatch(base_t.REF_NAME_STRICT_REGEX, input.vs_var_name):
+        raise ValueError(
+            f"Variable {input.vs_var_name} has invalid syntax. See REF_NAME_STRICT_REGEX"
+        )
+
+
 def emit_query_index(
     ctx: emit_typescript_util.EmitTypescriptContext,
     defn_infos: list[EntryActionTypeInfo],
@@ -182,6 +192,7 @@ def emit_query_index(
             "type": BaseTypeName.s_object,
             "properties": {
                 "value_spec_var": {"type": "String"},
+                "type": {"type": "ui_entry_actions.UiEntryActionDataType"},
                 "variable": {"type": "ui_entry_actions.UiEntryActionVariable"},
             },
         },
@@ -215,10 +226,12 @@ def emit_query_index(
     for scope, defn in definitions.items():
         inputs = []
         outputs = []
-        for name, input in defn.inputs.items():
+        for input in defn.inputs.values():
+            _validate_input(input)
             inputs.append(
                 serialize_for_api({
-                    "value_spec_var": snake_to_camel_case(name),
+                    "value_spec_var": input.vs_var_name,
+                    "type": input.type,
                     "variable": input,
                 })
             )
@@ -269,6 +282,7 @@ def generate_entry_actions_typescript(
     ctx = emit_typescript_util.EmitTypescriptContext(
         out=definition_buffer,
         namespace=index_namespace,
+        api_endpoints={},
     )
     builder.namespaces[index_namespace.name] = index_namespace
 

uncountable/core/environment.py

@@ -18,7 +18,7 @@ def get_server_env() -> str | None:
     return os.environ.get("UNC_SERVER_ENV")
 
 
-def get_webhook_server_port() -> int:
+def get_http_server_port() -> int:
     return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT", "5001"))
 
 

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,67 @@
+import base64
+import functools
+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={"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,15 +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.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, entity_t, webhook_job_t
-from uncountable.types.job_definition_t import JobDefinition, JobResult, ProfileMetadata
+from uncountable.types import (
+    base_t,
+    entity_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)
@@ -26,9 +54,6 @@ class JobArguments:
 CronJobArguments = JobArguments
 
 
-PT = typing.TypeVar("PT")
-
-
 class Job[PT](ABC):
     _unc_job_registered: bool = False
 
@@ -63,6 +88,51 @@ class CronJob(Job):
 WPT = typing.TypeVar("WPT")
 
 
+@dataclass(kw_only=True)
+class WebhookResponse:
+    pass
+
+
+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
+    @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]:
@@ -72,6 +142,60 @@ class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload]):
     @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(

uncountable/integration/server.py

@@ -16,7 +16,7 @@ from uncountable.integration.telemetry import Logger
 from uncountable.types import base_t, job_definition_t
 from uncountable.types.job_definition_t import (
     CronJobDefinition,
-    WebhookJobDefinition,
+    HttpJobDefinitionBase,
 )
 
 _MAX_APSCHEDULER_CONCURRENT_JOBS = 1
@@ -86,7 +86,7 @@ class IntegrationServer:
                                 misfire_grace_time=None,
                                 **job_opts,
                             )
-                    case WebhookJobDefinition():
+                    case HttpJobDefinitionBase():
                         pass
                     case _:
                         assert_never(job_defn)

uncountable/integration/telemetry.py

@@ -177,7 +177,7 @@ class JobLogger(Logger):
                 patched_attributes["job.definition.cron_spec"] = (
                     self.job_definition.cron_spec
                 )
-            case job_definition_t.WebhookJobDefinition():
+            case job_definition_t.HttpJobDefinitionBase():
                 pass
             case _:
                 assert_never(self.job_definition)

uncountable/integration/webhook_server/entrypoint.py

@@ -1,146 +1,71 @@
-import hmac
-import typing
-from dataclasses import dataclass
+import base64
 
 import flask
-import simplejson
 from flask.typing import ResponseReturnValue
-from flask.wrappers import Response
 from opentelemetry.trace import get_current_span
 from uncountable.core.environment import (
-    get_local_admin_server_port,
+    get_http_server_port,
     get_server_env,
-    get_webhook_server_port,
-)
-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.executors.script_executor import resolve_script_executor
+from uncountable.integration.http_server import GenericHttpRequest, HttpException
+from uncountable.integration.job import CustomHttpJob, WebhookJob
 from uncountable.integration.scan_profiles import load_profiles
-from uncountable.integration.secret_retrieval.retrieve_secret import retrieve_secret
 from uncountable.integration.telemetry import Logger
-from uncountable.types import base_t, job_definition_t, queued_job_t, webhook_job_t
-
-from pkgs.argument_parser import CachedParser
+from uncountable.types import job_definition_t
 
 app = flask.Flask(__name__)
 
 
-@dataclass(kw_only=True)
-class WebhookResponse:
-    pass
-
-
-webhook_payload_parser = CachedParser(webhook_job_t.WebhookEventBody)
-
-
-class WebhookException(BaseException):
-    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() -> "WebhookException":
-        return WebhookException(
-            error_code=401, message="webhook payload did not match signature"
-        )
-
-    @staticmethod
-    def no_signature_passed() -> "WebhookException":
-        return WebhookException(error_code=400, message="missing signature")
-
-    @staticmethod
-    def body_parse_error() -> "WebhookException":
-        return WebhookException(error_code=400, message="body parse error")
-
-    @staticmethod
-    def unknown_error() -> "WebhookException":
-        return WebhookException(error_code=500, message="internal server error")
-
-    def __str__(self) -> str:
-        return f"[{self.error_code}]: {self.message}"
-
-    def make_error_response(self) -> Response:
-        return Response(
-            status=self.error_code, response={"error": {"message": str(self)}}
-        )
-
-
-def _parse_webhook_payload(
-    *, raw_request_body: bytes, signature_key: str, passed_signature: str
-) -> base_t.JsonValue:
-    request_body_signature = hmac.new(
-        signature_key.encode("utf-8"), msg=raw_request_body, digestmod="sha256"
-    ).hexdigest()
-
-    if request_body_signature != passed_signature:
-        raise WebhookException.payload_failed_signature()
-
-    try:
-        request_body = simplejson.loads(raw_request_body.decode())
-        return typing.cast(base_t.JsonValue, request_body)
-    except (simplejson.JSONDecodeError, ValueError) as e:
-        raise WebhookException.body_parse_error() from e
-
-
 def register_route(
     *,
     server_logger: Logger,
     profile_meta: job_definition_t.ProfileMetadata,
-    job: job_definition_t.WebhookJobDefinition,
+    job: job_definition_t.HttpJobDefinitionBase,
 ) -> None:
     route = f"/{profile_meta.name}/{job.id}"
 
-    def handle_webhook() -> ResponseReturnValue:
+    def handle_request() -> ResponseReturnValue:
         with server_logger.push_scope(route):
             try:
-                signature_key = retrieve_secret(
-                    profile_metadata=profile_meta,
-                    secret_retrieval=job.signature_key_secret,
+                if not isinstance(job.executor, job_definition_t.JobExecutorScript):
+                    raise HttpException.configuration_error(
+                        message="[internal] http job must use a script executor"
+                    )
+                job_instance = resolve_script_executor(
+                    executor=job.executor, profile_metadata=profile_meta
                 )
-
-                passed_signature = flask.request.headers.get(
-                    "Uncountable-Webhook-Signature"
+                if not isinstance(job_instance, (CustomHttpJob, WebhookJob)):
+                    raise HttpException.configuration_error(
+                        message="[internal] http job must descend from CustomHttpJob"
+                    )
+                http_request = GenericHttpRequest(
+                    body_base64=base64.b64encode(flask.request.get_data()).decode(),
+                    headers=dict(flask.request.headers),
                 )
-                if passed_signature is None:
-                    raise WebhookException.no_signature_passed()
-
-                webhook_payload = _parse_webhook_payload(
-                    raw_request_body=flask.request.data,
-                    signature_key=signature_key,
-                    passed_signature=passed_signature,
+                job_instance.validate_request(
+                    request=http_request, job_definition=job, profile_meta=profile_meta
+                )
+                http_response = job_instance.handle_request(
+                    request=http_request, job_definition=job, profile_meta=profile_meta
                 )
 
-                try:
-                    send_job_queue_message(
-                        job_ref_name=job.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 WebhookException.unknown_error() from e
-
-                return flask.jsonify(WebhookResponse())
-            except WebhookException as e:
+                return flask.make_response(
+                    http_response.response,
+                    http_response.status_code,
+                    http_response.headers,
+                )
+            except HttpException as e:
                 server_logger.log_exception(e)
                 return e.make_error_response()
             except Exception as e:
                 server_logger.log_exception(e)
-                return WebhookException.unknown_error().make_error_response()
+                return HttpException.unknown_error().make_error_response()
 
     app.add_url_rule(
         route,
-        endpoint=f"handle_webhook_{job.id}",
-        view_func=handle_webhook,
+        endpoint=f"handle_request_{job.id}",
+        view_func=handle_request,
         methods=["POST"],
     )
 
@@ -152,7 +77,7 @@ def main() -> None:
     for profile_metadata in profiles:
         server_logger = Logger(get_current_span())
         for job in profile_metadata.jobs:
-            if isinstance(job, job_definition_t.WebhookJobDefinition):
+            if isinstance(job, job_definition_t.HttpJobDefinitionBase):
                 register_route(
                     server_logger=server_logger, profile_meta=profile_metadata, job=job
                 )
@@ -164,7 +89,7 @@ main()
 if __name__ == "__main__":
     app.run(
         host="0.0.0.0",
-        port=get_webhook_server_port(),
+        port=get_http_server_port(),
         debug=get_server_env() == "playground",
         exclude_patterns=[],
     )

uncountable/types/api/entity/create_or_update_entity.py

@@ -35,6 +35,7 @@ class Arguments:
     definition_key: identifier_t.IdentifierKey
     field_values: list[field_values_t.FieldArgumentValue]
     entity_key: identifier_t.IdentifierKey | None = None
+    on_create_init_field_values: list[field_values_t.FieldArgumentValue] | None = None
 
 
 # DO NOT MODIFY -- This file is generated by type_spec

uncountable/types/api/entity/lookup_entity.py

@@ -21,6 +21,7 @@ __all__: list[str] = [
     "Data",
     "ENDPOINT_METHOD",
     "ENDPOINT_PATH",
+    "LookupEntityCompositeFieldValues",
     "LookupEntityFieldValue",
     "LookupEntityQuery",
     "LookupEntityQueryBase",
@@ -35,6 +36,7 @@ ENDPOINT_PATH = "api/external/entity/lookup_entity"
 # DO NOT MODIFY -- This file is generated by type_spec
 class LookupEntityQueryType(StrEnum):
     FIELD_VALUE = "field_value"
+    COMPOSITE_FIELD_VALUES = "composite_field_values"
 
 
 # DO NOT MODIFY -- This file is generated by type_spec
@@ -69,14 +71,26 @@ class LookupEntityFieldValue:
     value: LookupFieldArgumentValue
 
 
+# DO NOT MODIFY -- This file is generated by type_spec
+@serial_class(
+    named_type_path="sdk.api.entity.lookup_entity.LookupEntityCompositeFieldValues",
+    parse_require={"type"},
+)
+@dataclasses.dataclass(slots=base_t.ENABLE_SLOTS, kw_only=True)  # type: ignore[literal-required]
+class LookupEntityCompositeFieldValues:
+    type: typing.Literal[LookupEntityQueryType.COMPOSITE_FIELD_VALUES] = LookupEntityQueryType.COMPOSITE_FIELD_VALUES
+    values: list[LookupFieldArgumentValue]
+
+
 # DO NOT MODIFY -- This file is generated by type_spec
 LookupEntityQuery = typing.Annotated[
-    typing.Union[LookupEntityFieldValue],
+    LookupEntityFieldValue | LookupEntityCompositeFieldValues,
     serial_union_annotation(
     named_type_path="sdk.api.entity.lookup_entity.LookupEntityQuery",
         discriminator="type",
         discriminator_map={
             "field_value": LookupEntityFieldValue,
+            "composite_field_values": LookupEntityCompositeFieldValues,
         },
     ),
 ]

uncountable/types/async_batch_processor.py

@@ -258,10 +258,12 @@ class AsyncBatchProcessorBase(ABC):
         definition_key: identifier_t.IdentifierKey,
         field_values: list[field_values_t.FieldArgumentValue],
         entity_key: identifier_t.IdentifierKey | None = None,
+        on_create_init_field_values: list[field_values_t.FieldArgumentValue] | None = None,
         depends_on: list[str] | None = None,
     ) -> async_batch_t.QueuedAsyncBatchRequest:
         """Creates or updates field values for an entity
 
+        :param on_create_init_field_values: Field values set only when the entity is created (will be ignored if entity already exists)
         :param depends_on: A list of batch reference keys to process before processing this request
         """
         args = create_or_update_entity_t.Arguments(
@@ -269,6 +271,7 @@ class AsyncBatchProcessorBase(ABC):
             entity_type=entity_type,
             definition_key=definition_key,
             field_values=field_values,
+            on_create_init_field_values=on_create_init_field_values,
         )
         json_data = serialize_for_api(args)