env.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. #
  2. # Licensed to the Apache Software Foundation (ASF) under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. The ASF licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing,
  13. # software distributed under the License is distributed on an
  14. # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15. # KIND, either express or implied. See the License for the
  16. # specific language governing permissions and limitations
  17. # under the License.
  18. from __future__ import annotations
  19. import contextlib
  20. import sys
  21. from logging.config import fileConfig
  22. from alembic import context
  23. from airflow import models, settings
  24. from airflow.utils.db import compare_server_default, compare_type
  25. def include_object(_, name, type_, *args):
  26. """Filter objects for autogenerating revisions."""
  27. # Ignore the sqlite_sequence table, which is an internal SQLite construct
  28. if name == "sqlite_sequence":
  29. return False
  30. # Ignore _anything_ to do with Celery, or FlaskSession's tables
  31. if type_ == "table" and (name.startswith("celery_") or name == "session"):
  32. return False
  33. else:
  34. return True
  35. # Make sure everything is imported so that alembic can find it all
  36. models.import_all_models()
  37. # this is the Alembic Config object, which provides
  38. # access to the values within the .ini file in use.
  39. config = context.config
  40. # Interpret the config file for Python logging.
  41. # This line sets up loggers basically.
  42. fileConfig(config.config_file_name, disable_existing_loggers=False)
  43. # add your model's MetaData object here
  44. # for 'autogenerate' support
  45. # from myapp import mymodel
  46. # target_metadata = mymodel.Base.metadata
  47. target_metadata = models.base.Base.metadata
  48. # other values from the config, defined by the needs of env.py,
  49. # can be acquired:
  50. # my_important_option = config.get_main_option("my_important_option")
  51. # ... etc.
  52. def run_migrations_offline():
  53. """
  54. Run migrations in 'offline' mode.
  55. This configures the context with just a URL
  56. and not an Engine, though an Engine is acceptable
  57. here as well. By skipping the Engine creation
  58. we don't even need a DBAPI to be available.
  59. Calls to context.execute() here emit the given string to the
  60. script output.
  61. """
  62. context.configure(
  63. url=settings.SQL_ALCHEMY_CONN,
  64. target_metadata=target_metadata,
  65. literal_binds=True,
  66. compare_type=compare_type,
  67. compare_server_default=compare_server_default,
  68. render_as_batch=True,
  69. )
  70. with context.begin_transaction():
  71. context.run_migrations()
  72. def run_migrations_online():
  73. """
  74. Run migrations in 'online' mode.
  75. In this scenario we need to create an Engine
  76. and associate a connection with the context.
  77. """
  78. def process_revision_directives(context, revision, directives):
  79. if getattr(config.cmd_opts, "autogenerate", False):
  80. script = directives[0]
  81. if script.upgrade_ops and script.upgrade_ops.is_empty():
  82. directives[:] = []
  83. print("No change detected in ORM schema, skipping revision.")
  84. with contextlib.ExitStack() as stack:
  85. connection = config.attributes.get("connection", None)
  86. if not connection:
  87. connection = stack.push(settings.engine.connect())
  88. context.configure(
  89. connection=connection,
  90. transaction_per_migration=True,
  91. target_metadata=target_metadata,
  92. compare_type=compare_type,
  93. compare_server_default=compare_server_default,
  94. include_object=include_object,
  95. render_as_batch=True,
  96. process_revision_directives=process_revision_directives,
  97. )
  98. with context.begin_transaction():
  99. context.run_migrations()
  100. if context.is_offline_mode():
  101. run_migrations_offline()
  102. else:
  103. run_migrations_online()
  104. if "airflow.www.app" in sys.modules:
  105. # Already imported, make sure we clear out any cached app
  106. from airflow.www.app import purge_cached_app
  107. purge_cached_app()