labrun-checks 0.2.0__py3-none-any.whl

labrun_checks/__init__.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import Callable
+
+from ._api import post_checkpoint, post_cleanup_complete
+from ._cleanup import run_cleanup
+from ._state import get_session, set_session
+from ._types import (
+    CheckpointResult,
+    CheckpointSession,
+    CleanupAdapter,
+    CleanupResource,
+    CleanupResult,
+)
+
+__all__ = [
+    "register",
+    "check",
+    "cleanup",
+    "run_cleanup",
+    "CheckpointResult",
+    "CheckpointSession",
+    "CleanupAdapter",
+    "CleanupResource",
+    "CleanupResult",
+]
+
+
+def register(session_token: str, lab_id: str, credentials: dict | None = None) -> None:
+    set_session(CheckpointSession(session_token=session_token, lab_id=lab_id, credentials=credentials or {}))
+
+
+def check(checkpoint_id: int, validator_fn: Callable[[CheckpointSession], CheckpointResult]) -> CheckpointResult:
+    session = get_session()
+
+    try:
+        result = validator_fn(session)
+    except Exception as exc:
+        result = CheckpointResult("error", error_reason=repr(exc))
+
+    try:
+        post_checkpoint(session.session_token, session.lab_id, checkpoint_id, result)
+    except Exception as exc:
+        return CheckpointResult(
+            "error",
+            error_reason=f"Network error recording checkpoint: {repr(exc)}",
+        )
+
+    return result
+
+
+def cleanup(status: str) -> CheckpointResult:
+    if status not in ("clean", "dirty"):
+        return CheckpointResult(
+            "error",
+            error_reason=f"status must be 'clean' or 'dirty', got {status!r}",
+        )
+
+    session = get_session()
+
+    try:
+        post_cleanup_complete(session.session_token, session.lab_id, status)
+    except Exception as exc:
+        return CheckpointResult(
+            "error",
+            error_reason=f"Network error posting cleanup event: {repr(exc)}",
+        )
+
+    return CheckpointResult("pass" if status == "clean" else "fail")

labrun_checks/_api.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import os
+import time
+
+import requests
+
+from ._types import CheckpointResult
+
+_BASE_URL = os.environ.get("LABRUN_API_BASE_URL", "https://labrun.dev")
+
+
+def post_checkpoint(
+    session_token: str,
+    lab_id: str,
+    checkpoint_id: int,
+    result: CheckpointResult,
+) -> None:
+    payload = {
+        "lab_id": lab_id,
+        "checkpoint_id": checkpoint_id,
+        "status": result.status,
+        "timestamp": int(time.time()),
+    }
+    headers = {"Authorization": f"Bearer {session_token}"}
+    response = requests.post(
+        f"{_BASE_URL}/api/labs/checkpoint",
+        json=payload,
+        headers=headers,
+        timeout=10,
+    )
+    response.raise_for_status()
+
+
+def post_cleanup_complete(session_token: str, lab_id: str, status: str) -> None:
+    payload = {
+        "lab_id": lab_id,
+        "status": status,
+    }
+    headers = {"Authorization": f"Bearer {session_token}"}
+    response = requests.post(
+        f"{_BASE_URL}/api/labs/cleanup-complete",
+        json=payload,
+        headers=headers,
+        timeout=10,
+    )
+    response.raise_for_status()

labrun_checks/_cleanup.py

@@ -0,0 +1,97 @@
+"""Three-phase cleanup orchestrator per RESOURCE-SAFETY-SPEC Section 3.6."""
+from __future__ import annotations
+
+from ._state import get_session
+from ._types import CleanupAdapter, CleanupResource, CleanupResult
+
+
+def run_cleanup(
+    manifest: list[CleanupResource],
+    adapter: CleanupAdapter | None = None,
+) -> CleanupResult:
+    """Run the three-phase cleanup protocol.
+
+    Args:
+        manifest: Resources to clean up, ordered critical-first.
+        adapter: Cloud adapter implementing delete/describe/tag_scan.
+                 Defaults to AwsCleanupAdapter if None.
+
+    Returns:
+        CleanupResult with status 'clean' or 'dirty', plus any warnings/orphans.
+    """
+    session = get_session()
+    lab_slug = session.lab_id
+
+    if adapter is None:
+        from .adapters.aws import AwsCleanupAdapter
+        adapter = AwsCleanupAdapter()
+
+    credentials = session.credentials
+    warnings: list[str] = []
+    orphans: list[str] = []
+
+    # Phase 1: Teardown
+    for resource in manifest:
+        try:
+            adapter.delete_resource(credentials, resource)
+        except Exception as exc:
+            cmd = resource.cli_delete_command or f"(no CLI command provided for {resource.type} {resource.id!r})"
+            warnings.append(
+                f"FAILED TO DELETE: {resource.type} {resource.id!r}. "
+                f"Error: {exc}. Run manually: {cmd}"
+            )
+
+    # Phase 2: Verification scan
+    for resource in manifest:
+        try:
+            still_exists = adapter.describe_resource(credentials, resource)
+        except Exception as exc:
+            warnings.append(
+                f"VERIFICATION ERROR: Could not check {resource.type} {resource.id!r}. "
+                f"Error: {exc}"
+            )
+            continue
+
+        if still_exists:
+            cmd = resource.cli_delete_command or f"(no CLI command provided for {resource.type} {resource.id!r})"
+            warnings.append(
+                f"WARNING: {resource.type} {resource.id!r} is still alive. "
+                f"Run this to kill it: {cmd}"
+            )
+
+    # Phase 3: Tag audit
+    region = credentials.get("region", "us-east-1")
+    manifest_ids = {r.id for r in manifest}
+    try:
+        tagged_arns = adapter.tag_scan(credentials, lab_slug, region)
+    except Exception as exc:
+        warnings.append(f"TAG AUDIT ERROR: Could not scan tags. Error: {exc}")
+        tagged_arns = []
+
+    for arn in tagged_arns:
+        resource_id = _extract_resource_id_from_arn(arn)
+        if resource_id not in manifest_ids:
+            orphans.append(
+                f"ORPHAN FOUND: {arn} is tagged as belonging to this lab "
+                f"but was not in this session's manifest. It may be left over "
+                f"from a previous attempt."
+            )
+
+    status = "clean" if (not warnings and not orphans) else "dirty"
+    return CleanupResult(status=status, warnings=warnings, orphans=orphans)
+
+
+def _extract_resource_id_from_arn(arn: str) -> str:
+    """Best-effort extraction of a human-readable resource ID from an ARN.
+
+    ARN formats vary by service:
+        arn:aws:s3:::bucket-name          -> bucket-name
+        arn:aws:glue:region:acct:database/db-name  -> db-name
+        arn:aws:kinesis:region:acct:stream/name     -> name
+    Falls back to the full ARN if no pattern matches.
+    """
+    if ":::" in arn:
+        return arn.split(":::")[-1]
+    if "/" in arn:
+        return arn.split("/")[-1]
+    return arn

labrun_checks/_state.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from ._types import CheckpointSession
+
+_session: Optional[CheckpointSession] = None
+
+
+def get_session() -> CheckpointSession:
+    if _session is None:
+        raise RuntimeError(
+            "Call labrun_checks.register(session_token, lab_id) before using check() or cleanup()."
+        )
+    return _session
+
+
+def set_session(session: CheckpointSession) -> None:
+    global _session
+    _session = session
+
+
+def clear() -> None:
+    global _session
+    _session = None

labrun_checks/_types.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal, Optional, Protocol, runtime_checkable
+
+
+Status = Literal["pass", "fail", "error"]
+
+
+@dataclass
+class CheckpointResult:
+    status: Status
+    hint_tier_available: int = 0
+    error_reason: Optional[str] = None
+
+
+@dataclass
+class CheckpointSession:
+    session_token: str
+    lab_id: str
+    credentials: dict = field(default_factory=dict)
+
+
+@dataclass
+class CleanupResource:
+    """Single resource to be cleaned up. Ordering in manifest is critical-first."""
+    type: str
+    id: str
+    region: str = "us-east-1"
+    parent: Optional[str] = None
+    cli_delete_command: Optional[str] = None
+
+
+@dataclass
+class CleanupResult:
+    """Returned by run_cleanup. Notebook passes result.status to cleanup()."""
+    status: Literal["clean", "dirty"]
+    warnings: list[str] = field(default_factory=list)
+    orphans: list[str] = field(default_factory=list)
+
+
+@runtime_checkable
+class CleanupAdapter(Protocol):
+    """Interface that cloud-specific adapters must implement."""
+
+    def delete_resource(self, credentials: dict, resource: CleanupResource) -> None:
+        """Delete a single resource. Raise on failure. Swallow already-gone errors."""
+        ...
+
+    def describe_resource(self, credentials: dict, resource: CleanupResource) -> bool:
+        """Return True if resource still exists, False if confirmed gone."""
+        ...
+
+    def tag_scan(self, credentials: dict, lab_slug: str, region: str) -> list[str]:
+        """Return list of ARNs tagged with labrun:lab-slug=<lab_slug>."""
+        ...

labrun_checks/adapters/__init__.py

@@ -0,0 +1,23 @@
+from .aws import (
+    AwsCleanupAdapter,
+    make_glue_client,
+    make_kinesis_client,
+    make_s3_client,
+    make_sts_client,
+    make_tagging_client,
+)
+from .azure import make_azure_resource_client
+from .databricks import make_databricks_client
+from .snowflake import make_snowflake_connection
+
+__all__ = [
+    "AwsCleanupAdapter",
+    "make_s3_client",
+    "make_glue_client",
+    "make_kinesis_client",
+    "make_sts_client",
+    "make_tagging_client",
+    "make_databricks_client",
+    "make_snowflake_connection",
+    "make_azure_resource_client",
+]

labrun_checks/adapters/aws.py

@@ -0,0 +1,256 @@
+"""AWS boto3 adapter for labrun-checks checkpoints and cleanup."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import boto3
+from botocore.exceptions import ClientError
+
+from .._types import CleanupResource
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3 import S3Client
+
+
+_account_id_cache: Optional[str] = None
+
+
+def _make_client(service: str, credentials: dict):
+    return boto3.client(
+        service,
+        aws_access_key_id=credentials["aws_access_key_id"],
+        aws_secret_access_key=credentials["aws_secret_access_key"],
+        region_name=credentials.get("region", "us-east-1"),
+    )
+
+
+def make_s3_client(credentials: dict) -> "S3Client":
+    return _make_client("s3", credentials)
+
+
+def make_glue_client(credentials: dict):
+    return _make_client("glue", credentials)
+
+
+def make_kinesis_client(credentials: dict):
+    return _make_client("kinesis", credentials)
+
+
+def make_sts_client(credentials: dict):
+    return _make_client("sts", credentials)
+
+
+def make_tagging_client(credentials: dict):
+    return _make_client("resourcegroupstaggingapi", credentials)
+
+
+def make_iam_client(credentials: dict):
+    return _make_client("iam", credentials)
+
+
+def _get_account_id(credentials: dict) -> str:
+    global _account_id_cache
+    if _account_id_cache is None:
+        sts = make_sts_client(credentials)
+        _account_id_cache = sts.get_caller_identity()["Account"]
+    return _account_id_cache
+
+
+def clear_account_id_cache() -> None:
+    global _account_id_cache
+    _account_id_cache = None
+
+
+def _is_already_gone(exc: ClientError) -> bool:
+    code = exc.response.get("Error", {}).get("Code", "")
+    gone_codes = {
+        "NoSuchBucket",
+        "NoSuchKey",
+        "NoSuchEntity",
+        "EntityNotFoundException",
+        "ResourceNotFoundException",
+        "ResourceNotFoundFault",
+        "StreamNotFoundFault",
+    }
+    return code in gone_codes
+
+
+class AwsCleanupAdapter:
+    """Implements CleanupAdapter protocol for AWS resources."""
+
+    def delete_resource(self, credentials: dict, resource: CleanupResource) -> None:
+        try:
+            dispatcher = _DELETE_DISPATCH.get(resource.type)
+            if dispatcher is None:
+                raise ValueError(f"Unknown resource type: {resource.type!r}")
+            dispatcher(credentials, resource)
+        except ClientError as exc:
+            if _is_already_gone(exc):
+                return
+            raise
+
+    def describe_resource(self, credentials: dict, resource: CleanupResource) -> bool:
+        try:
+            dispatcher = _DESCRIBE_DISPATCH.get(resource.type)
+            if dispatcher is None:
+                raise ValueError(f"Unknown resource type: {resource.type!r}")
+            dispatcher(credentials, resource)
+            return True
+        except ClientError as exc:
+            if _is_already_gone(exc):
+                return False
+            raise
+
+    def tag_scan(self, credentials: dict, lab_slug: str, region: str) -> list[str]:
+        creds_with_region = {**credentials, "region": region}
+        client = make_tagging_client(creds_with_region)
+        tag_filters = [{"Key": "labrun:lab-slug", "Values": [lab_slug]}]
+        arns: list[str] = []
+        paginator = client.get_paginator("get_resources")
+        for page in paginator.paginate(TagFilters=tag_filters):
+            for item in page.get("ResourceTagMappingList", []):
+                arn = item.get("ResourceARN", "")
+                if _is_transitional_resource(creds_with_region, arn):
+                    continue
+                arns.append(arn)
+        return arns
+
+
+def _is_transitional_resource(credentials: dict, arn: str) -> bool:
+    """Phase 3 false-positive guard: skip resources in a deleting state."""
+    try:
+        if ":s3:::" in arn:
+            bucket_name = arn.split(":::")[-1]
+            s3 = make_s3_client(credentials)
+            s3.head_bucket(Bucket=bucket_name)
+            return False
+        if ":glue:" in arn and "/database/" in arn:
+            db_name = arn.split("/database/")[-1].split("/")[0]
+            glue = make_glue_client(credentials)
+            glue.get_database(Name=db_name)
+            return False
+    except ClientError as exc:
+        if _is_already_gone(exc):
+            return True
+    return False
+
+
+def _delete_s3_bucket(credentials: dict, resource: CleanupResource) -> None:
+    s3 = make_s3_client(credentials)
+    objects = s3.list_objects_v2(Bucket=resource.id)
+    for obj in objects.get("Contents", []):
+        s3.delete_object(Bucket=resource.id, Key=obj["Key"])
+    s3.delete_bucket(Bucket=resource.id)
+
+
+def _delete_glue_database(credentials: dict, resource: CleanupResource) -> None:
+    glue = make_glue_client(credentials)
+    glue.delete_database(Name=resource.id)
+
+
+def _delete_glue_crawler(credentials: dict, resource: CleanupResource) -> None:
+    glue = make_glue_client(credentials)
+    glue.delete_crawler(Name=resource.id)
+
+
+def _delete_glue_table(credentials: dict, resource: CleanupResource) -> None:
+    glue = make_glue_client(credentials)
+    if resource.parent is None:
+        raise ValueError(f"Glue table {resource.id!r} requires parent (database name)")
+    glue.delete_table(DatabaseName=resource.parent, Name=resource.id)
+
+
+def _delete_kinesis_stream(credentials: dict, resource: CleanupResource) -> None:
+    kinesis = make_kinesis_client(credentials)
+    kinesis.delete_stream(StreamName=resource.id, EnforceConsumerDeletion=True)
+
+
+def _delete_iam_role(credentials: dict, resource: CleanupResource) -> None:
+    iam = make_iam_client(credentials)
+    role_name = resource.id
+
+    try:
+        attached = iam.list_attached_role_policies(RoleName=role_name)
+    except ClientError as exc:
+        if _is_already_gone(exc):
+            return
+        raise
+    for policy in attached.get("AttachedPolicies", []):
+        try:
+            iam.detach_role_policy(RoleName=role_name, PolicyArn=policy["PolicyArn"])
+        except ClientError as exc:
+            if _is_already_gone(exc):
+                return
+            raise
+
+    try:
+        inline = iam.list_role_policies(RoleName=role_name)
+    except ClientError as exc:
+        if _is_already_gone(exc):
+            return
+        raise
+    for policy_name in inline.get("PolicyNames", []):
+        try:
+            iam.delete_role_policy(RoleName=role_name, PolicyName=policy_name)
+        except ClientError as exc:
+            if _is_already_gone(exc):
+                return
+            raise
+
+    try:
+        iam.delete_role(RoleName=role_name)
+    except ClientError as exc:
+        if _is_already_gone(exc):
+            return
+        raise
+
+
+def _describe_s3_bucket(credentials: dict, resource: CleanupResource) -> None:
+    s3 = make_s3_client(credentials)
+    s3.head_bucket(Bucket=resource.id)
+
+
+def _describe_glue_database(credentials: dict, resource: CleanupResource) -> None:
+    glue = make_glue_client(credentials)
+    glue.get_database(Name=resource.id)
+
+
+def _describe_glue_crawler(credentials: dict, resource: CleanupResource) -> None:
+    glue = make_glue_client(credentials)
+    glue.get_crawler(Name=resource.id)
+
+
+def _describe_glue_table(credentials: dict, resource: CleanupResource) -> None:
+    glue = make_glue_client(credentials)
+    if resource.parent is None:
+        raise ValueError(f"Glue table {resource.id!r} requires parent (database name)")
+    glue.get_table(DatabaseName=resource.parent, Name=resource.id)
+
+
+def _describe_kinesis_stream(credentials: dict, resource: CleanupResource) -> None:
+    kinesis = make_kinesis_client(credentials)
+    kinesis.describe_stream(StreamName=resource.id)
+
+
+def _describe_iam_role(credentials: dict, resource: CleanupResource) -> None:
+    iam = make_iam_client(credentials)
+    iam.get_role(RoleName=resource.id)
+
+
+_DELETE_DISPATCH = {
+    "s3_bucket": _delete_s3_bucket,
+    "glue_database": _delete_glue_database,
+    "glue_crawler": _delete_glue_crawler,
+    "glue_table": _delete_glue_table,
+    "kinesis_stream": _delete_kinesis_stream,
+    "iam_role": _delete_iam_role,
+}
+
+_DESCRIBE_DISPATCH = {
+    "s3_bucket": _describe_s3_bucket,
+    "glue_database": _describe_glue_database,
+    "glue_crawler": _describe_glue_crawler,
+    "glue_table": _describe_glue_table,
+    "kinesis_stream": _describe_kinesis_stream,
+    "iam_role": _describe_iam_role,
+}

labrun_checks/adapters/azure.py

@@ -0,0 +1,9 @@
+"""Azure adapter stub for labrun-checks. Not yet implemented."""
+from __future__ import annotations
+
+
+def make_azure_resource_client(credentials: dict):
+    raise NotImplementedError(
+        "The Azure adapter is not yet available in labrun-checks. "
+        "Azure labs ship in a future release. Contact support@labrun.dev if this blocks you."
+    )

labrun_checks/adapters/databricks.py

@@ -0,0 +1,9 @@
+"""Databricks adapter stub for labrun-checks. Not yet implemented."""
+from __future__ import annotations
+
+
+def make_databricks_client(credentials: dict):
+    raise NotImplementedError(
+        "The Databricks adapter is not yet available in labrun-checks. "
+        "Databricks labs ship in a future release. Contact support@labrun.dev if this blocks you."
+    )

labrun_checks/adapters/snowflake.py

@@ -0,0 +1,9 @@
+"""Snowflake adapter stub for labrun-checks. Not yet implemented."""
+from __future__ import annotations
+
+
+def make_snowflake_connection(credentials: dict):
+    raise NotImplementedError(
+        "The Snowflake adapter is not yet available in labrun-checks. "
+        "Snowflake labs ship in a future release. Contact support@labrun.dev if this blocks you."
+    )

labrun_checks-0.2.0.dist-info/METADATA

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: labrun-checks
+Version: 0.2.0
+Summary: Checkpoint validation for Labrun labs
+Requires-Python: >=3.9
+Requires-Dist: requests>=2.28
+Requires-Dist: boto3>=1.26
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: responses>=0.23; extra == "dev"

labrun_checks-0.2.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+labrun_checks/__init__.py,sha256=lvQg4wG6zynSOPqgtKVljeNx6e6cQ54AuYHxr-G7d5c,1886
+labrun_checks/_api.py,sha256=L4EuI6oa1QwEaJUV6oHNdrHTzsRdKSD4rTYdaLxoo_Y,1114
+labrun_checks/_cleanup.py,sha256=6Ym_7nM4z2_-CedGEPeZH_h5RyVYGUcZJKoLxak2vWU,3447
+labrun_checks/_state.py,sha256=mSguVvIa1rrH5yEy8lKLzOOugXeFckmX7ltmRBgt258,535
+labrun_checks/_types.py,sha256=jLHOllr5aHo8t9jpIbAqw7UWHKY9qC00QZzuVNUwRk0,1607
+labrun_checks/adapters/__init__.py,sha256=jUozzi6Ph2z_ZhnHudRsTfuG8sAhURtm8ZtqgRB55Ww,559
+labrun_checks/adapters/aws.py,sha256=sTkw-GnSiGhmbFQbPIvior4YPKq2C2dMvOjzAYQw5eU,8304
+labrun_checks/adapters/azure.py,sha256=T0dp1zuze52bidcpZYOhvjVYNLXnDWninApZBbWN9Hg,352
+labrun_checks/adapters/databricks.py,sha256=fPb7zSdZG4lRMsQSsW6D24eU1FEdgoFmRMnYJERxuc0,363
+labrun_checks/adapters/snowflake.py,sha256=2b6dodrIv88QOVeB1W8Y9Urx6qPEPuG0SOv77bB8S5k,363
+labrun_checks-0.2.0.dist-info/METADATA,sha256=99M744eEDjnGMHO-gVIu3mD8ZIdTcIkuHER4v79a9cc,292
+labrun_checks-0.2.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+labrun_checks-0.2.0.dist-info/top_level.txt,sha256=o5yNRTbV-f95Wz896gNhVv03C9BXr1TV9T53m7su9G8,14
+labrun_checks-0.2.0.dist-info/RECORD,,

labrun_checks-0.2.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

labrun_checks-0.2.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+labrun_checks