Skip to content

Commit

Permalink
fix(airflow_db_cleanup): sessions cleanup (#11035)
Browse files Browse the repository at this point in the history
This fixes an issue where the python operator gets killed if the ssession table is large neough
  • Loading branch information
mokshasoul committed Jan 17, 2024
1 parent 980a984 commit 196be7a
Showing 1 changed file with 6 additions and 17 deletions.
23 changes: 6 additions & 17 deletions composer/workflows/airflow_db_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -468,24 +468,13 @@ def cleanup_sessions():

try:
logging.info("Deleting sessions...")
before = len(
session.execute(
text("SELECT * FROM session WHERE expiry < now()::timestamp(0);")
)
.mappings()
.all()
)
count_statement = "SELECT COUNT(*) AS cnt FROM session WHERE expiry < now()::timestamp(0);"
before = session.execute(text(count_statement)).one_or_none()["cnt"]
session.execute(text("DELETE FROM session WHERE expiry < now()::timestamp(0);"))
after = len(
session.execute(
text("SELECT * FROM session WHERE expiry < now()::timestamp(0);")
)
.mappings()
.all()
)
logging.info("Deleted {} expired sessions.".format(before - after))
except Exception as e:
logging.error(e)
after = session.execute(text(count_statement)).one_or_none()["cnt"]
logging.info("Deleted %s expired sessions.", (before - after))
except Exception as err:
logging.exception(err)

session.commit()
session.close()
Expand Down

0 comments on commit 196be7a

Please sign in to comment.