Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 1 commit into from
Aug 24, 2023
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions airflow/auth/managers/fab/security_manager/modules/db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
)
Comment on lines +260 to 264
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This exposes a problem that might not be entirely in scope of the PR. If there is a hash cache, this lookup may return a wrong user and cause a security vulnarability. We should really either add a unique constraint to registration_hash, or use one_or_none here instead; both would remove the need to add a limit.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should I keep this changes or better revert it back?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Keep it. We’re not sure how to change this code yet and adding this limit would be useful before we figure things out.


def update_user(self, user):
Expand Down
5 changes: 3 additions & 2 deletions airflow/models/dagrun.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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):
Expand Down Expand Up @@ -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))
uranusjr marked this conversation as resolved.
Show resolved Hide resolved
else:
template = session.get(LogTemplate, self.log_template_id)
if template is None:
Expand Down