Coverage for mlos_bench/mlos_bench/storage/sql/alembic/env.py: 61%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Alembic environment script.""" 

6# pylint: disable=no-member 

7 

8import sys 

9from logging.config import fileConfig 

10 

11from alembic import context 

12from sqlalchemy import engine_from_config, pool 

13 

14from mlos_bench.storage.sql.schema import DbSchema 

15 

16# this is the Alembic Config object, which provides 

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

18config = context.config 

19 

20# Interpret the config file for Python logging. 

21# This line sets up loggers basically. 

22# Don't override the mlos_bench or pytest loggers though. 

23if config.config_file_name is not None and "alembic" in sys.argv[0]: 

24 fileConfig(config.config_file_name) 

25 

26# add your model's MetaData object here 

27# for 'autogenerate' support 

28target_metadata = DbSchema(engine=None).meta 

29 

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

31# can be acquired: 

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

33# ... etc. 

34 

35 

36def run_migrations_offline() -> None: 

37 """ 

38 Run migrations in 'offline' mode. 

39 

40 This configures the context with just a URL and not an Engine, though an Engine is 

41 acceptable here as well. By skipping the Engine creation we don't even need a DBAPI 

42 to be available. 

43 

44 Calls to context.execute() here emit the given string to the script output. 

45 """ 

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

47 context.configure( 

48 url=url, 

49 target_metadata=target_metadata, 

50 literal_binds=True, 

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

52 ) 

53 

54 with context.begin_transaction(): 

55 context.run_migrations() 

56 

57 

58def run_migrations_online() -> None: 

59 """ 

60 Run migrations in 'online' mode. 

61 

62 In this scenario we need to create an Engine and associate a connection with the 

63 context. 

64 """ 

65 connectable = config.attributes.get("connection", None) 

66 

67 if connectable is None: 

68 # only create Engine if we don't have a Connection 

69 # from the outside 

70 connectable = engine_from_config( 

71 config.get_section(config.config_ini_section) or {}, 

72 prefix="sqlalchemy.", 

73 poolclass=pool.NullPool, 

74 ) 

75 

76 with connectable.connect() as connection: 

77 context.configure(connection=connection, target_metadata=target_metadata) 

78 

79 with context.begin_transaction(): 

80 context.run_migrations() 

81 else: 

82 context.configure(connection=connectable, target_metadata=target_metadata) 

83 

84 with context.begin_transaction(): 

85 context.run_migrations() 

86 

87 

88if context.is_offline_mode(): 

89 run_migrations_offline() 

90else: 

91 run_migrations_online()