This is an automated email from the ASF dual-hosted git repository.
taragolis 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 3bb1c07861 Refactor SqlAlchemy session.execute() calls to 2.0 style in
case of plain text SQL queries (#32857)
3bb1c07861 is described below
commit 3bb1c07861c7cbd3848e415c2cfc362f1334a479
Author: max <[email protected]>
AuthorDate: Fri Aug 25 08:07:52 2023 +0000
Refactor SqlAlchemy session.execute() calls to 2.0 style in case of plain
text SQL queries (#32857)
---
airflow/utils/db.py | 5 ++---
tests/utils/test_db_cleanup.py | 3 ++-
tests/utils/test_sqlalchemy.py | 13 +++++++------
3 files changed, 11 insertions(+), 10 deletions(-)
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 4429cb8dfe..a5252f4e70 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -1210,9 +1210,8 @@ def _create_table_as(
)
else:
# Postgres and SQLite both support the same "CREATE TABLE a AS SELECT
..." syntax
- session.execute(
- f"CREATE TABLE {target_table_name} AS
{source_query.selectable.compile(bind=session.get_bind())}"
- )
+ select_table = source_query.selectable.compile(bind=session.get_bind())
+ session.execute(text(f"CREATE TABLE {target_table_name} AS
{select_table}"))
def _move_dangling_data_to_new_table(
diff --git a/tests/utils/test_db_cleanup.py b/tests/utils/test_db_cleanup.py
index b6819342e7..53d6be704c 100644
--- a/tests/utils/test_db_cleanup.py
+++ b/tests/utils/test_db_cleanup.py
@@ -28,6 +28,7 @@ from uuid import uuid4
import pendulum
import pytest
from pytest import param
+from sqlalchemy import text
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.declarative import DeclarativeMeta
@@ -212,7 +213,7 @@ class TestDBCleanup:
)
stmt = CreateTableAs(target_table_name, query.selectable)
session.execute(stmt)
- res = session.execute(f"SELECT COUNT(1) FROM {target_table_name}")
+ res = session.execute(text(f"SELECT COUNT(1) FROM
{target_table_name}"))
for row in res:
assert row[0] == expected_to_delete
diff --git a/tests/utils/test_sqlalchemy.py b/tests/utils/test_sqlalchemy.py
index 84eb3c8093..76f1d41caa 100644
--- a/tests/utils/test_sqlalchemy.py
+++ b/tests/utils/test_sqlalchemy.py
@@ -26,6 +26,7 @@ from unittest.mock import MagicMock
import pytest
from kubernetes.client import models as k8s
from pytest import param
+from sqlalchemy import text
from sqlalchemy.exc import StatementError
from airflow import settings
@@ -54,7 +55,7 @@ class TestSqlAlchemyUtils:
# make sure NOT to run in UTC. Only postgres supports storing
# timezone information in the datetime field
if session.bind.dialect.name == "postgresql":
- session.execute("SET timezone='Europe/Amsterdam'")
+ session.execute(text("SET timezone='Europe/Amsterdam'"))
self.session = session
@@ -208,17 +209,17 @@ class TestSqlAlchemyUtils:
def test_prohibit_commit(self):
with prohibit_commit(self.session) as guard:
- self.session.execute("SELECT 1")
+ self.session.execute(text("SELECT 1"))
with pytest.raises(RuntimeError):
self.session.commit()
self.session.rollback()
- self.session.execute("SELECT 1")
+ self.session.execute(text("SELECT 1"))
guard.commit()
# Check the expected_commit is reset
with pytest.raises(RuntimeError):
- self.session.execute("SELECT 1")
+ self.session.execute(text("SELECT 1"))
self.session.commit()
def test_prohibit_commit_specific_session_only(self):
@@ -233,12 +234,12 @@ class TestSqlAlchemyUtils:
assert other_session is not self.session
with prohibit_commit(self.session):
- self.session.execute("SELECT 1")
+ self.session.execute(text("SELECT 1"))
with pytest.raises(RuntimeError):
self.session.commit()
self.session.rollback()
- other_session.execute("SELECT 1")
+ other_session.execute(text("SELECT 1"))
other_session.commit()
def teardown_method(self):
