feat(engine)!: Implement actions registry and template action managem…

…ent (#408)

.github/workflows/test-python.yml

@@ -65,6 +65,8 @@ jobs:
         run: bash env.sh
 
       - name: Start Docker services
+        env:
+          TRACECAT__UNSAFE_DISABLE_SM_MASKING: "true"
         run: docker compose -f docker-compose.dev.yml up --build --no-deps -d api worker postgres_db caddy
 
       - name: Verify Tracecat API is running

.vscode/settings.json

@@ -4,5 +4,6 @@
   "python.testing.pytestEnabled": true,
   "yaml.schemas": {
     "https://json.schemastore.org/yamllint.json": "playbooks/*.yml"
-  }
+  },
+  "python.analysis.extraPaths": ["./registry"]
 }

alembic/env.py

@@ -8,15 +8,16 @@
 from alembic import context
 from tracecat.db import schemas  # noqa: F401
 
-
 TRACECAT__DB_URI = os.getenv("TRACECAT__DB_URI")
 if not TRACECAT__DB_URI:
     username = os.getenv("TRACECAT__DB_USER", "postgres")
     password = os.getenv("TRACECAT__DB_PASS")
     host = os.getenv("TRACECAT__DB_ENDPOINT")
     port = os.getenv("TRACECAT__DB_PORT", 5432)
     database = os.getenv("TRACECAT__DB_NAME", "postgres")
-    TRACECAT__DB_URI = f"postgresql+psycopg://{username}:{password}@{host}:{port!s}/{database}"
+    TRACECAT__DB_URI = (
+        f"postgresql+psycopg://{username}:{password}@{host}:{port!s}/{database}"
+    )
 
 
 # this is the Alembic Config object, which provides
@@ -59,7 +60,7 @@ def run_migrations_online() -> None:
         configuration=config.get_section(config.config_ini_section, {}),
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
-        url=TRACECAT__DB_URI
+        url=TRACECAT__DB_URI,
     )
 
     with connectable.connect() as connection:

alembic/versions/5273ec029b1b_add_registry_repository_and_action_.py

@@ -0,0 +1,139 @@
+"""Add registry repository and action tables
+
+Revision ID: 5273ec029b1b
+Revises: 5308ffdd79f3
+Create Date: 2024-10-03 19:54:16.459544
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+import sqlmodel
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "5273ec029b1b"
+down_revision: str | None = "5308ffdd79f3"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "registryrepository",
+        sa.Column("surrogate_id", sa.Integer(), nullable=False),
+        sa.Column("owner_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("(now() AT TIME ZONE 'utc'::text)"),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("(now() AT TIME ZONE 'utc'::text)"),
+            nullable=False,
+        ),
+        sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("version", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("origin", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.PrimaryKeyConstraint("surrogate_id"),
+        sa.UniqueConstraint("id"),
+    )
+    op.create_table(
+        "registryaction",
+        sa.Column("surrogate_id", sa.Integer(), nullable=False),
+        sa.Column("owner_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("(now() AT TIME ZONE 'utc'::text)"),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("(now() AT TIME ZONE 'utc'::text)"),
+            nullable=False,
+        ),
+        sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("namespace", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("version", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("origin", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("type", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("default_title", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("display_group", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("secrets", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+        sa.Column("interface", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+        sa.Column(
+            "implementation", postgresql.JSONB(astext_type=sa.Text()), nullable=True
+        ),
+        sa.Column("options", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+        sa.Column("repository_id", sa.UUID(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["repository_id"], ["registryrepository.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("surrogate_id"),
+        sa.UniqueConstraint("id"),
+        sa.UniqueConstraint(
+            "namespace",
+            "name",
+            "version",
+            name="uq_registry_action_namespace_name_version",
+        ),
+    )
+    op.drop_index("ix_udfspec_id", table_name="udfspec")
+    op.drop_table("udfspec")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "udfspec",
+        sa.Column("surrogate_id", sa.INTEGER(), autoincrement=True, nullable=False),
+        sa.Column("owner_id", sa.UUID(), autoincrement=False, nullable=False),
+        sa.Column(
+            "created_at",
+            postgresql.TIMESTAMP(timezone=True),
+            server_default=sa.text("(now() AT TIME ZONE 'utc'::text)"),
+            autoincrement=False,
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            postgresql.TIMESTAMP(timezone=True),
+            server_default=sa.text("(now() AT TIME ZONE 'utc'::text)"),
+            autoincrement=False,
+            nullable=False,
+        ),
+        sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False),
+        sa.Column("description", sa.VARCHAR(), autoincrement=False, nullable=False),
+        sa.Column("namespace", sa.VARCHAR(), autoincrement=False, nullable=False),
+        sa.Column("key", sa.VARCHAR(), autoincrement=False, nullable=False),
+        sa.Column("version", sa.VARCHAR(), autoincrement=False, nullable=True),
+        sa.Column(
+            "json_schema",
+            postgresql.JSONB(astext_type=sa.Text()),
+            autoincrement=False,
+            nullable=True,
+        ),
+        sa.Column(
+            "meta",
+            postgresql.JSONB(astext_type=sa.Text()),
+            autoincrement=False,
+            nullable=True,
+        ),
+        sa.PrimaryKeyConstraint("surrogate_id", name="udfspec_pkey"),
+    )
+    op.create_index("ix_udfspec_id", "udfspec", ["id"], unique=True)
+    op.drop_table("registryaction")
+    op.drop_table("registryrepository")
+    # ### end Alembic commands ###

cli/pyproject.toml

@@ -78,6 +78,6 @@ markers = [
     "slow: marks tests as slow",
 ]
 
-[tool.ruff.pyupgrade]
+[tool.ruff.lint.pyupgrade]
 # Preserve types, even if a file imports `from __future__ import annotations`.
 keep-runtime-typing = true

cli/tracecat_cli/dev.py

@@ -94,13 +94,19 @@ def api(
     endpoint: str = typer.Argument(..., help="Endpoint to hit"),
     method: str = typer.Option("GET", "-X", help="HTTP Method"),
     data: str = typer.Option(None, "--data", "-d", help="JSON Payload to send"),
+    to_json: bool = typer.Option(
+        False, "--json", "-j", help="Output the response as JSON"
+    ),
 ):
     """Commit a workflow definition to the database."""
 
     payload = read_input(data) if data else None
     result = hit_api_endpoint(method, endpoint, payload)
-    rich.print("Hit the endpoint successfully!")
-    rich.print(result, len(result))
+    if to_json:
+        rich.print(orjson.dumps(result, option=orjson.OPT_INDENT_2).decode())
+    else:
+        rich.print("Hit the endpoint successfully!")
+        rich.print(result, len(result))
 
 
 @app.command(name="generate-spec", help="Generate OpenAPI specification. Requires npx.")
@@ -337,13 +343,13 @@ def generate_udf_docs():
         pass
     path.mkdir(parents=True, exist_ok=True)
 
-    from tracecat.registry import registry
+    from tracecat.registry import repository
 
-    registry.init()
+    repository.init()
 
     rich.print(f"Generating API reference paths in {config.docs_path!s}")
 
-    for key, udf in registry:
+    for key, udf in repository:
         if not udf.metadata["include_in_schema"]:
             continue
 
@@ -385,7 +391,9 @@ def generate_udf_docs():
     # Overwrite the 'navigation' property with the new JSON data
     gname = "Schemas"
     # Find 'Schemas' group
-    filtered_keys = [key for key in registry.keys if udf.metadata["include_in_schema"]]
+    filtered_keys = [
+        key for key in repository.keys if udf.metadata["include_in_schema"]
+    ]
     new_mint_pages = key_tree_to_pages(filtered_keys, int_relpath).model_dump()["pages"]
     try:
         schemas_ref = next(
@@ -434,15 +442,15 @@ def generate_secret_tables(
         help="Output file path",
     ),
 ):
-    from tracecat.registry import registry
+    from tracecat.registry import repository
 
-    registry.init()
+    repository.init()
 
     # Table of core UDFs required secrets
     secrets = defaultdict(list)
     # Get UDF -> Secrets mapping
     blacklist = {"example"}
-    for key, udf in registry:
+    for key, udf in repository:
         top_level_ns, *stem, func = key.split(".")
         if top_level_ns in blacklist:
             continue
@@ -458,7 +466,7 @@ def generate_secret_tables(
 
     # Get all secrets -> secret keys
     api_credentials = set()
-    for secret in chain.from_iterable(udf.secrets or [] for _, udf in registry):
+    for secret in chain.from_iterable(udf.secrets or [] for _, udf in repository):
         api_credentials.add(
             (
                 wrap(secret.name, "`"),

docker-compose.dev.yml

@@ -20,6 +20,7 @@ services:
     ports:
       - 8000:8000
     environment:
+      # App
       DUMP_TRACECAT_RESULT: 0
       LOG_LEVEL: ${LOG_LEVEL}
       TRACECAT__API_URL: ${TRACECAT__API_URL}
@@ -53,6 +54,8 @@ services:
       RUN_MIGRATIONS: "true"
       # Remote registry
       TRACECAT__REMOTE_REGISTRY_URL: ${TRACECAT__REMOTE_REGISTRY_URL}
+      # This is only used for testing
+      TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false}
     volumes:
       - ./tracecat:/app/tracecat
       - ./registry:/app/registry
@@ -75,21 +78,9 @@ services:
       TRACECAT__PUBLIC_RUNNER_URL: ${TRACECAT__PUBLIC_RUNNER_URL}
       TRACECAT__SERVICE_KEY: ${TRACECAT__SERVICE_KEY} # Sensitive
       TRACECAT__SIGNING_SECRET: ${TRACECAT__SIGNING_SECRET} # Sensitive
-      # Remote registry
-      TRACECAT__REMOTE_REGISTRY_URL: ${TRACECAT__REMOTE_REGISTRY_URL}
       # Temporal
       TEMPORAL__CLUSTER_URL: ${TEMPORAL__CLUSTER_URL}
       TEMPORAL__CLUSTER_QUEUE: ${TEMPORAL__CLUSTER_QUEUE}
-      # SMTP
-      SMTP_HOST: ${SMTP_HOST}
-      SMTP_PORT: ${SMTP_PORT}
-      SMTP_STARTTLS_ENABLED: ${SMTP_STARTTLS_ENABLED}
-      SMTP_SSL_ENABLED: ${SMTP_SSL_ENABLED}
-      SMTP_IGNORE_CERT_ERRORS: ${SMTP_IGNORE_CERT_ERRORS}
-      SMTP_AUTH_ENABLED: ${SMTP_AUTH_ENABLED}
-      SMTP_USER: ${SMTP_USER}
-      SMTP_PASS: ${SMTP_PASS}
-      TRACECAT__UNSAFE_DISABLE_SM_MASKING: ${TRACECAT__UNSAFE_DISABLE_SM_MASKING:-false}
     volumes:
       - ./tracecat:/app/tracecat
       - ./registry:/app/registry

docker-compose.yml

@@ -22,6 +22,7 @@ services:
       - core-db
       - temporal
     environment:
+      # App
       LOG_LEVEL: ${LOG_LEVEL}
       TRACECAT__API_URL: ${TRACECAT__API_URL}
       TRACECAT__API_ROOT_PATH: ${TRACECAT__API_ROOT_PATH}
@@ -42,6 +43,20 @@ services:
       # Temporal
       TEMPORAL__CLUSTER_URL: ${TEMPORAL__CLUSTER_URL}
       TEMPORAL__CLUSTER_QUEUE: ${TEMPORAL__CLUSTER_QUEUE}
+      # SMTP
+      SMTP_HOST: ${SMTP_HOST}
+      SMTP_PORT: ${SMTP_PORT}
+      SMTP_STARTTLS_ENABLED: ${SMTP_STARTTLS_ENABLED}
+      SMTP_SSL_ENABLED: ${SMTP_SSL_ENABLED}
+      SMTP_IGNORE_CERT_ERRORS: ${SMTP_IGNORE_CERT_ERRORS}
+      SMTP_AUTH_ENABLED: ${SMTP_AUTH_ENABLED}
+      SMTP_USER: ${SMTP_USER}
+      SMTP_PASS: ${SMTP_PASS}
+      # LDAP
+      LDAP_HOST: ${LDAP_HOST}
+      LDAP_PORT: ${LDAP_PORT}
+      LDAP_SSL: ${LDAP_SSL}
+      LDAP_TYPE: ${LDAP_TYPE}
       # Remote registry
       TRACECAT__REMOTE_REGISTRY_URL: ${TRACECAT__REMOTE_REGISTRY_URL}
 
@@ -66,22 +81,6 @@ services:
       # Temporal
       TEMPORAL__CLUSTER_URL: ${TEMPORAL__CLUSTER_URL}
       TEMPORAL__CLUSTER_QUEUE: ${TEMPORAL__CLUSTER_QUEUE}
-      # SMTP
-      SMTP_HOST: ${SMTP_HOST}
-      SMTP_PORT: ${SMTP_PORT}
-      SMTP_STARTTLS_ENABLED: ${SMTP_STARTTLS_ENABLED}
-      SMTP_SSL_ENABLED: ${SMTP_SSL_ENABLED}
-      SMTP_IGNORE_CERT_ERRORS: ${SMTP_IGNORE_CERT_ERRORS}
-      SMTP_AUTH_ENABLED: ${SMTP_AUTH_ENABLED}
-      SMTP_USER: ${SMTP_USER}
-      SMTP_PASS: ${SMTP_PASS}
-      # LDAP
-      LDAP_HOST: ${LDAP_HOST}
-      LDAP_PORT: ${LDAP_PORT}
-      LDAP_SSL: ${LDAP_SSL}
-      LDAP_TYPE: ${LDAP_TYPE}
-      # Remote registry
-      TRACECAT__REMOTE_REGISTRY_URL: ${TRACECAT__REMOTE_REGISTRY_URL}
 
     command: ["python", "tracecat/dsl/worker.py"]