From 1f7f5384fb744546c70ca81159075665651a3ce3 Mon Sep 17 00:00:00 2001 From: RT Date: Sun, 21 Jan 2024 01:56:46 -0500 Subject: [PATCH] fix: Remove the hard-coded database URI in the test helper to use an env var (#865) --- tests/clear_postgres_db.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/clear_postgres_db.py b/tests/clear_postgres_db.py index 60618363..457beea7 100644 --- a/tests/clear_postgres_db.py +++ b/tests/clear_postgres_db.py @@ -1,8 +1,19 @@ +import os + from sqlalchemy import create_engine, MetaData -engine = create_engine("postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt") -meta = MetaData() -meta.reflect(bind=engine) +def main(): + uri = os.environ.get( + "PGVECTOR_TEST_DB_URL", + "postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt", + ) -meta.drop_all(bind=engine) + engine = create_engine(uri) + meta = MetaData() + meta.reflect(bind=engine) + meta.drop_all(bind=engine) + + +if __name__ == "__main__": + main()