Coverage for slidge/migration.py: 47%

1import logging 

2import shutil 

3import sys 

4from pathlib import Path 

5 

6from alembic import command 

7from alembic.config import Config 

8from slixmpp import JID 

9 

10from .core import config 

11from .db.meta import get_engine 

12from .db.models import GatewayUser 

13from .db.store import SlidgeStore 

14 

15 

16def remove_avatar_cache_v1(): 

17 old_dir = config.HOME_DIR / "slidge_avatars" 

18 if old_dir.exists(): 

19 log.info("Avatar cache dir v1 found, clearing it.") 

20 shutil.rmtree(old_dir) 

21 

22 

23def get_alembic_cfg() -> Config: 

24 alembic_cfg = Config() 

25 alembic_cfg.set_section_option( 

26 "alembic", 

27 "script_location", 

28 str(Path(__file__).parent / "db" / "alembic"), 

29 ) 

30 return alembic_cfg 

31 

32 

33def remove_resource_parts_from_users() -> None: 

34 with SlidgeStore(get_engine(config.DB_URL)).session() as orm: 

35 for user in orm.query(GatewayUser).all(): 

36 if user.jid.resource: 

37 user.jid = JID(user.jid.bare) 

38 orm.add(user) 

39 orm.commit() 

40 

41 

42def migrate() -> None: 

43 remove_avatar_cache_v1() 

44 command.upgrade(get_alembic_cfg(), "head") 

45 remove_resource_parts_from_users() 

46 

47 

48def main(): 

49 """ 

50 Updates the (dev) database in ./dev/slidge.sqlite and generates a revision 

51 

52 Usage: python -m slidge.migration "Revision message blah blah blah" 

53 """ 

54 alembic_cfg = get_alembic_cfg() 

55 command.upgrade(alembic_cfg, "head") 

56 command.revision(alembic_cfg, sys.argv[1], autogenerate=True) 

57 

58 

59log = logging.getLogger(__name__) 

60 

61if __name__ == "__main__": 

62 main()