Skip to content

Commit 0c069eb

Browse files
committed
Address copilot's comments
1 parent cc066bf commit 0c069eb

1 file changed

Lines changed: 9 additions & 4 deletions

File tree

  • airflow-core/src/airflow/utils

airflow-core/src/airflow/utils/db.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -816,7 +816,7 @@ def _single_connection_pool() -> Generator[None, None, None]:
816816
yield
817817
finally:
818818
os.environ.pop("AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE", None)
819-
if previous_pool_size:
819+
if previous_pool_size is not None:
820820
os.environ["AIRFLOW__DATABASE__SQL_ALCHEMY_MAX_SIZE"] = previous_pool_size
821821
settings.reconfigure_orm()
822822

@@ -1175,8 +1175,8 @@ def _run_upgradedb(config, to_revision: str | None, session: Session) -> None:
11751175
external_db_manager = RunDBManager()
11761176
external_db_manager.upgradedb(work_session)
11771177

1178-
add_default_pool_if_not_exists(session=work_session)
1179-
synchronize_log_template(session=work_session)
1178+
add_default_pool_if_not_exists(session=work_session)
1179+
synchronize_log_template(session=work_session)
11801180

11811181

11821182
@provide_session
@@ -1365,7 +1365,7 @@ def downgrade(*, to_revision, from_revision=None, show_sql_only=False, session:
13651365
revision_range = f"{from_revision}:{to_revision}"
13661366
_offline_migration(command.downgrade, config=config, revision=revision_range)
13671367
else:
1368-
dialect_label = " (MySQL)" if get_dialect_name(session) == "mysql" else ""
1368+
dialect_label = " (MySQL)" if get_dialect_name(work_session) == "mysql" else ""
13691369
log.info("Applying downgrade migrations to Airflow database%s.", dialect_label)
13701370
command.downgrade(config, revision=to_revision, sql=show_sql_only)
13711371

@@ -1538,8 +1538,10 @@ def _create_global_lock_postgresql(
15381538
bind = session.get_bind()
15391539
if hasattr(bind, "connect"):
15401540
conn = bind.connect()
1541+
owns_connection = True
15411542
else:
15421543
conn = bind
1544+
owns_connection = False
15431545

15441546
try:
15451547
if _USE_PSYCOPG3:
@@ -1567,6 +1569,9 @@ def _create_global_lock_postgresql(
15671569
if not unlocked:
15681570
raise RuntimeError("Error releasing DB lock!")
15691571

1572+
if owns_connection:
1573+
conn.close()
1574+
15701575

15711576
@contextlib.contextmanager
15721577
def create_global_lock(

0 commit comments

Comments
 (0)