fix(airflow_db_cleanup): sessions cleanup (GoogleCloudPlatform#11035)

mokshasoul · web-flow · commit 196be7a900be · 2024-01-17T12:09:20.000-08:00
This fixes an issue where the python operator gets killed if the ssession table is large neough
diff --git a/composer/workflows/airflow_db_cleanup.py b/composer/workflows/airflow_db_cleanup.py
@@ -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()
