Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/check_migration.py: 9%

1"""Module for checking differences between Prisma schema and database.""" 

2 

3import os 

4import subprocess 

5from typing import Final 

6 

7from litellm._logging import verbose_logger 

8 

9 

10def extract_sql_commands(diff_output: str) -> list[str]: 

11 """ 

12 Extract SQL commands from the Prisma migrate diff output. 

13 Args: 

14 diff_output (str): The full output from prisma migrate diff. 

15 Returns: 

16 List[str]: A list of SQL commands extracted from the diff output. 

17 """ 

18 # Split the output into lines and remove empty lines 

19 lines: Final = [line.strip() for line in diff_output.split("\n") if line.strip()] 

20 

21 sql_commands: Final = [] 

22 current_command = "" 

23 in_sql_block = False 

24 

25 for line in lines: 

26 if line.startswith("-- "): # Comment line, likely a table operation description 

27 if in_sql_block and current_command: 

28 sql_commands.append(current_command.strip()) 

29 current_command = "" 

30 in_sql_block = True 

31 elif in_sql_block: 

32 if line.endswith(";"): 

33 current_command += line 

34 sql_commands.append(current_command.strip()) 

35 current_command = "" 

36 in_sql_block = False 

37 else: 

38 current_command += line + " " 

39 

40 # Add any remaining command 

41 if current_command: 

42 sql_commands.append(current_command.strip()) 

43 

44 return sql_commands 

45 

46 

47def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: 

48 """Checks for differences between current database and Prisma schema. 

49 

50 Never raises: a diff that cannot be produced, because the runner is missing, 

51 because the command failed, or because it outlived its budget, is reported as 

52 "no diff" so boot continues. 

53 

54 Returns: 

55 A tuple containing: 

56 - A boolean indicating if differences were found (True) or not (False). 

57 - The SQL commands that would close the diff, empty when there is none. 

58 """ 

59 try: 

60 from litellm_proxy_extras.prisma_toolchain import ( 

61 PRISMA_COMMAND_TIMEOUT_ENV_VAR, 

62 prisma_command_timeout, 

63 run_prisma, 

64 ) 

65 except ImportError as e: 

66 print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages 

67 f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" 

68 ) 

69 return False, [] 

70 

71 verbose_logger.debug("Checking for Prisma schema diff...") 

72 timeout: Final = prisma_command_timeout() 

73 try: 

74 result: Final = run_prisma( 

75 [ 

76 "prisma", 

77 "migrate", 

78 "diff", 

79 "--from-url", 

80 db_url, 

81 "--to-schema-datamodel", 

82 "./schema.prisma", 

83 "--script", 

84 ], 

85 timeout=timeout, 

86 env=os.environ.copy(), 

87 ) 

88 

89 sql_commands: Final = extract_sql_commands(result.stdout) 

90 

91 if sql_commands: 

92 print("Changes to DB Schema detected") # noqa: T201 

93 print("Required SQL commands:") # noqa: T201 

94 for command in sql_commands: 

95 print(command) # noqa: T201 

96 return True, sql_commands 

97 else: 

98 return False, [] 

99 except subprocess.TimeoutExpired: 

100 print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages 

101 f"Timed out after {timeout}s generating the migration diff. " 

102 f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." 

103 ) 

104 return False, [] 

105 except subprocess.CalledProcessError as e: 

106 error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" 

107 print(error_message) # noqa: T201 

108 return False, [] 

109 

110 

111def check_prisma_schema_diff(db_url: str | None = None) -> None: 

112 """Main function to run the Prisma schema diff check.""" 

113 if db_url is None: 

114 db_url = os.getenv("DATABASE_URL") 

115 if db_url is None: 

116 raise Exception("DATABASE_URL not set") 

117 has_diff, message = check_prisma_schema_diff_helper(db_url) 

118 if has_diff: 

119 verbose_logger.exception( 

120 "🚨🚨🚨 prisma schema out of sync with db. Consider running these sql_commands to sync the two - %s", 

121 message, 

122 )