Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/sql/alembic/env.py: 100.00%

14 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2025-02-11 15:07 +0000

1# Copyright (C) 2019 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15import sys 

16from logging import getLogger 

17from logging.config import fileConfig 

18 

19from alembic import context 

20from sqlalchemy import engine_from_config, pool 

21from sqlalchemy.engine.base import Connection 

22 

23from buildgrid.server.sql.models import Base 

24 

25# this is the Alembic Config object, which provides 

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

27config = context.config 

28 

29# Interpret the config file for Python logging. 

30# This line sets up loggers basically. 

31try: 

32 fileConfig(config.config_file_name) # type: ignore[arg-type] # covariance in union issue. Not a real error. 

33except TypeError: 

34 pass 

35 

36LOGGER = getLogger(__name__) 

37 

38target_metadata = Base.metadata 

39 

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

41# can be acquired: 

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

43# ... etc. 

44 

45 

46def run_migrations_offline() -> None: # pragma: no cover 

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

48 

49 This configures the context with just a URL 

50 and not an Engine, though an Engine is acceptable 

51 here as well. By skipping the Engine creation 

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

53 

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

55 script output. 

56 

57 """ 

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

59 context.configure(url=url, target_metadata=target_metadata, literal_binds=True) 

60 

61 with context.begin_transaction(): 

62 context.run_migrations() 

63 

64 

65def _run_migrations_online(connection: Connection) -> None: # pragma: no cover 

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

67 

68 with context.begin_transaction(): 

69 context.run_migrations() 

70 

71 

72def run_migrations_online() -> None: # pragma: no cover 

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

74 

75 In this scenario we need to create an Engine 

76 and associate a connection with the context. 

77 

78 """ 

79 try: 

80 connection = config.attributes["connection"] 

81 _run_migrations_online(connection) 

82 

83 except KeyError: 

84 config_section = config.get_section(config.config_ini_section) 

85 if config_section is not None: 

86 connectable = engine_from_config(config_section, prefix="sqlalchemy.", poolclass=pool.NullPool) 

87 with connectable.connect() as connection: 

88 _run_migrations_online(connection) 

89 else: 

90 LOGGER.error( 

91 "Configuration file must define either `connection` or SQLAlchemy config " 

92 f"keys prefixed with `sqlalchemy` in the {config.config_ini_section} section." 

93 ) 

94 sys.exit(1) 

95 

96 

97if context.is_offline_mode(): # pragma: no cover 

98 run_migrations_offline() 

99else: # pragma: no cover 

100 run_migrations_online()