Coverage for mindsdb / migrations / migrate.py: 81%

1from pathlib import Path 

2 

3from alembic.command import upgrade, autogen # noqa 

4from alembic.config import Config 

5from alembic.script import ScriptDirectory 

6from alembic.script.revision import ResolutionError 

7from alembic.operations import Operations 

8from alembic.migration import MigrationContext 

9from alembic import util 

10 

11import mindsdb.interfaces.storage.db as db 

12from mindsdb.utilities import log 

13 

14logger = log.getLogger(__name__) 

15 

16# This is a migration that is like a 'base version'. Applying only this 

17# migration to a fresh DB is equivalent to applying all previous migrations. 

18current_checkpoint = "9f150e4f9a05" 

19 

20 

21def apply_checkpoint_migration(script) -> None: 

22 """Apply the checkpoint migration to the database.""" 

23 with db.engine.begin() as connection: 

24 context = MigrationContext.configure( 

25 connection, 

26 opts={ 

27 "as_sql": False, 

28 "starting_rev": None, # ignore current version 

29 "destination_rev": current_checkpoint, 

30 }, 

31 ) 

32 revision = script.get_revision(current_checkpoint) 

33 if not revision: 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 raise util.CommandError(f"Migration {current_checkpoint} not found.") 

35 

36 op = Operations(context) 

37 revision.module.upgrade(op) 

38 context.stamp(script, current_checkpoint) 

39 connection.commit() 

40 

41 

42def get_current_revision() -> str | None: 

43 """Get the current revision of the database. 

44 

45 Returns: 

46 str | None: The current revision of the database. 

47 """ 

48 with db.engine.begin() as conn: 

49 mc = MigrationContext.configure(conn) 

50 return mc.get_current_revision() 

51 

52 

53def migrate_to_head(): 

54 """Trying to update database to head revision. 

55 If alembic unable to recognize current revision (In case when database version is newer than backend) 

56 then do nothing. 

57 """ 

58 logger.debug("Applying database migrations") 

59 

60 config_file = Path(__file__).parent / "alembic.ini" 

61 config = Config(config_file) 

62 

63 # mindsdb can runs not from project directory 

64 script_location_abc = config_file.parent / config.get_main_option("script_location") 

65 config.set_main_option("script_location", str(script_location_abc)) 

66 

67 script = ScriptDirectory.from_config(config) 

68 cur_revision = get_current_revision() 

69 if cur_revision is None: 69 ↛ 73line 69 didn't jump to line 73 because the condition on line 69 was always true

70 apply_checkpoint_migration(script) 

71 cur_revision = get_current_revision() 

72 

73 try: 

74 script.revision_map.get_revision(cur_revision) 

75 except ResolutionError: 

76 raise Exception("Database version higher than application.") 

77 

78 head_rev = script.get_current_head() 

79 if cur_revision == head_rev: 79 ↛ 80line 79 didn't jump to line 80 because the condition on line 79 was never true

80 logger.debug("The database is in its current state, no updates are required.") 

81 return 

82 

83 logger.info("Migrations are available. Applying updates to the database.") 

84 upgrade(config=config, revision="head") 

85 

86 

87if __name__ == "__main__": 87 ↛ 90line 87 didn't jump to line 90 because the condition on line 87 was never true

88 # have to import this because 

89 # all env initialization happens here 

90 db.init() 

91 migrate_to_head()