mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-03-27 02:02:18 +01:00
73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
import os
|
|
import sys
|
|
|
|
import psycopg2
|
|
from sqlalchemy.orm import Session
|
|
|
|
from onyx.db.engine import get_sqlalchemy_engine
|
|
|
|
# makes it so `PYTHONPATH=.` is not required when running this script
|
|
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.append(parent_dir)
|
|
|
|
from onyx.configs.app_configs import POSTGRES_DB # noqa: E402
|
|
from onyx.configs.app_configs import POSTGRES_HOST # noqa: E402
|
|
from onyx.configs.app_configs import POSTGRES_PASSWORD # noqa: E402
|
|
from onyx.configs.app_configs import POSTGRES_PORT # noqa: E402
|
|
from onyx.configs.app_configs import POSTGRES_USER # noqa: E402
|
|
from onyx.db.credentials import create_initial_public_credential # noqa: E402
|
|
|
|
|
|
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()
|
|
|
|
# Disable triggers to prevent foreign key constraints from being checked
|
|
cur.execute("SET session_replication_role = 'replica';")
|
|
|
|
# Fetch all table names in the current database
|
|
cur.execute(
|
|
"""
|
|
SELECT tablename
|
|
FROM pg_tables
|
|
WHERE schemaname = 'public'
|
|
"""
|
|
)
|
|
|
|
tables = cur.fetchall()
|
|
|
|
for table in tables:
|
|
table_name = table[0]
|
|
|
|
# Don't touch migration history
|
|
if table_name == "alembic_version":
|
|
continue
|
|
|
|
print(f"Deleting all rows from {table_name}...")
|
|
cur.execute(f'DELETE FROM "{table_name}"')
|
|
|
|
# Re-enable triggers
|
|
cur.execute("SET session_replication_role = 'origin';")
|
|
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
print("Finished wiping all rows.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Cleaning up all Onyx tables")
|
|
wipe_all_rows(POSTGRES_DB)
|
|
with Session(get_sqlalchemy_engine(), expire_on_commit=False) as db_session:
|
|
create_initial_public_credential(db_session)
|
|
print("To keep data consistent, it's best to wipe the document index as well.")
|
|
print(
|
|
"To be safe, it's best to restart the Onyx services (API Server and Background Tasks"
|
|
)
|