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

This fixes an issue were a database session got stuck #31128

Merged
merged 4 commits into from
Jun 14, 2023
Merged

This fixes an issue were a database session got stuck #31128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9d918d6
This fixes an issue were a database session got stuck
stijndehaes May 8, 2023
11b7568
Apply suggestions from code review
stijndehaes Jun 8, 2023
9e2701c
Make comment clearer
stijndehaes Jun 8, 2023
a162752
Merge branch 'main' into bugfix/session-not-properly-closed
eladkal Jun 13, 2023
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
17 changes: 10 additions & 7 deletions airflow/lineage/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,11 @@

from airflow.configuration import conf
from airflow.lineage.backend import LineageBackend
from airflow.utils.session import create_session

if TYPE_CHECKING:
from airflow.utils.context import Context


PIPELINE_OUTLETS = "pipeline_outlets"
PIPELINE_INLETS = "pipeline_inlets"
AUTO = "auto"
Expand Down Expand Up @@ -133,13 +133,16 @@ def wrapper(self, context, *args, **kwargs):

# Remove auto and task_ids
self.inlets = [i for i in self.inlets if not isinstance(i, str)]
_inlets = self.xcom_pull(context, task_ids=task_ids, dag_id=self.dag_id, key=PIPELINE_OUTLETS)

# re-instantiate the obtained inlets
# xcom_pull returns a list of items for each given task_id
_inlets = [item for item in itertools.chain.from_iterable(_inlets)]

self.inlets.extend(_inlets)
# We manually create a session here since xcom_pull returns a LazyXComAccess iterator.
# If we do not pass a session a new session will be created, however that session will not be
# properly closed and will remain open. After we are done iterating we can safely close this
# session.
with create_session() as session:
_inlets = self.xcom_pull(
context, task_ids=task_ids, dag_id=self.dag_id, key=PIPELINE_OUTLETS, session=session
)
uranusjr marked this conversation as resolved.
Show resolved Hide resolved
Comment on lines +141 to +144
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
with create_session() as session:
_inlets = self.xcom_pull(
context, task_ids=task_ids, dag_id=self.dag_id, key=PIPELINE_OUTLETS, session=session
)
_inlets = self.xcom_pull(
context, task_ids=task_ids, dag_id=self.dag_id, key=PIPELINE_OUTLETS
)

The call to create_session is no longer needed since you decorated xcom_pull with @provide_session. The session will be automatically provided

Copy link
Contributor

Choose a reason for hiding this comment

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

Am i missing something @uranusjr ?

Copy link
Member

@uranusjr uranusjr Jun 7, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

I think both work (and using the default session from @provide_session is cleaner so you’re correct)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Removing the create_session will actually remove the bug fix. As explained in the comment above the create session: the issue is that xcom_pull returns a LazyXComAccess iterator. Since we have returned from the xcom_pull function, the session that would have been automatically created there will be closed and thus a new one that leaks will be started when we iterate. By wrapping this in a manually created session the leak is fixed

Copy link
Member

Choose a reason for hiding this comment

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

Ah I see, the point is to make the session live slightly longer so line 147 below can still use it.

Copy link
Member

@uranusjr uranusjr Jun 8, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

Maybe a comment above this block would help future readers understand the context better.

(Also I think the blank line + comments in lines 144-146 might actually hurt readability a bit since they make the two lines feel less associated. Maybe moving them above the entire block to make the create_session, xcom_pull, andinlets.extend lines stick close together would help as well.)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I remove the empty lines and 2 comments. I don't think the added much here.

There already is a comment above the create_session block to explain things, so I am not sure I correctly understand what you mean with Maybe a comment above this block would help future readers understand the context better.. Can you clarify it for me?

self.inlets.extend(i for i in itertools.chain.from_iterable(_inlets))

elif self.inlets:
raise AttributeError("inlets is not a list, operator, string or attr annotated object")
Expand Down
8 changes: 7 additions & 1 deletion airflow/models/baseoperator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1477,12 +1477,14 @@ def xcom_push(
context["ti"].xcom_push(key=key, value=value, execution_date=execution_date)

@staticmethod
@provide_session
def xcom_pull(
context: Any,
task_ids: str | list[str] | None = None,
dag_id: str | None = None,
key: str = XCOM_RETURN_KEY,
include_prior_dates: bool | None = None,
session: Session = NEW_SESSION,
stijndehaes marked this conversation as resolved.
Show resolved Hide resolved
) -> Any:
"""
Pull XComs that optionally meet certain criteria.
Expand Down Expand Up @@ -1511,7 +1513,11 @@ def xcom_pull(
are returned as well.
"""
return context["ti"].xcom_pull(
key=key, task_ids=task_ids, dag_id=dag_id, include_prior_dates=include_prior_dates
key=key,
task_ids=task_ids,
dag_id=dag_id,
include_prior_dates=include_prior_dates,
session=session,
)

@classmethod
Expand Down