From c7e56fbce6111b1af2812905e85fd11ee760402d Mon Sep 17 00:00:00 2001 From: Hongbo Wu Date: Sat, 23 Sep 2023 11:55:22 +0800 Subject: [PATCH] add PG_TIMEOUT to env var --- packages/db/elastic_migrations/migrate_from_elastic.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/db/elastic_migrations/migrate_from_elastic.py b/packages/db/elastic_migrations/migrate_from_elastic.py index 42f8bb3f6..bfbd3e034 100755 --- a/packages/db/elastic_migrations/migrate_from_elastic.py +++ b/packages/db/elastic_migrations/migrate_from_elastic.py @@ -15,6 +15,7 @@ PG_USER = os.getenv('PG_USER', 'app_user') PG_PASSWORD = os.getenv('PG_PASSWORD', 'app_pass') PG_DB = os.getenv('PG_DB', 'omnivore') PG_COOLDOWN_TIME = os.getenv('PG_COOLDOWN_TIME', 1) +PG_TIMEOUT = os.getenv('PG_TIMEOUT', 60) ES_URL = os.getenv('ES_URL', 'http://localhost:9200') ES_USERNAME = os.getenv('ES_USERNAME', 'elastic') ES_PASSWORD = os.getenv('ES_PASSWORD', 'password') @@ -226,7 +227,7 @@ async def insert_recommendations(db_conn, recommendations): async def insert_into_postgres(insert_query, db_conn, records): - await db_conn.executemany(insert_query, records, timeout=60) + await db_conn.executemany(insert_query, records, timeout=int(PG_TIMEOUT)) # cool down for PG_COOLDOWN_TIME seconds if PG_COOLDOWN_TIME > 0: await asyncio.sleep(float(PG_COOLDOWN_TIME)) @@ -243,7 +244,8 @@ async def main(): # postgres connection db_conn = await asyncpg.connect(user=PG_USER, password=PG_PASSWORD, - database=PG_DB, host=PG_HOST, port=PG_PORT) + database=PG_DB, host=PG_HOST, port=PG_PORT, + timeout=int(PG_TIMEOUT)) # elastic client es_client = AsyncElasticsearch(ES_URL, http_auth=(