Coverage for src/couchers/migrations/env.py: 61%

1from logging.config import fileConfig 

2from typing import Any 

3 

4from alembic import context 

5from alembic.config import Config 

6from sqlalchemy import engine_from_config, pool 

7from sqlalchemy.schema import MetaData 

8 

9from couchers import models 

10from couchers.config import config as couchers_config 

11 

12# this is the Alembic Config object, which provides 

13# access to the values within the .ini file in use. 

14config: Config = context.config 

15 

16config.set_main_option("sqlalchemy.url", couchers_config["DATABASE_CONNECTION_STRING"]) 

17 

18 

19# Interpret the config file for Python logging. 

20# This line sets up loggers basically. 

21if config.get_main_option("dont_mess_up_logging", "dont_care") == "dont_care": 

22 if not config.config_file_name: 

23 raise RuntimeError(config.config_file_name) 

24 fileConfig(config.config_file_name) 

25 

26# add your model's MetaData object here 

27# for 'autogenerate' support 

28# from myapp import mymodel 

29# target_metadata = mymodel.Base.metadata 

30target_metadata: MetaData = models.Base.metadata 

31 

32# other values from the config, defined by the needs of env.py, 

33# can be acquired: 

34# my_important_option = config.get_main_option("my_important_option") 

35# ... etc. 

36 

37 

38exclude_tables = config.get_section("alembic:exclude", {}).get("tables", "").split(",") 

39 

40 

41def include_name(name: str | None, type_: str, parent_names: Any) -> bool: 

42 if type_ == "schema": 

43 return name in [None, "logging"] 

44 if type_ == "table": 

45 return name not in exclude_tables 

46 if type_ == "index": 

47 return not (name is not None and name.startswith("idx_") and name.endswith("_geom")) 

48 return True 

49 

50 

51def run_migrations_offline() -> None: 

52 """Run migrations in 'offline' mode. 

53 

54 This configures the context with just a URL 

55 and not an Engine, though an Engine is acceptable 

56 here as well. By skipping the Engine creation 

57 we don't even need a DBAPI to be available. 

58 

59 Calls to context.execute() here emit the given string to the 

60 script output. 

61 

62 """ 

63 url = config.get_main_option("sqlalchemy.url") 

64 context.configure( 

65 url=url, 

66 target_metadata=target_metadata, 

67 literal_binds=True, 

68 dialect_opts={"paramstyle": "named"}, 

69 include_schemas=True, 

70 include_name=include_name, 

71 compare_type=True, 

72 ) 

73 

74 with context.begin_transaction(): 

75 context.run_migrations() 

76 

77 

78def run_migrations_online() -> None: 

79 """Run migrations in 'online' mode. 

80 

81 In this scenario we need to create an Engine 

82 and associate a connection with the context. 

83 

84 """ 

85 connectable = engine_from_config( 

86 config.get_section(config.config_ini_section, {}), 

87 prefix="sqlalchemy.", 

88 poolclass=pool.NullPool, 

89 ) 

90 

91 with connectable.connect() as connection: 

92 context.configure( 

93 connection=connection, target_metadata=target_metadata, include_schemas=True, include_name=include_name 

94 ) 

95 

96 with context.begin_transaction(): 

97 context.run_migrations() 

98 

99 

100if context.is_offline_mode(): 

101 run_migrations_offline() 

102else: 

103 run_migrations_online()