diff --git a/src/backend/src/controller/jobs_manager.py b/src/backend/src/controller/jobs_manager.py index 3e6349175..b2e88a0b7 100644 --- a/src/backend/src/controller/jobs_manager.py +++ b/src/backend/src/controller/jobs_manager.py @@ -70,11 +70,22 @@ def list_available_workflows(self) -> List[Dict[str, str]]: return sorted(items, key=lambda x: x["name"].lower()) def install_workflow(self, workflow_id: str, *, job_cluster_id: Optional[str] = None) -> int: + wf_def = self._get_workflow_definition(workflow_id, job_cluster_id=job_cluster_id) + # Deploy workflow code to workspace if deployer is configured if self._workspace_deployer: workflow_dir = self._workflows_root / workflow_id if workflow_dir.exists(): - self._workspace_deployer.deploy_workflow(workflow_id, workflow_dir) + python_package_dir = ( + self._workflows_root.parent + if wf_def.get('deploy_backend_source') + else None + ) + self._workspace_deployer.deploy_workflow( + workflow_id, + workflow_dir, + python_package_dir=python_package_dir, + ) logger.info(f"Deployed workflow '{workflow_id}' to workspace") elif self._settings and self._settings.WORKSPACE_DEPLOYMENT_PATH: logger.error( @@ -83,8 +94,6 @@ def install_workflow(self, workflow_id: str, *, job_cluster_id: Optional[str] = f"The job will be created but will fail at runtime." ) - wf_def = self._get_workflow_definition(workflow_id, job_cluster_id=job_cluster_id) - # Build job settings kwargs from workflow definition tasks = self._build_tasks_from_definition(wf_def) job_settings_kwargs = { @@ -590,6 +599,21 @@ def _get_workflow_definition(self, workflow_id: str, *, job_cluster_id: Optional # Derive from __file__ (works when app runs in workspace) base_path = str(Path(__file__).parent.parent) + if wf.get('deploy_backend_source'): + parameters = wf.setdefault('parameters', {}) + if isinstance(parameters, dict): + if self._settings and self._settings.WORKSPACE_DEPLOYMENT_PATH: + backend_source_path = f"{base_path}/{workflow_id}/backend_src.zip" + else: + # Both non-deployment base paths name the ``src`` package directory: + # WORKSPACE_APP_PATH is configured as ``.../src/backend/src``; the + # fallback derives ``.../backend/src`` from this file at + # ``.../backend/src/controller/jobs_manager.py``. ``from src.*`` + # therefore needs exactly their parent, ``.../src/backend`` or + # ``.../backend``, on sys.path. + backend_source_path = str(Path(base_path).parent) + parameters['backend_source_path'] = backend_source_path + # Helper: detect URI scheme like file:, dbfs:, s3:, etc. def _has_scheme(path: str) -> bool: return bool(re.match(r'^[a-zA-Z][a-zA-Z0-9+\-.]*:', path)) @@ -1299,4 +1323,3 @@ def _create_job_success_notification(self, job_name: str, workflow_id: str, run_ logger.info(f"Triggered {len(executions)} workflow(s) for job success (run {run_id})") except Exception as e: logger.error(f"Failed to trigger workflow for job success: {e}", exc_info=True) - diff --git a/src/backend/src/utils/workspace_deployer.py b/src/backend/src/utils/workspace_deployer.py index 46249fbdf..68af4762c 100644 --- a/src/backend/src/utils/workspace_deployer.py +++ b/src/backend/src/utils/workspace_deployer.py @@ -4,9 +4,10 @@ for use by job tasks. This is necessary for containerized Databricks Apps where job clusters cannot access the container filesystem. """ -from pathlib import Path -from typing import Optional, Set import base64 +from io import BytesIO +from pathlib import Path +from zipfile import ZIP_DEFLATED, ZipFile from databricks.sdk import WorkspaceClient from databricks.sdk.service import workspace @@ -19,7 +20,7 @@ class WorkspaceDeployer: """Deploy workflow code to Databricks workspace.""" - def __init__(self, ws_client: WorkspaceClient, deployment_path: Optional[str] = None): + def __init__(self, ws_client: WorkspaceClient, deployment_path: str | None = None): """Initialize the deployer. Args: @@ -28,9 +29,15 @@ def __init__(self, ws_client: WorkspaceClient, deployment_path: Optional[str] = """ self._client = ws_client self._deployment_path = deployment_path - self._deployed_workflows: Set[str] = set() # Track deployed workflow IDs - - def deploy_workflow(self, workflow_id: str, workflow_dir: Path) -> str: + self._deployed_workflows: set[str] = set() # Track deployed workflow IDs + + def deploy_workflow( + self, + workflow_id: str, + workflow_dir: Path, + *, + python_package_dir: Path | None = None, + ) -> str: """Deploy a workflow folder to the workspace. Copies all files from the local workflow directory to the workspace deployment path. @@ -39,6 +46,7 @@ def deploy_workflow(self, workflow_id: str, workflow_dir: Path) -> str: Args: workflow_id: Unique workflow identifier workflow_dir: Local path to workflow directory + python_package_dir: Optional Python package directory to archive beside the workflow Returns: Workspace path where workflow was deployed @@ -79,6 +87,13 @@ def deploy_workflow(self, workflow_id: str, workflow_dir: Path) -> str: # Recursively upload subdirectories self._upload_directory(file_path, f"{target_path}/{file_path.name}") + if python_package_dir is not None: + self._upload_content( + self._build_python_package_archive(python_package_dir), + f"{target_path}/backend_src.zip", + workspace.ImportFormat.RAW, + ) + # Track successful deployment self._deployed_workflows.add(workflow_id) logger.info(f"Successfully deployed workflow '{workflow_id}' to {target_path}") @@ -121,26 +136,7 @@ def _upload_file(self, local_path: Path, workspace_path: str) -> None: workspace_path: Target workspace path """ try: - # Read file content - with open(local_path, 'rb') as f: - content = f.read() - - # Encode content as base64 - encoded_content = base64.b64encode(content).decode('utf-8') - - # Determine format based on file extension - if local_path.suffix in ['.py', '.yaml', '.yml', '.txt', '.json', '.md']: - format_type = workspace.ImportFormat.AUTO - else: - format_type = workspace.ImportFormat.AUTO - - # Upload using workspace import API - self._client.workspace.import_( - path=workspace_path, - content=encoded_content, - format=format_type, - overwrite=True - ) + self._upload_content(local_path.read_bytes(), workspace_path, workspace.ImportFormat.AUTO) logger.debug(f"Uploaded file: {local_path.name} -> {workspace_path}") @@ -148,6 +144,36 @@ def _upload_file(self, local_path: Path, workspace_path: str) -> None: logger.error(f"Failed to upload file {local_path}: {e}") raise + def _upload_content( + self, + content: bytes, + workspace_path: str, + format_type: workspace.ImportFormat, + ) -> None: + encoded_content = base64.b64encode(content).decode("utf-8") + self._client.workspace.import_( + path=workspace_path, + content=encoded_content, + format=format_type, + overwrite=True, + ) + + @staticmethod + def _build_python_package_archive(package_dir: Path) -> bytes: + if not package_dir.exists() or not package_dir.is_dir(): + raise ValueError(f"Python package directory does not exist: {package_dir}") + + archive_buffer = BytesIO() + with ZipFile(archive_buffer, "w", ZIP_DEFLATED) as archive: + for file_path in package_dir.rglob("*"): + if not file_path.is_file(): + continue + relative_path = file_path.relative_to(package_dir) + if "__pycache__" in relative_path.parts or file_path.suffix in {".pyc", ".pyo"}: + continue + archive.write(file_path, Path("src") / relative_path) + return archive_buffer.getvalue() + def _upload_directory(self, local_dir: Path, workspace_path: str) -> None: """Recursively upload a directory to the workspace. diff --git a/src/backend/src/workflows/compliance_checks/compliance_checks.py b/src/backend/src/workflows/compliance_checks/compliance_checks.py index 420a14f33..9dafbe3f0 100644 --- a/src/backend/src/workflows/compliance_checks/compliance_checks.py +++ b/src/backend/src/workflows/compliance_checks/compliance_checks.py @@ -9,13 +9,27 @@ import sys import json import argparse +import shutil +import tempfile +from io import BytesIO from typing import Any, Dict, List, Optional, Tuple from datetime import datetime from uuid import uuid4 from pathlib import Path +from zipfile import BadZipFile, ZipFile -# Add parent directory to path to enable imports from src.* -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) +# Bootstrap the backend package parent for local and file-based execution. The +# serverless runner may execute this source without binding ``__file__``; main() +# handles that case using the deployed backend source archive path. +_entry_file = globals().get("__file__") +if _entry_file is not None: + try: + sys.path.insert(0, str(Path(_entry_file).parent.parent.parent.parent)) + except Exception as _exc: # pragma: no cover - defensive; must never abort import + # Narrowly guarded so a surprising failure is diagnosable but never + # crashes module load. Use print because the logging stack isn't wired + # up at this point (matches the rest of this script's stdout logging). + print(f"WARNING: compliance_checks sys.path bootstrap failed: {_exc}", file=sys.stderr) from sqlalchemy import text, create_engine from sqlalchemy.engine import Engine @@ -24,6 +38,61 @@ from databricks.sdk import WorkspaceClient +_extracted_backend_sources: Dict[str, str] = {} + + +def _add_backend_source_path(backend_source_path: Optional[str]) -> Optional[str]: + if not backend_source_path: + return None + + source_path = Path(backend_source_path) + cached_path = _extracted_backend_sources.get(str(source_path)) + if cached_path and Path(cached_path).is_dir(): + import_path = cached_path + source_description = f"cached extraction of {source_path}" + elif source_path.is_dir(): + import_path = str(source_path) + source_description = "source directory" + elif source_path.suffix.lower() == ".zip": + if not source_path.is_file(): + print(f"WARNING: backend source archive not found: {source_path}", file=sys.stderr) + return None + + extracted_path = tempfile.mkdtemp(prefix="ontos-compliance-backend-") + try: + with source_path.open("rb") as archive_file: + archive_bytes = archive_file.read() + with ZipFile(BytesIO(archive_bytes)) as archive: + root = Path(extracted_path).resolve() + for member in archive.infolist(): + member_dest = (root / member.filename).resolve() + if not member_dest.is_relative_to(root): + raise RuntimeError( + f"unsafe zip member escapes extraction root: {member.filename}" + ) + archive.extractall(extracted_path) + if not (Path(extracted_path) / "src" / "__init__.py").is_file(): + raise RuntimeError("archive does not contain the src package") + except (BadZipFile, OSError, RuntimeError) as exc: + shutil.rmtree(extracted_path, ignore_errors=True) + raise RuntimeError( + f"Failed to extract backend source archive {source_path}: {exc}" + ) from exc + + import_path = extracted_path + _extracted_backend_sources[str(source_path)] = import_path + source_description = f"local extraction of {source_path}" + else: + print(f"WARNING: unsupported backend source path: {source_path}", file=sys.stderr) + return None + + while import_path in sys.path: + sys.path.remove(import_path) + sys.path.insert(0, import_path) + print(f"Backend source added to sys.path: {import_path} ({source_description})") + return import_path + + # ============================================================================ # OAuth Token Generation & Database Connection (for Lakebase Postgres) # ============================================================================ @@ -269,6 +338,7 @@ def main() -> None: parser.add_argument("--policy_severities", type=str, default=None) # JSON array of severities parser.add_argument("--entity_limit", type=str, default=None) # Limit entities checked per policy parser.add_argument("--verbose", type=str, default="false") + parser.add_argument("--backend_source_path", type=str, default=None) # Database connection parameters parser.add_argument("--lakebase_instance_name", type=str, required=True) @@ -282,6 +352,8 @@ def main() -> None: args, _ = parser.parse_known_args() + _add_backend_source_path(args.backend_source_path) + # Parse arguments policy_filter = args.policy_filter verbose = args.verbose.lower() == "true" diff --git a/src/backend/src/workflows/compliance_checks/compliance_checks.yaml b/src/backend/src/workflows/compliance_checks/compliance_checks.yaml index 147e10639..ea597f85e 100644 --- a/src/backend/src/workflows/compliance_checks/compliance_checks.yaml +++ b/src/backend/src/workflows/compliance_checks/compliance_checks.yaml @@ -1,11 +1,19 @@ name: "Compliance Checks" description: "Run user-defined compliance DSL rules from compliance_policies table" format: "MULTI_TASK" +deploy_backend_source: true environments: - environment_key: "compliance-checks-env" spec: dependencies: - "databricks-sdk>=0.81.0" + - "databricks-sql-connector" + - "alembic" + - "fastapi" + - "gitpython" + - "pydantic[email]" + - "pydantic-settings" + - "pyyaml" - "sqlalchemy" - "psycopg2-binary" tasks: @@ -23,6 +31,8 @@ tasks: - "{{job.parameters.entity_limit}}" - "--verbose" - "{{job.parameters.verbose}}" + - "--backend_source_path" + - "{{job.parameters.backend_source_path}}" - "--lakebase_instance_name" - "{{job.parameters.lakebase_instance_name}}" - "--postgres_host" @@ -51,6 +61,7 @@ parameters: policy_severities: "null" entity_limit: "null" verbose: "false" + backend_source_path: "" lakebase_instance_name: "" postgres_host: "" postgres_db: "" diff --git a/src/backend/tests/test_compliance_checks_workflow.py b/src/backend/tests/test_compliance_checks_workflow.py new file mode 100644 index 000000000..d70103aea --- /dev/null +++ b/src/backend/tests/test_compliance_checks_workflow.py @@ -0,0 +1,292 @@ +"""Serverless deployment tests for the compliance-checks workflow.""" + +import base64 +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from zipfile import ZipFile + +import pytest +from databricks.sdk.service import workspace + +from src.controller.jobs_manager import JobsManager +from src.utils.workspace_deployer import WorkspaceDeployer + +_MODULE_PATH = ( + Path(__file__).resolve().parent.parent + / "src" + / "workflows" + / "compliance_checks" + / "compliance_checks.py" +) +_BACKEND_PACKAGE_DIR = _MODULE_PATH.parent.parent.parent +_EXPECTED_SRC_ROOT = str(_BACKEND_PACKAGE_DIR.parent) + + +def _read_source() -> str: + return _MODULE_PATH.read_text(encoding="utf-8") + + +@pytest.fixture(autouse=True) +def restore_import_state(): + original_path = list(sys.path) + original_src_modules = { + name: module + for name, module in sys.modules.items() + if name == "src" or name.startswith("src.") + } + try: + yield + finally: + sys.path[:] = original_path + for name in list(sys.modules): + if name == "src" or name.startswith("src."): + sys.modules.pop(name, None) + sys.modules.update(original_src_modules) + + +@pytest.fixture +def stub_third_party(monkeypatch): + def ensure_module(name: str, attrs: dict) -> None: + try: + __import__(name) + return + except Exception: + pass + parts = name.split(".") + for index in range(1, len(parts)): + parent = ".".join(parts[:index]) + if parent not in sys.modules: + package = types.ModuleType(parent) + package.__path__ = [] + monkeypatch.setitem(sys.modules, parent, package) + module = types.ModuleType(name) + for attr, value in attrs.items(): + setattr(module, attr, value) + monkeypatch.setitem(sys.modules, name, module) + + ensure_module("sqlalchemy", {"text": object, "create_engine": object}) + ensure_module("sqlalchemy.engine", {"Engine": object}) + ensure_module("sqlalchemy.orm", {"Session": object, "sessionmaker": object}) + ensure_module("databricks.sdk", {"WorkspaceClient": object}) + + +def test_serverless_load_without_file_inserts_no_fabricated_path( + stub_third_party, monkeypatch, tmp_path +): + code = compile(_read_source(), "", "exec") + workflow_dir = tmp_path / "workflows" / "compliance_checks" + workflow_dir.mkdir(parents=True) + monkeypatch.chdir(workflow_dir) + cwd_grandparent = str(Path.cwd().parent.parent) + + module_globals = {"__name__": "not_main"} + sys_path_before = list(sys.path) + exec(code, module_globals) # noqa: S102 - intentionally simulates serverless execution + + added_paths = [path for path in sys.path if path not in sys_path_before] + assert "main" in module_globals + assert not added_paths + assert cwd_grandparent not in sys.path + assert "__file__" not in module_globals + + +def test_normal_load_with_file_inserts_backend_package_parent(stub_third_party): + code = compile(_read_source(), str(_MODULE_PATH), "exec") + module_globals = {"__name__": "not_main", "__file__": str(_MODULE_PATH)} + sys.path[:] = [path for path in sys.path if path != _EXPECTED_SRC_ROOT] + + exec(code, module_globals) # noqa: S102 - intentionally executes the entry source + + assert "main" in module_globals + assert sys.path[0] == _EXPECTED_SRC_ROOT + + +def test_serverless_main_imports_backend_modules_from_deployed_archive( + monkeypatch, tmp_path +): + archive_path = tmp_path / "backend_src.zip" + extracted_path = tmp_path / "extracted_backend" + archive_path.write_bytes( + WorkspaceDeployer._build_python_package_archive(_BACKEND_PACKAGE_DIR) + ) + + for name in list(sys.modules): + if name == "src" or name.startswith("src."): + sys.modules.pop(name, None) + blocked_roots = {_EXPECTED_SRC_ROOT, str(_BACKEND_PACKAGE_DIR)} + sys.path[:] = [path for path in sys.path if path not in blocked_roots] + + module_globals = {"__name__": "not_main"} + exec( # noqa: S102 - intentionally simulates serverless execution + compile(_read_source(), "", "exec"), + module_globals, + ) + + def make_extract_dir(prefix): + extracted_path.mkdir() + return str(extracted_path) + + module_globals["tempfile"] = SimpleNamespace(mkdtemp=make_extract_dir) + + class FakeSession: + def get(self, model, policy_id): + return None + + def close(self): + pass + + fake_session = FakeSession() + module_globals["WorkspaceClient"] = lambda **kwargs: object() + module_globals["create_engine_from_params"] = lambda **kwargs: object() + module_globals["sessionmaker"] = lambda **kwargs: lambda: fake_session + module_globals["load_policies"] = lambda *args, **kwargs: [ + { + "id": "policy-1", + "name": "Hermetic policy", + "category": "Governance", + "severity": "low", + } + ] + monkeypatch.setattr( + sys, + "argv", + [ + "compliance_checks.py", + "--backend_source_path", + str(archive_path), + "--lakebase_instance_name", + "instance", + "--postgres_host", + "host", + "--postgres_db", + "database", + ], + ) + + module_globals["main"]() + + assert sys.path[0] == str(extracted_path) + assert str(archive_path) not in sys.path + assert _EXPECTED_SRC_ROOT not in sys.path + assert str(_BACKEND_PACKAGE_DIR) not in sys.path + compliance_manager_file = Path(sys.modules["src.controller.compliance_manager"].__file__) + compliance_model_file = Path(sys.modules["src.db_models.compliance"].__file__) + assert compliance_manager_file.is_relative_to(extracted_path) + assert compliance_model_file.is_relative_to(extracted_path) + assert module_globals["_add_backend_source_path"](str(archive_path)) == str(extracted_path) + + +def test_backend_source_directory_is_added_without_extraction(tmp_path): + module_globals = {"__name__": "not_main"} + exec( # noqa: S102 - intentionally executes the workflow entry source + compile(_read_source(), "", "exec"), + module_globals, + ) + source_directory = tmp_path / "backend" + source_directory.mkdir() + + added_path = module_globals["_add_backend_source_path"](str(source_directory)) + + assert added_path == str(source_directory) + assert sys.path[0] == str(source_directory) + + +def test_workspace_deployer_uploads_backend_source_archive(tmp_path): + workflow_dir = tmp_path / "compliance_checks" + workflow_dir.mkdir() + (workflow_dir / "compliance_checks.py").write_text("print('ok')\n", encoding="utf-8") + package_dir = tmp_path / "package" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + (package_dir / "module.py").write_text("VALUE = 1\n", encoding="utf-8") + + uploads = {} + + class FakeWorkspaceApi: + def get_status(self, path): + return object() + + def mkdirs(self, path): + pass + + def import_(self, *, path, content, format, overwrite): + uploads[path] = (base64.b64decode(content), format, overwrite) + + deployer = WorkspaceDeployer( + SimpleNamespace(workspace=FakeWorkspaceApi()), + "/Workspace/Shared/ontos-workflows", + ) + target_path = deployer.deploy_workflow( + "compliance_checks", + workflow_dir, + python_package_dir=package_dir, + ) + + archive_bytes, format_type, overwrite = uploads[f"{target_path}/backend_src.zip"] + archive_path = tmp_path / "uploaded.zip" + archive_path.write_bytes(archive_bytes) + with ZipFile(archive_path) as archive: + assert set(archive.namelist()) == {"src/__init__.py", "src/module.py"} + assert format_type == workspace.ImportFormat.RAW + assert overwrite is True + + +def test_workflow_definition_points_serverless_job_at_deployed_archive(): + settings = SimpleNamespace( + WORKSPACE_DEPLOYMENT_PATH="/Workspace/Shared/ontos-workflows", + WORKSPACE_APP_PATH=None, + ) + manager = JobsManager( + db=object(), + ws_client=object(), + settings=settings, + workflows_root=_BACKEND_PACKAGE_DIR / "workflows", + ) + + definition = manager._get_workflow_definition( + "compliance_checks", + job_cluster_id=None, + ) + + assert definition["parameters"]["backend_source_path"] == ( + "/Workspace/Shared/ontos-workflows/compliance_checks/backend_src.zip" + ) + task_parameters = definition["tasks"][0]["spark_python_task"]["parameters"] + backend_path_index = task_parameters.index("--backend_source_path") + assert task_parameters[backend_path_index + 1] == "{{job.parameters.backend_source_path}}" + + +@pytest.mark.parametrize( + ("workspace_app_path", "expected_backend_source_path"), + [ + ( + "/Workspace/Users/user@example.com/ontos/src/backend/src", + "/Workspace/Users/user@example.com/ontos/src/backend", + ), + (None, str(_BACKEND_PACKAGE_DIR.parent)), + ], +) +def test_workflow_definition_uses_backend_package_parent_without_deployer( + workspace_app_path, + expected_backend_source_path, +): + settings = SimpleNamespace( + WORKSPACE_DEPLOYMENT_PATH=None, + WORKSPACE_APP_PATH=workspace_app_path, + ) + manager = JobsManager( + db=object(), + ws_client=object(), + settings=settings, + workflows_root=_BACKEND_PACKAGE_DIR / "workflows", + ) + + definition = manager._get_workflow_definition( + "compliance_checks", + job_cluster_id=None, + ) + + assert definition["parameters"]["backend_source_path"] == expected_backend_source_path + assert Path(expected_backend_source_path).name == "backend"