diff --git a/README.md b/README.md index aa2a313..e30665d 100644 --- a/README.md +++ b/README.md @@ -217,6 +217,7 @@ Zip a transformation job in preparation to upload to Data Cloud. Options: - `--path TEXT`: Path to the code directory (default: ".") +- `--network TEXT`: docker network (default: "default") #### `datacustomcode deploy` @@ -228,6 +229,7 @@ Options: - `--name TEXT`: Name of the transformation job [required] - `--version TEXT`: Version of the transformation job (default: "0.0.1") - `--description TEXT`: Description of the transformation job (default: "") +- `--network TEXT`: docker network (default: "default") - `--cpu-size TEXT`: CPU size for the deployment (default: "CPU_XL"). Available options: CPU_L(Large), CPU_XL(Extra Large), CPU_2XL(2X Large), CPU_4XL(4X Large) diff --git a/src/datacustomcode/cli.py b/src/datacustomcode/cli.py index 437674a..0876ce3 100644 --- a/src/datacustomcode/cli.py +++ b/src/datacustomcode/cli.py @@ -71,11 +71,12 @@ def configure( @cli.command() @click.argument("path", default="payload") -def zip(path: str): +@click.option("--network", default="default") +def zip(path: str, network: str): from datacustomcode.deploy import zip logger.debug("Zipping project") - zip(path) + zip(path, network) @cli.command() @@ -84,6 +85,7 @@ def zip(path: str): @click.option("--version", default="0.0.1") @click.option("--description", default="Custom Data Transform Code") @click.option("--profile", default="default") +@click.option("--network", default="default") @click.option( "--cpu-size", default="CPU_2XL", @@ -98,7 +100,13 @@ def zip(path: str): Choose based on your workload requirements.""", ) def deploy( - path: str, name: str, version: str, description: str, cpu_size: str, profile: str + path: str, + name: str, + version: str, + description: str, + cpu_size: str, + profile: str, + network: str, ): from datacustomcode.credentials import Credentials from datacustomcode.deploy import TransformationJobMetadata, deploy_full @@ -132,7 +140,7 @@ def deploy( fg="red", ) raise click.Abort() from None - deploy_full(path, metadata, credentials) + deploy_full(path, metadata, credentials, network) @cli.command() diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index 46109db..869803e 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -163,27 +163,22 @@ def create_deployment( ZIP_FILE_NAME = "deployment.zip" -def prepare_dependency_archive(directory: str) -> None: +def prepare_dependency_archive(directory: str, docker_network: 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 ." - ) + logger.info(f"Building docker image with docker network: {docker_network}...") + cmd = docker_build_cmd(docker_network) cmd_output(cmd) with tempfile.TemporaryDirectory() as temp_dir: - logger.info("Building dependencies archive") + logger.info( + f"Building dependencies archive with docker network: {docker_network}" + ) 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 = docker_run_cmd(docker_network, temp_dir) 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) @@ -192,6 +187,31 @@ def prepare_dependency_archive(directory: str) -> None: logger.info(f"Dependencies archived to {DEPENDENCIES_ARCHIVE_PATH}") +def docker_build_cmd(network: str) -> str: + cmd = ( + f"{PLATFORM_ENV_VAR} docker build -t {DOCKER_IMAGE_NAME} " + f"--file Dockerfile.dependencies . " + ) + + if network != "default": + cmd = cmd + f"--network {network}" + logger.debug(f"Docker build command: {cmd}") + return cmd + + +def docker_run_cmd(network: str, temp_dir) -> str: + cmd = ( + f"{PLATFORM_ENV_VAR} docker run --rm " + f"-v {temp_dir}:/workspace " + f"{DOCKER_IMAGE_NAME} " + ) + + if network != "default": + cmd = cmd + f"--network {network} " + logger.debug(f"Docker run command: {cmd}") + return cmd + + class DeploymentsResponse(BaseModel): deploymentStatus: str @@ -366,13 +386,14 @@ def upload_zip(file_upload_url: str) -> None: def zip( directory: str, + docker_network: 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) + prepare_dependency_archive(directory, docker_network) else: logger.info( f"Skipping dependency archive: requirements.txt is missing or empty " @@ -396,6 +417,7 @@ def deploy_full( directory: str, metadata: TransformationJobMetadata, credentials: Credentials, + docker_network: str, callback=None, ) -> AccessTokenResponse: """Deploy a data transform in the DataCloud.""" @@ -406,7 +428,7 @@ def deploy_full( # create deployment and upload payload deployment = create_deployment(access_token, metadata) - zip(directory) + zip(directory, docker_network) upload_zip(deployment.fileUploadUrl) wait_for_deployment(access_token, metadata, callback) diff --git a/tests/test_deploy.py b/tests/test_deploy.py index e4d524d..8fe3219 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -2,7 +2,6 @@ from unittest.mock import ( MagicMock, - call, mock_open, patch, ) @@ -46,12 +45,12 @@ class TestPrepareDependencyArchive: ) EXPECTED_BUILD_CMD = ( "DOCKER_DEFAULT_PLATFORM=linux/amd64 docker build " - "-t datacloud-custom-code-dependency-builder -f Dockerfile.dependencies ." + "-t datacloud-custom-code-dependency-builder -f Dockerfile.dependencies . " ) EXPECTED_DOCKER_RUN_CMD = ( "DOCKER_DEFAULT_PLATFORM=linux/amd64 docker run --rm " "-v /tmp/test_dir:/workspace " - "datacloud-custom-code-dependency-builder" + "datacloud-custom-code-dependency-builder " ) @patch("datacustomcode.deploy.cmd_output") @@ -59,8 +58,17 @@ class TestPrepareDependencyArchive: @patch("datacustomcode.deploy.tempfile.TemporaryDirectory") @patch("datacustomcode.deploy.os.path.join") @patch("datacustomcode.deploy.os.makedirs") + @patch("datacustomcode.deploy.docker_build_cmd") + @patch("datacustomcode.deploy.docker_run_cmd") def test_prepare_dependency_archive_image_exists( - self, mock_makedirs, mock_join, mock_temp_dir, mock_copy, mock_cmd_output + self, + mock_docker_run_cmd, + mock_docker_build_cmd, + mock_makedirs, + mock_join, + mock_temp_dir, + mock_copy, + mock_cmd_output, ): """Test prepare_dependency_archive when Docker image already exists.""" # Mock the temporary directory context manager @@ -75,20 +83,25 @@ def test_prepare_dependency_archive_image_exists( # Mock os.path.join for archive path mock_join.return_value = "/tmp/test_dir/native_dependencies.tar.gz" - prepare_dependency_archive("/test/dir") + # Mock the docker command functions + mock_docker_build_cmd.return_value = "mock build command" + mock_docker_run_cmd.return_value = "mock run command" + + prepare_dependency_archive("/test/dir", "default") # Verify docker images command was called mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_IMAGES_CMD) # Verify docker build command was not called (since image already exists) - assert call(self.EXPECTED_BUILD_CMD) not in mock_cmd_output.call_args_list + mock_docker_build_cmd.assert_not_called() # Verify files were copied to temp directory mock_copy.assert_any_call("requirements.txt", "/tmp/test_dir") mock_copy.assert_any_call("build_native_dependencies.sh", "/tmp/test_dir") # Verify docker run command was called - mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_RUN_CMD) + mock_docker_run_cmd.assert_called_once_with("default", "/tmp/test_dir") + mock_cmd_output.assert_any_call("mock run command") # Verify archives directory was created mock_makedirs.assert_called_once_with("payload/archives", exist_ok=True) @@ -104,8 +117,17 @@ def test_prepare_dependency_archive_image_exists( @patch("datacustomcode.deploy.tempfile.TemporaryDirectory") @patch("datacustomcode.deploy.os.path.join") @patch("datacustomcode.deploy.os.makedirs") + @patch("datacustomcode.deploy.docker_build_cmd") + @patch("datacustomcode.deploy.docker_run_cmd") def test_prepare_dependency_archive_build_image( - self, mock_makedirs, mock_join, mock_temp_dir, mock_copy, mock_cmd_output + self, + mock_docker_run_cmd, + mock_docker_build_cmd, + mock_makedirs, + mock_join, + mock_temp_dir, + mock_copy, + mock_cmd_output, ): """Test prepare_dependency_archive when Docker image needs to be built.""" # Mock the temporary directory context manager @@ -121,20 +143,26 @@ def test_prepare_dependency_archive_build_image( # Mock os.path.join for archive path mock_join.return_value = "/tmp/test_dir/native_dependencies.tar.gz" - prepare_dependency_archive("/test/dir") + # Mock the docker command functions + mock_docker_build_cmd.return_value = "mock build command" + mock_docker_run_cmd.return_value = "mock run command" + + prepare_dependency_archive("/test/dir", "default") # Verify docker images command was called mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_IMAGES_CMD) # Verify docker build command was called - mock_cmd_output.assert_any_call(self.EXPECTED_BUILD_CMD) + mock_docker_build_cmd.assert_called_once_with("default") + mock_cmd_output.assert_any_call("mock build command") # Verify files were copied to temp directory mock_copy.assert_any_call("requirements.txt", "/tmp/test_dir") mock_copy.assert_any_call("build_native_dependencies.sh", "/tmp/test_dir") # Verify docker run command was called - mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_RUN_CMD) + mock_docker_run_cmd.assert_called_once_with("default", "/tmp/test_dir") + mock_cmd_output.assert_any_call("mock run command") # Verify archives directory was created mock_makedirs.assert_called_once_with("payload/archives", exist_ok=True) @@ -150,8 +178,17 @@ def test_prepare_dependency_archive_build_image( @patch("datacustomcode.deploy.tempfile.TemporaryDirectory") @patch("datacustomcode.deploy.os.path.join") @patch("datacustomcode.deploy.os.makedirs") + @patch("datacustomcode.deploy.docker_build_cmd") + @patch("datacustomcode.deploy.docker_run_cmd") def test_prepare_dependency_archive_docker_build_failure( - self, mock_makedirs, mock_join, mock_temp_dir, mock_copy, mock_cmd_output + self, + mock_docker_run_cmd, + mock_docker_build_cmd, + mock_makedirs, + mock_join, + mock_temp_dir, + mock_copy, + mock_cmd_output, ): """Test prepare_dependency_archive when Docker build fails.""" # Mock the temporary directory context manager @@ -171,21 +208,30 @@ def test_prepare_dependency_archive_docker_build_failure( ] with pytest.raises(CalledProcessError, match="Build failed"): - prepare_dependency_archive("/test/dir") + prepare_dependency_archive("/test/dir", "default") # Verify docker images command was called mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_IMAGES_CMD) # Verify docker build command was called - mock_cmd_output.assert_any_call(self.EXPECTED_BUILD_CMD) + mock_docker_build_cmd.assert_called_once_with("default") @patch("datacustomcode.deploy.cmd_output") @patch("datacustomcode.deploy.shutil.copy") @patch("datacustomcode.deploy.tempfile.TemporaryDirectory") @patch("datacustomcode.deploy.os.path.join") @patch("datacustomcode.deploy.os.makedirs") + @patch("datacustomcode.deploy.docker_build_cmd") + @patch("datacustomcode.deploy.docker_run_cmd") def test_prepare_dependency_archive_docker_run_failure( - self, mock_makedirs, mock_join, mock_temp_dir, mock_copy, mock_cmd_output + self, + mock_docker_run_cmd, + mock_docker_build_cmd, + mock_makedirs, + mock_join, + mock_temp_dir, + mock_copy, + mock_cmd_output, ): """Test prepare_dependency_archive when Docker run fails.""" # Mock the temporary directory context manager @@ -205,7 +251,7 @@ def test_prepare_dependency_archive_docker_run_failure( ] with pytest.raises(CalledProcessError, match="Run failed"): - prepare_dependency_archive("/test/dir") + prepare_dependency_archive("/test/dir", "default") # Verify docker images command was called mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_IMAGES_CMD) @@ -215,15 +261,24 @@ def test_prepare_dependency_archive_docker_run_failure( mock_copy.assert_any_call("build_native_dependencies.sh", "/tmp/test_dir") # Verify docker run command was called - mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_RUN_CMD) + mock_docker_run_cmd.assert_called_once_with("default", "/tmp/test_dir") @patch("datacustomcode.deploy.cmd_output") @patch("datacustomcode.deploy.shutil.copy") @patch("datacustomcode.deploy.tempfile.TemporaryDirectory") @patch("datacustomcode.deploy.os.path.join") @patch("datacustomcode.deploy.os.makedirs") + @patch("datacustomcode.deploy.docker_build_cmd") + @patch("datacustomcode.deploy.docker_run_cmd") def test_prepare_dependency_archive_file_copy_failure( - self, mock_makedirs, mock_join, mock_temp_dir, mock_copy, mock_cmd_output + self, + mock_docker_run_cmd, + mock_docker_build_cmd, + mock_makedirs, + mock_join, + mock_temp_dir, + mock_copy, + mock_cmd_output, ): """Test prepare_dependency_archive when file copy fails.""" # Mock the temporary directory context manager @@ -239,7 +294,7 @@ def test_prepare_dependency_archive_file_copy_failure( mock_copy.side_effect = FileNotFoundError("File not found") with pytest.raises(FileNotFoundError, match="File not found"): - prepare_dependency_archive("/test/dir") + prepare_dependency_archive("/test/dir", "default") # Verify docker images command was called mock_cmd_output.assert_any_call(self.EXPECTED_DOCKER_IMAGES_CMD) @@ -494,10 +549,10 @@ def test_zip_with_requirements( ("/test/dir/subdir", [], ["file3.py"]), ] - zip("/test/dir") + zip("/test/dir", "default") mock_has_requirements.assert_called_once_with("/test/dir") - mock_prepare.assert_called_once_with("/test/dir") + mock_prepare.assert_called_once_with("/test/dir", "default") mock_zipfile.assert_called_once_with( "deployment.zip", "w", zipfile.ZIP_DEFLATED ) @@ -522,7 +577,7 @@ def test_zip_without_requirements( ("/test/dir/subdir", [], ["file3.py"]), ] - zip("/test/dir") + zip("/test/dir", "default") mock_has_requirements.assert_called_once_with("/test/dir") mock_prepare.assert_not_called() @@ -797,13 +852,13 @@ def test_deploy_full( ) # Call function - result = deploy_full("/test/dir", metadata, credentials, callback) + result = deploy_full("/test/dir", metadata, credentials, "default", callback) # Assertions mock_retrieve_token.assert_called_once_with(credentials) mock_verify_config.assert_called_once_with("/test/dir") mock_create_deployment.assert_called_once_with(access_token, metadata) - mock_zip.assert_called_once_with("/test/dir") + mock_zip.assert_called_once_with("/test/dir", "default") mock_upload_zip.assert_called_once_with("https://upload.example.com") mock_wait.assert_called_once_with(access_token, metadata, callback) mock_create_transform.assert_called_once_with( @@ -883,13 +938,13 @@ def test_deploy_full_happy_path( mock_has_requirements.return_value = True # Call function - result = deploy_full("/test/dir", metadata, credentials, callback) + result = deploy_full("/test/dir", metadata, credentials, "default", callback) # Assertions mock_retrieve_token.assert_called_once_with(credentials) mock_verify_config.assert_called_once_with("/test/dir") mock_create_deployment.assert_called_once_with(access_token, metadata) - mock_zip.assert_called_once_with("/test/dir") + mock_zip.assert_called_once_with("/test/dir", "default") mock_upload_zip.assert_called_once_with("https://upload.example.com") mock_wait.assert_called_once_with(access_token, metadata, callback) mock_create_transform.assert_called_once_with(