danswer/backend/scripts/reset_postgres.py
Yuhong Sun 6fe54a4eed
DAN-115 Document Polling (#91)
Includes updated document counting for polling
2023-06-15 21:07:05 -07:00

43 lines
1.1 KiB
Python

import psycopg2
from danswer.configs.app_configs import POSTGRES_DB
from danswer.configs.app_configs import POSTGRES_HOST
from danswer.configs.app_configs import POSTGRES_PASSWORD
from danswer.configs.app_configs import POSTGRES_PORT
from danswer.configs.app_configs import POSTGRES_USER
def wipe_all_rows(database: str) -> None:
conn = psycopg2.connect(
dbname=database,
user=POSTGRES_USER,
password=POSTGRES_PASSWORD,
host=POSTGRES_HOST,
port=POSTGRES_PORT,
)
cur = conn.cursor()
cur.execute(
"""
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_type = 'BASE TABLE'
"""
)
table_names = cur.fetchall()
for table_name in table_names:
if table_name[0] == "alembic_version":
continue
cur.execute(f'DELETE FROM "{table_name[0]}"')
print(f"Deleted all rows from table {table_name[0]}")
conn.commit()
cur.close()
conn.close()
if __name__ == "__main__":
wipe_all_rows(POSTGRES_DB)