Coverage for src/commands/cluster_init_tools.py: 100%
39 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
1"""
2Cluster init script upload helper for stitch setup.
4This module contains the upload logic for cluster init scripts with versioning
5to prevent conflicts between concurrent stitch runs.
6"""
8import logging
9import datetime
10from typing import Dict, Any
12from src.clients.databricks import DatabricksAPIClient
15def _helper_upload_cluster_init_logic(
16 client: DatabricksAPIClient,
17 target_catalog: str,
18 target_schema: str,
19 init_script_content: str,
20) -> Dict[str, Any]:
21 """Internal logic for uploading cluster init script."""
22 if not target_catalog or not target_schema:
23 return {
24 "error": "Target catalog and schema are required for cluster init upload."
25 }
27 if not init_script_content.strip():
28 return {"error": "Init script content cannot be empty."}
30 # Step 1: Check/Create "chuck" volume (same as stitch)
31 volume_name = "chuck"
32 volume_exists = False
34 # Check if volume exists
35 try:
36 volumes_response = client.list_volumes(
37 catalog_name=target_catalog, schema_name=target_schema
38 )
39 for volume_info in volumes_response.get("volumes", []):
40 if volume_info.get("name") == volume_name:
41 volume_exists = True
42 break
43 except Exception as e:
44 return {"error": f"Failed to list volumes: {str(e)}"}
46 if not volume_exists:
47 logging.debug(
48 f"Volume '{volume_name}' not found in {target_catalog}.{target_schema}. Attempting to create."
49 )
50 try:
51 volume_response = client.create_volume(
52 catalog_name=target_catalog, schema_name=target_schema, name=volume_name
53 )
54 if not volume_response:
55 return {"error": f"Failed to create volume '{volume_name}'"}
56 logging.debug(f"Volume '{volume_name}' created successfully.")
57 except Exception as e:
58 return {"error": f"Failed to create volume '{volume_name}': {str(e)}"}
60 # Step 2: Generate timestamped filename (matching stitch pattern)
61 current_datetime = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M")
62 init_script_filename = f"cluster_init-{current_datetime}.sh"
64 # Step 3: Upload the init script to the volume
65 volume_path = f"/Volumes/{target_catalog}/{target_schema}/{volume_name}/{init_script_filename}"
67 try:
68 # Upload file to volume
69 upload_response = client.upload_file(
70 path=volume_path, content=init_script_content, overwrite=True
71 )
72 if not upload_response:
73 return {"error": f"Failed to upload init script to {volume_path}"}
75 logging.info(f"Cluster init script uploaded successfully to {volume_path}")
77 return {
78 "success": True,
79 "volume_path": volume_path,
80 "filename": init_script_filename,
81 "timestamp": current_datetime,
82 }
84 except Exception as e:
85 return {"error": f"Failed to upload init script: {str(e)}"}