feat: Add executions queries

agents-api/agents_api/queries/executions/count_executions.py

@@ -0,0 +1,39 @@
+from typing import Any, TypeVar
+from uuid import UUID
+
+import sqlvalidator
+from beartype import beartype
+
+from ..utils import (
+    pg_query,
+    wrap_in_class,
+)
+
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
+sql_query = sqlvalidator.parse(
+    """
+SELECT COUNT(*) FROM executions
+WHERE
+    developer_id = $1
+    AND task_id = $2
+"""
+)
+
+# @rewrap_exceptions(
+#     {
+#         QueryException: partialclass(HTTPException, status_code=400),
+#         ValidationError: partialclass(HTTPException, status_code=400),
+#         TypeError: partialclass(HTTPException, status_code=400),
+#     }
+# )
+@wrap_in_class(dict, one=True)
+@pg_query
+@beartype
+def count_executions(
+    *,
+    developer_id: UUID,
+    task_id: UUID,
+) -> tuple[list[str], dict]:
+    return (sql_query.format(), [developer_id, task_id])

agents-api/agents_api/queries/executions/get_execution.py

@@ -0,0 +1,52 @@
+from typing import Any, TypeVar
+from uuid import UUID
+
+from beartype import beartype
+
+import sqlvalidator
+from ...autogen.openapi_model import Execution
+from ..utils import (
+    pg_query,
+    wrap_in_class,
+)
+from .constants import OUTPUT_UNNEST_KEY
+
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
+sql_query = sqlvalidator.parse("""
+SELECT * FROM executions
+WHERE
+    execution_id = $1
+LIMIT 1
+""")
+
+
+# @rewrap_exceptions(
+#     {
+#         AssertionError: partialclass(HTTPException, status_code=404),
+#         QueryException: partialclass(HTTPException, status_code=400),
+#         ValidationError: partialclass(HTTPException, status_code=400),
+#         TypeError: partialclass(HTTPException, status_code=400),
+#     }
+# )
+@wrap_in_class(
+    Execution,
+    one=True,
+    transform=lambda d: {
+        **d,
+        "output": d["output"][OUTPUT_UNNEST_KEY]
+        if isinstance(d["output"], dict) and OUTPUT_UNNEST_KEY in d["output"]
+        else d["output"],
+    },
+)
+@pg_query
+@beartype
+def get_execution(
+    *,
+    execution_id: UUID,
+) -> tuple[str, dict]:
+    return (
+        sql_query.format(),
+        [execution_id],
+    )