deploy.py

# Copyright (c) 2025, Salesforce, Inc.
# SPDX-License-Identifier: Apache-2
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

from html import unescape
import json
import os
import shutil
import tempfile
import time
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Union,
)

from loguru import logger
import pydantic
from pydantic import BaseModel
import requests

from datacustomcode.cmd import cmd_output

if TYPE_CHECKING:
    from datacustomcode.credentials import Credentials

DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms"
AUTH_PATH = "services/oauth2/token"
WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000

# Available compute types for Data Cloud deployments.
# Nomenclature used by COMPUTE_TYPES keys align with
# compute instances provisioned by Data Cloud.
COMPUTE_TYPES = {
    "CPU_L": "CPU_XS",  # Large CPU instance
    "CPU_XL": "CPU_S",  # X-Large CPU instance
    "CPU_2XL": "CPU_M",  # 2X-Large CPU instance (default)
    "CPU_4XL": "CPU_L",  # 4X-Large CPU instance
}


class TransformationJobMetadata(BaseModel):
    name: str
    version: str
    description: str
    computeType: str

    def __init__(self, **data):
        super().__init__(**data)


def _join_strip_url(*args: str) -> str:
    return "/".join(arg.strip("/") for arg in args)


JSONValue = Union[
    Dict[str, "JSONValue"], List["JSONValue"], str, int, float, bool, None
]


def _make_api_call(
    url: str,
    method: str,
    headers: Union[dict, None] = None,
    token: Union[str, None] = None,
    **kwargs,
) -> dict[str, JSONValue]:
    """Make a request to Data Cloud Custom Code API."""
    headers = headers or {}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    logger.debug(f"Making API call: {method} {url}")
    logger.debug(f"Headers: {headers}")
    logger.debug(f"Request params: {kwargs}")

    response = requests.request(method=method, url=url, headers=headers, **kwargs)
    json_response = response.json()
    if response.status_code >= 400:
        logger.debug(f"Error Response: {json_response}")
    response.raise_for_status()
    assert isinstance(
        json_response, dict
    ), f"Unexpected response type: {type(json_response)}"
    return json_response


class AccessTokenResponse(BaseModel):
    access_token: str
    instance_url: str


def _retrieve_access_token(credentials: Credentials) -> AccessTokenResponse:
    """Get a token for the Salesforce API."""
    logger.debug("Getting oauth token...")

    url = f"{credentials.login_url.rstrip('/')}/{AUTH_PATH.lstrip('/')}"

    data = {
        "grant_type": "password",
        "username": credentials.username,
        "password": credentials.password,
        "client_id": credentials.client_id,
        "client_secret": credentials.client_secret,
    }
    response = _make_api_call(url, "POST", data=data)
    return AccessTokenResponse(**response)


class CreateDeploymentResponse(BaseModel):
    fileUploadUrl: str


def create_deployment(
    access_token: AccessTokenResponse, metadata: TransformationJobMetadata
) -> CreateDeploymentResponse:
    """Create a custom code deployment in the DataCloud."""
    url = _join_strip_url(access_token.instance_url, DATA_CUSTOM_CODE_PATH)
    body = {
        "label": metadata.name,
        "name": metadata.name,
        "description": metadata.description,
        "version": metadata.version,
        "computeType": metadata.computeType,
    }
    logger.debug(f"Creating deployment {metadata.name}...")
    try:
        response = _make_api_call(
            url, "POST", token=access_token.access_token, json=body
        )
        return CreateDeploymentResponse(**response)
    except requests.HTTPError as exc:
        if exc.response.status_code == 409:
            raise ValueError(
                f"Deployment {metadata.name} exists. Please use a different name."
            ) from exc
        raise


PLATFORM_ENV_VAR = "DOCKER_DEFAULT_PLATFORM=linux/amd64"
DOCKER_IMAGE_NAME = "datacloud-custom-code-dependency-builder"
DEPENDENCIES_ARCHIVE_NAME = "native_dependencies"
DEPENDENCIES_ARCHIVE_FULL_NAME = f"{DEPENDENCIES_ARCHIVE_NAME}.tar.gz"
DEPENDENCIES_ARCHIVE_PATH = os.path.join(
    "payload", "archives", DEPENDENCIES_ARCHIVE_FULL_NAME
)
ZIP_FILE_NAME = "deployment.zip"


def prepare_dependency_archive(directory: str) -> None:
    cmd = f"docker images -q {DOCKER_IMAGE_NAME}"
    image_exists = cmd_output(cmd)

    if not image_exists:
        logger.info("Building docker image...")
        cmd = (
            f"{PLATFORM_ENV_VAR} docker build -t {DOCKER_IMAGE_NAME} "
            f"-f Dockerfile.dependencies ."
        )
        cmd_output(cmd)

    with tempfile.TemporaryDirectory() as temp_dir:
        logger.info("Building dependencies archive")
        shutil.copy("requirements.txt", temp_dir)
        shutil.copy("build_native_dependencies.sh", temp_dir)
        cmd = (
            f"{PLATFORM_ENV_VAR} docker run --rm "
            f"-v {temp_dir}:/workspace "
            f"{DOCKER_IMAGE_NAME}"
        )
        cmd_output(cmd)
        archives_temp_path = os.path.join(temp_dir, DEPENDENCIES_ARCHIVE_FULL_NAME)
        os.makedirs(os.path.dirname(DEPENDENCIES_ARCHIVE_PATH), exist_ok=True)
        shutil.copy(archives_temp_path, DEPENDENCIES_ARCHIVE_PATH)

        logger.info(f"Dependencies archived to {DEPENDENCIES_ARCHIVE_PATH}")


class DeploymentsResponse(BaseModel):
    deploymentStatus: str


def get_deployments(
    access_token: AccessTokenResponse, metadata: TransformationJobMetadata
) -> DeploymentsResponse:
    """Get all custom code deployments from the DataCloud."""
    url = _join_strip_url(
        access_token.instance_url, DATA_CUSTOM_CODE_PATH, metadata.name
    )
    response = _make_api_call(url, "GET", token=access_token.access_token)
    return DeploymentsResponse(**response)


def wait_for_deployment(
    access_token: AccessTokenResponse,
    metadata: TransformationJobMetadata,
    callback: Union[Callable[[str], None], None] = None,
) -> None:
    """Wait for deployment to complete.

    Args:
        callback: Optional callback function that receives the deployment status
    """
    start_time = time.time()
    logger.debug("Waiting for deployment to complete")

    while True:
        deployment_status = get_deployments(access_token, metadata)
        status = deployment_status.deploymentStatus
        if (time.time() - start_time) > WAIT_FOR_DEPLOYMENT_TIMEOUT:
            raise TimeoutError("Deployment timed out.")

        if callback:
            callback(status)
        if status == "Deployed":
            logger.debug(
                f"Deployment completed.\nElapsed time: {time.time() - start_time}"
            )
            break
        time.sleep(1)


DATA_TRANSFORM_REQUEST_TEMPLATE: dict[str, Any] = {
    "nodes": {},
    "sources": {},
    "macros": {
        "macro.byoc": {
            "arguments": [{"name": "{SCRIPT_NAME}", "type": "BYOC_SCRIPT"}],
        }
    },
}


class DataTransformConfig(BaseModel):
    sdkVersion: str
    entryPoint: str
    dataspace: str
    permissions: Permissions


class Permissions(BaseModel):
    read: Union[DloPermission]
    write: Union[DloPermission]


class DloPermission(BaseModel):
    dlo: list[str]


def get_data_transform_config(directory: str) -> DataTransformConfig:
    """Get the data transform config from the config.json file."""
    config_path = os.path.join(directory, "config.json")
    try:
        with open(config_path, "r") as f:
            config = json.loads(f.read())
            return DataTransformConfig(**config)
    except FileNotFoundError as err:
        raise FileNotFoundError(f"config.json not found at {config_path}") from err
    except json.JSONDecodeError as err:
        raise ValueError(f"config.json at {config_path} is not valid JSON") from err
    except pydantic.ValidationError as err:
        missing_fields = [str(err["loc"][0]) for err in err.errors()]
        raise ValueError(
            f"config.json at {config_path} is missing required "
            f"fields: {', '.join(missing_fields)}"
        ) from err


def verify_data_transform_config(directory: str) -> None:
    """Verify the data transform config.json contents."""
    get_data_transform_config(directory)
    logger.debug(f"Verified data transform config in {directory}")


def create_data_transform(
    directory: str,
    access_token: AccessTokenResponse,
    metadata: TransformationJobMetadata,
) -> dict:
    """Create a data transform in the DataCloud."""
    script_name = metadata.name
    data_transform_config = get_data_transform_config(directory)
    request_hydrated = DATA_TRANSFORM_REQUEST_TEMPLATE.copy()

    # Add nodes for each write DLO
    for i, dlo in enumerate(data_transform_config.permissions.write.dlo, 1):
        request_hydrated["nodes"][f"node{i}"] = {
            "relation_name": dlo,
            "config": {"materialized": "table"},
            "compiled_code": "",
        }

    # Add sources for each read DLO
    for i, dlo in enumerate(data_transform_config.permissions.read.dlo, 1):
        request_hydrated["sources"][f"source{i}"] = {"relation_name": dlo}

    request_hydrated["macros"]["macro.byoc"]["arguments"][0]["name"] = script_name

    body = {
        "definition": {
            "type": "DCSQL",
            "manifest": request_hydrated,
            "version": "56.0",
        },
        "label": f"{metadata.name}",
        "name": f"{metadata.name}",
        "type": "BATCH",
        "dataSpaceName": data_transform_config.dataspace,
    }

    url = _join_strip_url(access_token.instance_url, DATA_TRANSFORMS_PATH)
    response = _make_api_call(url, "POST", token=access_token.access_token, json=body)
    return response


def has_nonempty_requirements_file(directory: str) -> bool:
    """
    Check if requirements.txt exists in the given directory and has at least
    one non-comment line.
    Args:
        directory (str): The directory to check for requirements.txt.
    Returns:
        bool: True if requirements.txt exists and has a non-comment line,
        False otherwise.
    """
    # Look for requirements.txt in the parent directory of the given directory
    requirements_path = os.path.join(os.path.dirname(directory), "requirements.txt")

    try:
        if os.path.isfile(requirements_path):
            with open(requirements_path, "r", encoding="utf-8") as f:
                for line in f:
                    # Consider non-empty if any line is not a comment (ignoring
                    # leading whitespace)
                    if line.strip() and not line.lstrip().startswith("#"):
                        return True
    except Exception as e:
        logger.error(f"Error reading requirements.txt: {e}")
    return False


def upload_zip(file_upload_url: str) -> None:
    file_upload_url = unescape(file_upload_url)
    with open(ZIP_FILE_NAME, "rb") as zip_file:
        response = requests.put(
            file_upload_url, data=zip_file, headers={"Content-Type": "application/zip"}
        )
        response.raise_for_status()


def zip(
    directory: str,
):
    # Create a zip file excluding .DS_Store files
    import zipfile

    # prepare payload only if requirements.txt is non-empty
    if has_nonempty_requirements_file(directory):
        prepare_dependency_archive(directory)
    else:
        logger.info(
            f"Skipping dependency archive: requirements.txt is missing or empty "
            f"in {directory}"
        )

    logger.debug(f"Zipping directory... {directory}")

    with zipfile.ZipFile(ZIP_FILE_NAME, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(directory):
            # Skip .DS_Store files when adding to zip
            for file in files:
                if file != ".DS_Store":
                    file_path = os.path.join(root, file)
                    zipf.write(file_path)

    logger.debug(f"Created zip file: {ZIP_FILE_NAME}")


def deploy_full(
    directory: str,
    metadata: TransformationJobMetadata,
    credentials: Credentials,
    callback=None,
) -> AccessTokenResponse:
    """Deploy a data transform in the DataCloud."""
    access_token = _retrieve_access_token(credentials)

    # prepare payload
    verify_data_transform_config(directory)

    # create deployment and upload payload
    deployment = create_deployment(access_token, metadata)
    zip(directory)
    upload_zip(deployment.fileUploadUrl)
    wait_for_deployment(access_token, metadata, callback)

    # create data transform
    create_data_transform(directory, access_token, metadata)
    return access_token


def run_data_transform(
    access_token: AccessTokenResponse, metadata: TransformationJobMetadata
) -> dict:
    logger.debug(f"Triggering data transform {metadata.name}")
    url = _join_strip_url(
        access_token.instance_url, DATA_TRANSFORMS_PATH, metadata.name, "actions", "run"
    )
    return _make_api_call(url, "POST")