Coverage for src/utils.py: 100%

31 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Utility functions for Chuck TUI. 

3""" 

4 

5import time 

6import logging 

7 

8 

9def build_query_params(params): 

10 """ 

11 Build URL query parameters string from a dictionary. 

12 

13 Args: 

14 params: Dictionary of parameter names and values 

15 

16 Returns: 

17 URL query string starting with '?' if parameters exist, empty string otherwise 

18 """ 

19 if not params: 

20 return "" 

21 

22 # Filter out None values 

23 filtered_params = {k: v for k, v in params.items() if v is not None} 

24 

25 # Convert boolean values to strings 

26 for key, value in filtered_params.items(): 

27 if isinstance(value, bool): 

28 filtered_params[key] = "true" if value else "false" 

29 else: 

30 filtered_params[key] = str(value) 

31 

32 # Build query string 

33 query_string = "&".join( 

34 [f"{key}={value}" for key, value in filtered_params.items()] 

35 ) 

36 return f"?{query_string}" if query_string else "" 

37 

38 

39def execute_sql_statement( 

40 client, warehouse_id, sql_text, catalog=None, wait_timeout="5s" 

41): 

42 """ 

43 Execute SQL statement and poll for results. 

44 

45 Args: 

46 client: DatabricksAPIClient instance 

47 warehouse_id: ID of the SQL warehouse 

48 sql_text: SQL query text 

49 catalog: Optional catalog name 

50 wait_timeout: Wait timeout value 

51 

52 Returns: 

53 SQL statement result data 

54 """ 

55 # Prepare request data 

56 data = { 

57 "on_wait_timeout": "CONTINUE", 

58 "statement": sql_text, 

59 "wait_timeout": wait_timeout, 

60 "warehouse_id": warehouse_id, 

61 } 

62 

63 if catalog: 

64 data["catalog"] = catalog 

65 

66 logging.debug(f"Executing SQL: {sql_text}") 

67 

68 # Submit the SQL statement 

69 response = client.post("/api/2.0/sql/statements", data) 

70 statement_id = response.get("statement_id") 

71 

72 # Poll until complete 

73 state = None 

74 while True: 

75 status = client.get(f"/api/2.0/sql/statements/{statement_id}") 

76 state = status.get("status", {}).get("state", status.get("state")) 

77 if state not in ["PENDING", "RUNNING"]: 

78 break 

79 time.sleep(1) 

80 

81 # Check result 

82 if state != "SUCCEEDED": 

83 error = status.get("status", {}).get("error", status.get("error", {})) 

84 message = error.get("message", "Unknown error") 

85 raise ValueError(f"SQL statement failed: {message}") 

86 

87 return status.get("result", {})