Coverage for app/alembic/utils.py: 97%

1import inspect 

2from pathlib import Path 

3 

4from sqlalchemy import Engine, text 

5 

6from app.core.logger import logger 

7 

8 

9def read_sql_file(sql_file: Path) -> tuple[str, ...]: 

10 if not sql_file.exists(): 

11 raise FileNotFoundError(f"{sql_file} is missing.") 

12 

13 with open(sql_file, "r") as f: 

14 contents = f.read().strip() 

15 

16 sql_statements = tuple( 

17 statement.strip() for statement in contents.split(";") if statement.strip() 

18 ) 

19 

20 if not sql_statements: 

21 raise ValueError( 

22 f"The file '{sql_file}' exists but contains no SQL statements." 

23 ) 

24 

25 return sql_statements 

26 

27 

28def get_sql_migration_file() -> Path: 

29 caller_frame = inspect.stack()[1] 

30 caller_function_name = caller_frame.function 

31 version_file = Path(caller_frame.filename) 

32 

33 if caller_function_name not in ("upgrade", "downgrade"): 

34 raise ValueError( 

35 f"Function '{caller_function_name}' is not a valid Alembic migration function. " 

36 "Expected 'upgrade' or 'downgrade'." 

37 ) 

38 

39 return version_file.parent / f"{version_file.stem}_{caller_function_name}.sql" 

40 

41 

42def run_sql_statements(write_engine: Engine, sql_file: Path) -> tuple[str, ...]: 

43 executed_statements = [] 

44 with write_engine.begin() as conn: 

45 for sql in read_sql_file(sql_file): 

46 try: 

47 logger.info(f"Executing SQL: {sql}") 

48 conn.execute(text(sql)) 

49 executed_statements.append(sql) 

50 except Exception as e: 

51 logger.error(f"Failed to execute SQL: {sql}. Error: {e}") 

52 raise e 

53 return tuple(executed_statements)