Add computeType argument

Author: ppandit-sfdc

README.md

@@ -71,11 +71,16 @@ datacustomcode run ./payload/entrypoint.py
 After modifying the `entrypoint.py` as needed, using any dependencies you add in the `.venv` virtual environment, you can run this script in Data Cloud:
 ```zsh
 datacustomcode scan ./payload/entrypoint.py
-datacustomcode deploy --path ./payload --name my_custom_script
+datacustomcode deploy --path ./payload --name my_custom_script --compute-type CPU_L
 ```
 
 > [!TIP]
 > The `deploy` process can take several minutes.  If you'd like more feedback on the underlying process, you can add `--debug` to the command like `datacustomcode --debug deploy --path ./payload --name my_custom_script`
+>
+> [!NOTE]
+> **Compute Types**: Choose the appropriate compute type based on your workload requirements:
+> - **CPU_L/CPU_XL/CPU_2XL/CPU_4XL**: Large, X-Large, 2X-Large and 4X-Large CPU instances for data processing
+> - Default is `CPU_2XL` which provides a good balance of performance and cost for most use cases
 
 You can now use the Salesforce Data Cloud UI to find the created Data Transform and use the `Run Now` button to run it.
 Once the Data Transform run is successful, check the DLO your script is writing to and verify the correct records were added.
@@ -139,6 +144,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: "")
+- `--compute-type TEXT`: Compute type for the deployment (default: "CPU_M"). Available options: CPU_XS (extra-small), CPU_S (small), CPU_M (medium), CPU_L (large)
 
 #### `datacustomcode init`
 Initialize a new development environment with a template.

src/datacustomcode/cli.py

@@ -83,16 +83,30 @@ def zip(path: str):
 @click.option("--name", required=True)
 @click.option("--version", default="0.0.1")
 @click.option("--description", default="Custom Data Transform Code")
-def deploy(path: str, name: str, version: str, description: str):
+@click.option("--compute-type", default="CPU_M")
+def deploy(path: str, name: str, version: str, description: str, compute_type: str):
     from datacustomcode.credentials import Credentials
     from datacustomcode.deploy import TransformationJobMetadata, deploy_full
 
     logger.debug("Deploying project")
+    
+    # Validate compute type
+    from datacustomcode.deploy import COMPUTE_TYPES
+    if compute_type not in COMPUTE_TYPES:
+        click.secho(
+            f"Error: Invalid compute type '{compute_type}'. "
+            f"Available options: {', '.join(COMPUTE_TYPES.keys())}",
+            fg="red",
+        )
+        raise click.Abort()
 
+    logger.info(f"Deploying with compute type: {compute_type}")
+    
     metadata = TransformationJobMetadata(
         name=name,
         version=version,
         description=description,
+        computeType=compute_type,
     )
     try:
         credentials = Credentials.from_available()

src/datacustomcode/deploy.py

@@ -44,11 +44,25 @@
 AUTH_PATH = "services/oauth2/token"
 WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000
 
+# Available compute types for Data Cloud deployments
+COMPUTE_TYPES = {
+    "CPU_XS": "CPU_XS",      # Extra Small CPU instance
+    "CPU_S": "CPU_S",      # Small CPU instance
+    "CPU_M": "CPU_M",     # Medium CPU instance (default)
+    "CPU_L": "CPU_L"      # Large CPU instance
+}
+
 
 class TransformationJobMetadata(BaseModel):
     name: str
     version: str
     description: str
+    computeType: str = "CPU_M"  # Default to CPU_M for backward compatibility
+    
+    def __init__(self, **data):
+        super().__init__(**data)
+        if self.computeType not in COMPUTE_TYPES:
+            raise ValueError(f"Invalid compute type '{self.computeType}'. Available options: {', '.join(COMPUTE_TYPES.keys())}")
 
 
 def _join_strip_url(*args: str) -> str:
@@ -123,7 +137,7 @@ def create_deployment(
         "name": metadata.name,
         "description": metadata.description,
         "version": metadata.version,
-        "computeType": "CPU_M",
+        "computeType": metadata.computeType,
     }
     logger.debug(f"Creating deployment {metadata.name}...")
     try: