Coverage for src/chuck_data/commands/cluster_init_tools.py: 0%

39 statements  

« 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. 

3 

4This module contains the upload logic for cluster init scripts with versioning 

5to prevent conflicts between concurrent stitch runs. 

6""" 

7 

8import logging 

9import datetime 

10from typing import Dict, Any 

11 

12from ..clients.databricks import DatabricksAPIClient 

13 

14 

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 } 

26 

27 if not init_script_content.strip(): 

28 return {"error": "Init script content cannot be empty."} 

29 

30 # Step 1: Check/Create "chuck" volume (same as stitch) 

31 volume_name = "chuck" 

32 volume_exists = False 

33 

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)}"} 

45 

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)}"} 

59 

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" 

63 

64 # Step 3: Upload the init script to the volume 

65 volume_path = f"/Volumes/{target_catalog}/{target_schema}/{volume_name}/{init_script_filename}" 

66 

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}"} 

74 

75 logging.info(f"Cluster init script uploaded successfully to {volume_path}") 

76 

77 return { 

78 "success": True, 

79 "volume_path": volume_path, 

80 "filename": init_script_filename, 

81 "timestamp": current_datetime, 

82 } 

83 

84 except Exception as e: 

85 return {"error": f"Failed to upload init script: {str(e)}"}