102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Clean the test database by removing all data while preserving the schema.
|
|
|
|
Usage:
|
|
poetry run python clean_test_db.py [--yes]
|
|
|
|
Options:
|
|
--yes Skip confirmation prompt
|
|
"""
|
|
|
|
import asyncio
|
|
import sys
|
|
|
|
from prisma import Prisma
|
|
|
|
|
|
async def main():
|
|
db = Prisma()
|
|
await db.connect()
|
|
|
|
print("=" * 60)
|
|
print("Cleaning Test Database")
|
|
print("=" * 60)
|
|
print()
|
|
|
|
# Get initial counts
|
|
user_count = await db.user.count()
|
|
agent_count = await db.agentgraph.count()
|
|
|
|
print(f"Current data: {user_count} users, {agent_count} agent graphs")
|
|
|
|
if user_count == 0 and agent_count == 0:
|
|
print("Database is already clean!")
|
|
await db.disconnect()
|
|
return
|
|
|
|
# Check for --yes flag
|
|
skip_confirm = "--yes" in sys.argv
|
|
|
|
if not skip_confirm:
|
|
response = input("\nDo you want to clean all data? (yes/no): ")
|
|
if response.lower() != "yes":
|
|
print("Aborted.")
|
|
await db.disconnect()
|
|
return
|
|
|
|
print("\nCleaning database...")
|
|
|
|
# Delete in reverse order of dependencies
|
|
tables = [
|
|
("UserNotificationBatch", db.usernotificationbatch),
|
|
("NotificationEvent", db.notificationevent),
|
|
("CreditRefundRequest", db.creditrefundrequest),
|
|
("StoreListingReview", db.storelistingreview),
|
|
("StoreListingVersion", db.storelistingversion),
|
|
("StoreListing", db.storelisting),
|
|
("AgentNodeExecutionInputOutput", db.agentnodeexecutioninputoutput),
|
|
("AgentNodeExecution", db.agentnodeexecution),
|
|
("AgentGraphExecution", db.agentgraphexecution),
|
|
("AgentNodeLink", db.agentnodelink),
|
|
("LibraryAgent", db.libraryagent),
|
|
("AgentPreset", db.agentpreset),
|
|
("IntegrationWebhook", db.integrationwebhook),
|
|
("AgentNode", db.agentnode),
|
|
("AgentGraph", db.agentgraph),
|
|
("AgentBlock", db.agentblock),
|
|
("APIKey", db.apikey),
|
|
("CreditTransaction", db.credittransaction),
|
|
("AnalyticsMetrics", db.analyticsmetrics),
|
|
("AnalyticsDetails", db.analyticsdetails),
|
|
("Profile", db.profile),
|
|
("UserOnboarding", db.useronboarding),
|
|
("User", db.user),
|
|
]
|
|
|
|
for table_name, table in tables:
|
|
try:
|
|
count = await table.count()
|
|
if count > 0:
|
|
await table.delete_many()
|
|
print(f"✓ Deleted {count} records from {table_name}")
|
|
except Exception as e:
|
|
print(f"⚠ Error cleaning {table_name}: {e}")
|
|
|
|
# Refresh materialized views (they should be empty now)
|
|
try:
|
|
await db.execute_raw("SELECT refresh_store_materialized_views();")
|
|
print("\n✓ Refreshed materialized views")
|
|
except Exception as e:
|
|
print(f"\n⚠ Could not refresh materialized views: {e}")
|
|
|
|
await db.disconnect()
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Database cleaned successfully!")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|