diff --git a/api/chalicelib/core/sessions.py b/api/chalicelib/core/sessions.py index 7543f9c9d..5f9a6d3ab 100644 --- a/api/chalicelib/core/sessions.py +++ b/api/chalicelib/core/sessions.py @@ -1199,7 +1199,7 @@ def get_session_ids_by_user_ids(project_id, user_ids): def delete_sessions_by_session_ids(session_ids): - with pg_client.PostgresClient(long_query=True) as cur: + with pg_client.PostgresClient(unlimited_query=True) as cur: query = cur.mogrify( """\ DELETE FROM public.sessions @@ -1213,7 +1213,7 @@ def delete_sessions_by_session_ids(session_ids): def delete_sessions_by_user_ids(project_id, user_ids): - with pg_client.PostgresClient(long_query=True) as cur: + with pg_client.PostgresClient(unlimited_query=True) as cur: query = cur.mogrify( """\ DELETE FROM public.sessions @@ -1227,6 +1227,6 @@ def delete_sessions_by_user_ids(project_id, user_ids): def count_all(): - with pg_client.PostgresClient(long_query=True) as cur: + with pg_client.PostgresClient(unlimited_query=True) as cur: row = cur.execute(query="SELECT COUNT(session_id) AS count FROM public.sessions") return row.get("count", 0) diff --git a/api/chalicelib/utils/pg_client.py b/api/chalicelib/utils/pg_client.py index 1c4625873..2abc9f6c7 100644 --- a/api/chalicelib/utils/pg_client.py +++ b/api/chalicelib/utils/pg_client.py @@ -76,12 +76,17 @@ class PostgresClient: cursor = None long_query = False - def __init__(self, long_query=False): + def __init__(self, long_query=False, unlimited_query=False): self.long_query = long_query - if long_query: + if unlimited_query: + long_config = dict(_PG_CONFIG) + long_config["application_name"] += "-UNLIMITED" + self.connection = psycopg2.connect(**long_config) + elif long_query: long_config = dict(_PG_CONFIG) long_config["application_name"] += "-LONG" - self.connection = psycopg2.connect(**_PG_CONFIG) + long_config["options"] = f"-c statement_timeout={config('pg_long_timeout', cast=int, default=5*60) * 1000}" + self.connection = psycopg2.connect(**long_config) else: self.connection = postgreSQL_pool.getconn()