deploy sf-cli-org

joroscoSF · joroscoSF · commit 0aecc3a94e41 · 2026-02-26T16:22:59.000-08:00
diff --git a/src/datacustomcode/cli.py b/src/datacustomcode/cli.py
@@ -160,6 +160,11 @@ def zip(path: str, network: str):
     Choose based on your workload requirements.""",
 )
 @click.option("--function-invoke-opt")
+@click.option(
+    "--sf-cli-org",
+    default=None,
+    help="SF CLI org alias or username. Fetches credentials via `sf org display`.",
+)
 def deploy(
     path: str,
     name: str,
@@ -169,15 +174,19 @@ def deploy(
     profile: str,
     network: str,
     function_invoke_opt: str,
+    sf_cli_org: Optional[str],
 ):
     from datacustomcode.credentials import Credentials
-    from datacustomcode.deploy import CodeExtensionMetadata, deploy_full
+    from datacustomcode.deploy import (
+        COMPUTE_TYPES,
+        AccessTokenResponse,
+        CodeExtensionMetadata,
+        _retrieve_access_token_from_sf_cli,
+        deploy_full,
+    )
 
     logger.debug("Deploying project")
 
-    # Validate compute type
-    from datacustomcode.deploy import COMPUTE_TYPES
-
     if cpu_size not in COMPUTE_TYPES.keys():
         click.secho(
             f"Error: Invalid CPU size '{cpu_size}'. "
@@ -208,15 +217,23 @@ def deploy(
             function_invoke_options = function_invoke_opt.split(",")
             metadata.functionInvokeOptions = function_invoke_options
 
-    try:
-        credentials = Credentials.from_available(profile=profile)
-    except ValueError as e:
-        click.secho(
-            f"Error: {e}",
-            fg="red",
-        )
-        raise click.Abort() from None
-    deploy_full(path, metadata, credentials, network)
+    auth: Union[Credentials, AccessTokenResponse]
+    if sf_cli_org:
+        try:
+            auth = _retrieve_access_token_from_sf_cli(sf_cli_org)
+        except RuntimeError as e:
+            click.secho(f"Error: {e}", fg="red")
+            raise click.Abort() from None
+    else:
+        try:
+            auth = Credentials.from_available(profile=profile)
+        except ValueError as e:
+            click.secho(
+                f"Error: {e}",
+                fg="red",
+            )
+            raise click.Abort() from None
+    deploy_full(path, metadata, auth, network)
 
 
 @cli.command()
diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py
@@ -18,6 +18,7 @@
 import json
 import os
 import shutil
+import subprocess
 import tempfile
 import time
 from typing import (
@@ -156,6 +157,54 @@ def _retrieve_access_token(credentials: Credentials) -> AccessTokenResponse:
     return AccessTokenResponse(**response)
 
 
+def _retrieve_access_token_from_sf_cli(sf_cli_org: str) -> AccessTokenResponse:
+    """Get an access token from the Salesforce CLI."""
+    try:
+        result = subprocess.run(
+            ["sf", "org", "display", "--target-org", sf_cli_org, "--json"],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=30,
+        )
+    except FileNotFoundError as exc:
+        raise RuntimeError(
+            "The 'sf' command was not found. "
+            "Please install Salesforce CLI: https://developer.salesforce.com/tools/salesforcecli"
+        ) from exc
+    except subprocess.TimeoutExpired as exc:
+        raise RuntimeError(
+            f"'sf org display' timed out for org '{sf_cli_org}'"
+        ) from exc
+    except subprocess.CalledProcessError as exc:
+        raise RuntimeError(
+            f"'sf org display' failed for org '{sf_cli_org}'.\n"
+            f"Ensure the org is authenticated via 'sf org login web'.\n"
+            f"stderr: {exc.stderr.strip()}"
+        ) from exc
+
+    try:
+        data = json.loads(result.stdout)
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(f"Failed to parse 'sf org display' output: {exc}") from exc
+
+    if data.get("status") != 0:
+        raise RuntimeError(
+            f"SF CLI error for org '{sf_cli_org}': "
+            f"{data.get('message', 'unknown error')}"
+        )
+
+    org_result = data.get("result", {})
+    access_token = org_result.get("accessToken")
+    instance_url = org_result.get("instanceUrl")
+    if not access_token or not instance_url:
+        raise RuntimeError(
+            f"'sf org display' did not return an access token or instance URL "
+            f"for org '{sf_cli_org}'"
+        )
+    return AccessTokenResponse(access_token=access_token, instance_url=instance_url)
+
+
 class CreateDeploymentResponse(BaseModel):
     fileUploadUrl: str
 
@@ -463,12 +512,15 @@ def zip(
 def deploy_full(
     directory: str,
     metadata: CodeExtensionMetadata,
-    credentials: Credentials,
+    credentials: Union["Credentials", AccessTokenResponse],
     docker_network: str,
     callback=None,
 ) -> AccessTokenResponse:
     """Deploy a data transform in the DataCloud."""
-    access_token = _retrieve_access_token(credentials)
+    if isinstance(credentials, AccessTokenResponse):
+        access_token = credentials
+    else:
+        access_token = _retrieve_access_token(credentials)
 
     # prepare payload
     config = get_config(directory)
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -5,6 +5,7 @@
 from click.testing import CliRunner
 
 from datacustomcode.cli import deploy, init
+from datacustomcode.deploy import AccessTokenResponse
 from datacustomcode.scan import write_sdk_config
 
 
@@ -180,3 +181,39 @@ def test_deploy_command_custom_description(
             # Check that deploy_full was called with custom description
             call_args = mock_deploy_full.call_args
             assert call_args[0][1].description == "Custom description"
+
+    @patch("datacustomcode.deploy.deploy_full")
+    @patch("datacustomcode.deploy._retrieve_access_token_from_sf_cli")
+    def test_deploy_command_sf_cli_org(self, mock_sf_cli_token, mock_deploy_full):
+        """Test deploy command with --sf-cli-org flag."""
+        mock_token = AccessTokenResponse(
+            access_token="test_token", instance_url="https://test.salesforce.com"
+        )
+        mock_sf_cli_token.return_value = mock_token
+
+        runner = CliRunner()
+        with runner.isolated_filesystem():
+            os.makedirs("payload", exist_ok=True)
+            result = runner.invoke(
+                deploy, ["--name", "test-job", "--sf-cli-org", "my-org"]
+            )
+
+            assert result.exit_code == 0
+            mock_sf_cli_token.assert_called_once_with("my-org")
+            mock_deploy_full.assert_called_once()
+            call_args = mock_deploy_full.call_args
+            assert call_args[0][2] == mock_token  # AccessTokenResponse passed directly
+
+    @patch("datacustomcode.deploy._retrieve_access_token_from_sf_cli")
+    def test_deploy_command_sf_cli_org_error(self, mock_sf_cli_token):
+        """Test deploy command when --sf-cli-org fails."""
+        mock_sf_cli_token.side_effect = RuntimeError("sf command not found")
+
+        runner = CliRunner()
+        with runner.isolated_filesystem():
+            os.makedirs("payload", exist_ok=True)
+            result = runner.invoke(
+                deploy, ["--name", "test-job", "--sf-cli-org", "bad-org"]
+            )
+            assert result.exit_code == 1
+            assert "sf command not found" in result.output
diff --git a/tests/test_deploy.py b/tests/test_deploy.py
@@ -1,5 +1,6 @@
 """Tests for the deploy module."""
 
+import json
 from unittest.mock import (
     MagicMock,
     mock_open,
@@ -27,6 +28,7 @@
         DeploymentsResponse,
         _make_api_call,
         _retrieve_access_token,
+        _retrieve_access_token_from_sf_cli,
         create_data_transform,
         create_deployment,
         deploy_full,
@@ -1107,3 +1109,144 @@ def test_deploy_full_happy_path(
             "/test/dir", access_token, metadata, data_transform_config
         )
         assert result == access_token
+
+
+class TestRetrieveAccessTokenFromSFCLI:
+    """Tests for _retrieve_access_token_from_sf_cli."""
+
+    SF_CLI_OUTPUT = json.dumps(
+        {
+            "status": 0,
+            "result": {
+                "accessToken": "sf_access_token",
+                "instanceUrl": "https://sf.salesforce.com",
+            },
+        }
+    )
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_happy_path(self, mock_run):
+        """Successful sf org display returns AccessTokenResponse."""
+        mock_run.return_value = MagicMock(stdout=self.SF_CLI_OUTPUT, returncode=0)
+
+        result = _retrieve_access_token_from_sf_cli("my-org")
+
+        assert result.access_token == "sf_access_token"
+        assert result.instance_url == "https://sf.salesforce.com"
+        mock_run.assert_called_once_with(
+            ["sf", "org", "display", "--target-org", "my-org", "--json"],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=30,
+        )
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_file_not_found(self, mock_run):
+        """FileNotFoundError raised when sf CLI is not installed."""
+        mock_run.side_effect = FileNotFoundError("No such file or directory: 'sf'")
+
+        with pytest.raises(RuntimeError, match="'sf' command was not found"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_timeout_expired(self, mock_run):
+        """TimeoutExpired raised when sf CLI times out."""
+        import subprocess
+
+        mock_run.side_effect = subprocess.TimeoutExpired(cmd="sf", timeout=30)
+
+        with pytest.raises(RuntimeError, match="timed out"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_called_process_error(self, mock_run):
+        """CalledProcessError raised when sf CLI exits non-zero."""
+        import subprocess
+
+        mock_run.side_effect = subprocess.CalledProcessError(
+            returncode=1, cmd="sf", stderr="Org not found"
+        )
+
+        with pytest.raises(RuntimeError, match="failed for org"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_json_decode_error(self, mock_run):
+        """RuntimeError raised when output is not valid JSON."""
+        mock_run.return_value = MagicMock(stdout="not-json", returncode=0)
+
+        with pytest.raises(RuntimeError, match="Failed to parse"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_nonzero_status_in_json(self, mock_run):
+        """RuntimeError raised when JSON status field is non-zero."""
+        output = json.dumps({"status": 1, "message": "org not found"})
+        mock_run.return_value = MagicMock(stdout=output, returncode=0)
+
+        with pytest.raises(RuntimeError, match="SF CLI error"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_missing_access_token(self, mock_run):
+        """RuntimeError raised when accessToken is absent."""
+        output = json.dumps(
+            {"status": 0, "result": {"instanceUrl": "https://sf.salesforce.com"}}
+        )
+        mock_run.return_value = MagicMock(stdout=output, returncode=0)
+
+        with pytest.raises(RuntimeError, match="did not return"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+    @patch("datacustomcode.deploy.subprocess.run")
+    def test_missing_instance_url(self, mock_run):
+        """RuntimeError raised when instanceUrl is absent."""
+        output = json.dumps({"status": 0, "result": {"accessToken": "sf_access_token"}})
+        mock_run.return_value = MagicMock(stdout=output, returncode=0)
+
+        with pytest.raises(RuntimeError, match="did not return"):
+            _retrieve_access_token_from_sf_cli("my-org")
+
+
+class TestDeployFullWithAccessTokenResponse:
+    """Test deploy_full when passed an AccessTokenResponse directly."""
+
+    @patch("datacustomcode.deploy.create_data_transform")
+    @patch("datacustomcode.deploy.wait_for_deployment")
+    @patch("datacustomcode.deploy.upload_zip")
+    @patch("datacustomcode.deploy.zip")
+    @patch("datacustomcode.deploy.create_deployment")
+    @patch("datacustomcode.deploy.get_config")
+    @patch("datacustomcode.deploy._retrieve_access_token")
+    def test_deploy_full_with_access_token_response_skips_token_exchange(
+        self,
+        mock_retrieve_token,
+        mock_get_config,
+        mock_create_deployment,
+        mock_zip,
+        mock_upload_zip,
+        mock_wait,
+        mock_create_transform,
+    ):
+        """deploy_full skips token exchange when given an AccessTokenResponse."""
+        access_token = AccessTokenResponse(
+            access_token="direct_token", instance_url="https://instance.example.com"
+        )
+        metadata = CodeExtensionMetadata(
+            name="test",
+            version="1.0.0",
+            description="desc",
+            computeType="CPU_M",
+            codeType="script",
+        )
+        mock_get_config.return_value = MagicMock(spec=[])  # not DataTransformConfig
+        mock_create_deployment.return_value = CreateDeploymentResponse(
+            fileUploadUrl="https://upload.example.com"
+        )
+
+        result = deploy_full("/test/dir", metadata, access_token, "default")
+
+        mock_retrieve_token.assert_not_called()
+        mock_create_deployment.assert_called_once_with(access_token, metadata)
+        assert result == access_token
