Add limit 1 if required first value from query result (#33672)

apache · Aug 24, 2023 · e8ba579 · e8ba579
1 parent 370348a
commit e8ba579
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/airflow/auth/managers/fab/security_manager/modules/db.py b/airflow/auth/managers/fab/security_manager/modules/db.py
@@ -22,7 +22,7 @@
 
 from flask_appbuilder import const
 from flask_appbuilder.models.sqla import Base
-from sqlalchemy import func, inspect
+from sqlalchemy import func, inspect, select
 from sqlalchemy.exc import MultipleResultsFound
 from werkzeug.security import generate_password_hash
 
@@ -257,10 +257,10 @@ def find_user(self, username=None, email=None):
                 return None
 
     def find_register_user(self, registration_hash):
-        return (
-            self.get_session.query(self.registeruser_model)
-            .filter(self.registeruser_model.registration_hash == registration_hash)
-            .scalar()
+        return self.get_session.scalar(
+            select(self.registeruser_mode)
+            .where(self.registeruser_model.registration_hash == registration_hash)
+            .limit(1)
         )
 
     def update_user(self, user):

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
@@ -530,7 +530,7 @@ def get_previous_dagrun(
         ]
         if state is not None:
             filters.append(DagRun.state == state)
-        return session.scalar(select(DagRun).where(*filters).order_by(DagRun.execution_date.desc()))
+        return session.scalar(select(DagRun).where(*filters).order_by(DagRun.execution_date.desc()).limit(1))
 
     @provide_session
     def get_previous_scheduled_dagrun(self, session: Session = NEW_SESSION) -> DagRun | None:
@@ -543,6 +543,7 @@ def get_previous_scheduled_dagrun(self, session: Session = NEW_SESSION) -> DagRu
                 DagRun.run_type != DagRunType.MANUAL,
             )
             .order_by(DagRun.execution_date.desc())
+            .limit(1)
         )
 
     def _tis_for_dagrun_state(self, *, dag, tis):
@@ -1387,7 +1388,7 @@ def schedule_tis(
     @provide_session
     def get_log_template(self, *, session: Session = NEW_SESSION) -> LogTemplate:
         if self.log_template_id is None:  # DagRun created before LogTemplate introduction.
-            template = session.scalar(select(LogTemplate).order_by(LogTemplate.id))
+            template = session.scalar(select(LogTemplate).order_by(LogTemplate.id).limit(1))
         else:
             template = session.get(LogTemplate, self.log_template_id)
         if template is None: