[airflow] branch main updated: Don't show alembic info logs at the start of every cli command (#21758)

[as...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 3c12c2e  Don't show alembic info logs at the start of every cli command (#21758)
3c12c2e is described below

commit 3c12c2e1e5c2d1d961addbe0452186d32800135e
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Wed Feb 23 14:51:23 2022 +0000

    Don't show alembic info logs at the start of every cli command (#21758)
    
    We recently merged a change where we check if DB migrations
    are pending before running the main command, but this had the
    side-effect of showing these two log lines from alembic:
    
    ```
    [2022-02-22 18:06:01,995] {{migration.py:201}} INFO - Context impl PostgresqlImpl.
    [2022-02-22 18:06:01,995] {{migration.py:204}} INFO - Will assume transactional DDL.
    ```
    
    Which is a) not useful information to a user, and b) "pollutes" the
    output if a command was producing JSON or some other structured format
---
 airflow/utils/db.py | 37 ++++++++++++++++++-------------------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index d452e23..3bd98f1 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -654,19 +654,12 @@ def check_migrations(timeout):
     :param timeout: Timeout for the migration in seconds
     :return: None
     """
-    from alembic.runtime.environment import EnvironmentContext
-
-    script_, config = _get_script_dir_and_config()
-    with EnvironmentContext(
-        config,
-        script_,
-    ) as env, settings.engine.connect() as connection:
-        env.configure(connection)
+    with _configured_alembic_environment() as env:
         context = env.get_context()
         source_heads = None
         db_heads = None
         for ticker in range(timeout):
-            source_heads = set(script_.get_heads())
+            source_heads = set(env.script.get_heads())
             db_heads = set(context.get_current_heads())
             if source_heads == db_heads:
                 return
@@ -678,27 +671,33 @@ def check_migrations(timeout):
         )
 
 
-def _get_script_dir_and_config():
-    """Get config and script directory"""
+@contextlib.contextmanager
+def _configured_alembic_environment():
+    from alembic.runtime.environment import EnvironmentContext
     from alembic.script import ScriptDirectory
 
     config = _get_alembic_config()
     script_ = ScriptDirectory.from_config(config)
-    return script_, config
-
 
-def check_and_run_migrations():
-    """Check and run migrations if necessary. Only use in a tty"""
-    from alembic.runtime.environment import EnvironmentContext
-
-    script_, config = _get_script_dir_and_config()
     with EnvironmentContext(
         config,
         script_,
     ) as env, settings.engine.connect() as connection:
+
+        alembic_logger = logging.getLogger('alembic')
+        level = alembic_logger.level
+        alembic_logger.setLevel(logging.WARNING)
         env.configure(connection)
+        alembic_logger.setLevel(level)
+
+        yield env
+
+
+def check_and_run_migrations():
+    """Check and run migrations if necessary. Only use in a tty"""
+    with _configured_alembic_environment() as env:
         context = env.get_context()
-        source_heads = set(script_.get_heads())
+        source_heads = set(env.script.get_heads())
         db_heads = set(context.get_current_heads())
         db_command = None
         command_name = None