From c19c76c3ad6f658cd2526cf19746e8b20c3affe5 Mon Sep 17 00:00:00 2001 From: Chris Weaver Date: Mon, 13 Oct 2025 14:19:02 -0700 Subject: [PATCH] feat: tenant cleanup (#5703) --- backend/onyx/document_index/vespa/index.py | 8 +- backend/scripts/tenant_cleanup/README.md | 59 ++ .../tenant_cleanup/analyze_current_tenants.py | 454 ++++++++++++ .../scripts/tenant_cleanup/cleanup_tenants.py | 693 ++++++++++++++++++ .../scripts/tenant_cleanup/cleanup_utils.py | 194 +++++ .../mark_connectors_for_deletion.py | 465 ++++++++++++ .../on_pod_scripts/check_documents_deleted.py | 107 +++ .../on_pod_scripts/cleanup_tenant_schema.py | 91 +++ .../on_pod_scripts/get_tenant_connectors.py | 109 +++ .../on_pod_scripts/get_tenant_index_name.py | 59 ++ .../on_pod_scripts/get_tenant_users.py | 78 ++ .../mark_connector_for_deletion.py | 144 ++++ .../on_pod_scripts/understand_tenants.py | 96 +++ 13 files changed, 2555 insertions(+), 2 deletions(-) create mode 100644 backend/scripts/tenant_cleanup/README.md create mode 100644 backend/scripts/tenant_cleanup/analyze_current_tenants.py create mode 100644 backend/scripts/tenant_cleanup/cleanup_tenants.py create mode 100644 backend/scripts/tenant_cleanup/cleanup_utils.py create mode 100644 backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/mark_connector_for_deletion.py create mode 100644 backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py diff --git a/backend/onyx/document_index/vespa/index.py b/backend/onyx/document_index/vespa/index.py index 46f8ec1bcd..99d57ed76a 100644 --- a/backend/onyx/document_index/vespa/index.py +++ b/backend/onyx/document_index/vespa/index.py @@ -1062,7 +1062,7 @@ class VespaIndex(DocumentIndex): *, tenant_id: str, index_name: str, - ) -> None: + ) -> int: """ Deletes all entries in the specified index with the given tenant_id. @@ -1072,6 +1072,9 @@ class VespaIndex(DocumentIndex): Parameters: tenant_id (str): The tenant ID whose documents are to be deleted. index_name (str): The name of the index from which to delete documents. + + Returns: + int: The number of documents deleted. """ logger.info( f"Deleting entries with tenant_id: {tenant_id} from index: {index_name}" @@ -1084,7 +1087,7 @@ class VespaIndex(DocumentIndex): logger.info( f"No documents found with tenant_id: {tenant_id} in index: {index_name}" ) - return + return 0 # Step 2: Delete documents in batches delete_requests = [ @@ -1093,6 +1096,7 @@ class VespaIndex(DocumentIndex): ] cls._apply_deletes_batched(delete_requests) + return len(document_ids) @classmethod def _get_all_document_ids_by_tenant_id( diff --git a/backend/scripts/tenant_cleanup/README.md b/backend/scripts/tenant_cleanup/README.md new file mode 100644 index 0000000000..fe515ff416 --- /dev/null +++ b/backend/scripts/tenant_cleanup/README.md @@ -0,0 +1,59 @@ +## How to Tenant Cleanup + +Three main steps. + +### Build a list of tenants to cleanup + +Use the `analyze_current_tenants.py` script: + +``` +PYTHONPATH=. \ +CONTROL_PLANE_RDS_HOST= \ +CONTROL_PLANE_RDS_PASSWORD= \ +BASTION_HOST= \ +PEM_FILE_LOCATION= \ +python scripts/tenant_cleanup/analyze_current_tenants.py +``` + +This will create a `.csv` called something like `gated_tenants_no_query_3mo_20251012_161102.csv` in the `backend` dir. + + +### Delete all documents within these tenants + +Use the `mark_connectors_for_deletion.py` script: + +``` +PYTHONPATH=. \ +CONTROL_PLANE_RDS_HOST= \ +CONTROL_PLANE_RDS_PASSWORD= \ +BASTION_HOST= \ +PEM_FILE_LOCATION= \ +python scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo_.csv --force +``` + +Replace `gated_tenants_no_query_3mo_.csv` with the CSV name from step (1). + +This will update the data plane database to 1/ cancel all index attempts 2/ mark all connectors as up for deletion. +We now need to wait for the deletion to run. + +It's done this way to re-use as much of the existing code + take advantage of existing infra for parallelized, long running jobs. These +deletion jobs can take a LONG time (>6hrs), so having it performed syncronously by a script is not really tenable. + + +### Cleanup the tenants + +Use the `cleanup_tenants.py` script: + +``` +PYTHONPATH=. \ +CONTROL_PLANE_RDS_HOST= \ +CONTROL_PLANE_RDS_PASSWORD= \ +BASTION_HOST= \ +PEM_FILE_LOCATION= \ +python scripts/tenant_cleanup/cleanup_tenants.py --csv gated_tenants_no_query_3mo_.csv --force +``` + +This will drop the tenant schema from the data plane DB, cleanup the `user_tenant_mapping` table, and +clean up any control plane DB tables associated with each tenant. + +NOTE: if the previous step has not completed, tenants with documents will throw an exception. diff --git a/backend/scripts/tenant_cleanup/analyze_current_tenants.py b/backend/scripts/tenant_cleanup/analyze_current_tenants.py new file mode 100644 index 0000000000..86e40d1708 --- /dev/null +++ b/backend/scripts/tenant_cleanup/analyze_current_tenants.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +""" +Full tenant analysis script that: +1. Finds a heavy worker pod +2. Runs the tenant data collection script on the pod +3. Analyzes the collected data +""" + +import csv +import json +import os +import subprocess +import sys +from datetime import datetime +from datetime import timedelta +from datetime import timezone +from pathlib import Path +from typing import Any + +from scripts.tenant_cleanup.cleanup_utils import find_worker_pod + + +def collect_tenant_data(pod_name: str) -> list[dict[str, Any]]: + """Run the understand_tenants script on the pod and return the data.""" + print(f"\nCollecting tenant data from pod {pod_name}...") + + # Get the path to the understand_tenants script + script_dir = Path(__file__).parent + understand_tenants_script = script_dir / "on_pod_scripts" / "understand_tenants.py" + + if not understand_tenants_script.exists(): + raise FileNotFoundError( + f"understand_tenants.py not found at {understand_tenants_script}" + ) + + # Copy script to pod + print("Copying script to pod...") + subprocess.run( + [ + "kubectl", + "cp", + str(understand_tenants_script), + f"{pod_name}:/tmp/understand_tenants.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + print("Executing script on pod (this may take a while)...") + result = subprocess.run( + ["kubectl", "exec", pod_name, "--", "python", "/tmp/understand_tenants.py"], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(result.stderr, file=sys.stderr) + + # Parse JSON from stdout + try: + tenant_data = json.loads(result.stdout) + print(f"Successfully collected data for {len(tenant_data)} tenants") + return tenant_data + except json.JSONDecodeError as e: + print(f"Failed to parse JSON output: {e}", file=sys.stderr) + print(f"stdout: {result.stdout[:500]}", file=sys.stderr) + raise + + +def collect_control_plane_data() -> list[dict[str, Any]]: + """Collect control plane data from the control plane database.""" + print("\nCollecting control plane data...") + + rds_host = os.environ.get("CONTROL_PLANE_RDS_HOST") + if not rds_host: + raise ValueError("CONTROL_PLANE_RDS_HOST is not set") + + rds_password = os.environ.get("CONTROL_PLANE_RDS_PASSWORD") + if not rds_password: + raise ValueError("CONTROL_PLANE_RDS_PASSWORD is not set") + + db_url = f"postgresql://postgres:{rds_password}@{rds_host}:5432/control" + + bastion_host = os.environ.get("BASTION_HOST") + if not bastion_host: + raise ValueError("BASTION_HOST is not set") + + pem_file_location = os.environ.get("PEM_FILE_LOCATION") + if not pem_file_location: + raise ValueError("PEM_FILE_LOCATION is not set") + + full_cmd = ( + f"ssh -i {pem_file_location} ec2-user@{bastion_host} " + f"\"psql {db_url} -c '\\copy (SELECT * FROM tenant) " + f"to '/tmp/control_plane_data.csv' with (format csv);'\"" + ) + + result = subprocess.run( + full_cmd, + shell=True, + check=True, + capture_output=True, + text=True, + ) + + # Copy the CSV file from the bastion to local machine + copy_cmd = f"scp -i {pem_file_location} ec2-user@{bastion_host}:/tmp/control_plane_data.csv ." + + copy_result = subprocess.run( + copy_cmd, shell=True, check=True, capture_output=True, text=True + ) + + if copy_result.stderr: + print(f"Copy warnings: {copy_result.stderr}", file=sys.stderr) + raise RuntimeError( + "Failed to copy control plane data from bastion to local machine" + ) + + print("Control plane data copied to local machine as control_plane_data.csv") + + print(result.stdout) + + # Read the CSV file and convert to list of dictionaries + control_plane_data = [] + with open("control_plane_data.csv", "r", newline="", encoding="utf-8") as csvfile: + reader = csv.DictReader(csvfile) + for row in reader: + control_plane_data.append(row) + + return control_plane_data + + +def analyze_tenants( + tenants: list[dict[str, Any]], control_plane_data: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Analyze tenant activity data and return gated tenants with no query in last 3 months.""" + + print(f"\n{'=' * 80}") + print(f"TENANT ANALYSIS REPORT - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") + print(f"{'=' * 80}") + print(f"Total tenants analyzed: {len(tenants)}\n") + + # Create a lookup dict for control plane data by tenant_id + control_plane_lookup = {} + for row in control_plane_data: + # CSV has no header, columns are: tenant_id, stripe_customer_id, created_at, + # stripe_subscription_quantity, contact_email, registration_origin, tenant_status + if len(row) >= 7: + tenant_id = list(row.values())[0] # First column is tenant_id + tenant_status = list(row.values())[6] # 7th column is tenant_status + control_plane_lookup[tenant_id] = tenant_status + + # Calculate cutoff dates + one_month_cutoff = datetime.now(timezone.utc) - timedelta(days=30) + three_month_cutoff = datetime.now(timezone.utc) - timedelta(days=90) + + # Categorize tenants into 4 groups + gated_no_query_3_months = [] # GATED_ACCESS + no query in last 3 months + gated_query_1_3_months = [] # GATED_ACCESS + query between 1-3 months + gated_query_1_month = [] # GATED_ACCESS + query in last 1 month + everyone_else = [] # All other tenants + + for tenant in tenants: + tenant_id = tenant.get("tenant_id") + last_query_time = tenant.get("last_query_time") + tenant_status = control_plane_lookup.get(tenant_id, "UNKNOWN") + + is_gated = tenant_status == "GATED_ACCESS" + + # Parse last query time + if last_query_time: + query_time = datetime.fromisoformat(last_query_time.replace("Z", "+00:00")) + else: + query_time = None + + # Categorize + if is_gated: + if query_time is None or query_time <= three_month_cutoff: + gated_no_query_3_months.append(tenant) + elif query_time <= one_month_cutoff: + gated_query_1_3_months.append(tenant) + else: # query_time > one_month_cutoff + gated_query_1_month.append(tenant) + else: + everyone_else.append(tenant) + + # Calculate document counts for each group + gated_no_query_docs = sum( + t.get("num_documents", 0) for t in gated_no_query_3_months + ) + gated_1_3_month_docs = sum( + t.get("num_documents", 0) for t in gated_query_1_3_months + ) + gated_1_month_docs = sum(t.get("num_documents", 0) for t in gated_query_1_month) + everyone_else_docs = sum(t.get("num_documents", 0) for t in everyone_else) + + print("=" * 80) + print("TENANT CATEGORIZATION BY GATED ACCESS STATUS AND ACTIVITY") + print("=" * 80) + + print("\n1. GATED_ACCESS + No query in last 3 months:") + print(f" Count: {len(gated_no_query_3_months):,}") + print(f" Total documents: {gated_no_query_docs:,}") + print( + f" Avg documents per tenant: {gated_no_query_docs / len(gated_no_query_3_months) if gated_no_query_3_months else 0:.2f}" + ) + + print("\n2. GATED_ACCESS + Query between 1-3 months ago:") + print(f" Count: {len(gated_query_1_3_months):,}") + print(f" Total documents: {gated_1_3_month_docs:,}") + print( + f" Avg documents per tenant: {gated_1_3_month_docs / len(gated_query_1_3_months) if gated_query_1_3_months else 0:.2f}" + ) + + print("\n3. GATED_ACCESS + Query in last 1 month:") + print(f" Count: {len(gated_query_1_month):,}") + print(f" Total documents: {gated_1_month_docs:,}") + print( + f" Avg documents per tenant: {gated_1_month_docs / len(gated_query_1_month) if gated_query_1_month else 0:.2f}" + ) + + print("\n4. Everyone else (non-GATED_ACCESS):") + print(f" Count: {len(everyone_else):,}") + print(f" Total documents: {everyone_else_docs:,}") + print( + f" Avg documents per tenant: {everyone_else_docs / len(everyone_else) if everyone_else else 0:.2f}" + ) + + total_docs = ( + gated_no_query_docs + + gated_1_3_month_docs + + gated_1_month_docs + + everyone_else_docs + ) + print(f"\nTotal documents across all tenants: {total_docs:,}") + + # Top 100 tenants by document count + print("\n" + "=" * 80) + print("TOP 100 TENANTS BY DOCUMENT COUNT") + print("=" * 80) + + # Sort all tenants by document count + sorted_tenants = sorted( + tenants, key=lambda t: t.get("num_documents", 0), reverse=True + ) + + top_100 = sorted_tenants[:100] + + print( + f"\n{'Rank':<6} {'Tenant ID':<45} {'Documents':>12} {'Users':>8} {'Last Query':<12} {'Group'}" + ) + print("-" * 130) + + for idx, tenant in enumerate(top_100, 1): + tenant_id = tenant.get("tenant_id", "Unknown") + num_docs = tenant.get("num_documents", 0) + num_users = tenant.get("num_users", 0) + last_query = tenant.get("last_query_time", "Never") + tenant_status = control_plane_lookup.get(tenant_id, "UNKNOWN") + + # Format the last query time + if last_query and last_query != "Never": + try: + query_dt = datetime.fromisoformat(last_query.replace("Z", "+00:00")) + last_query_str = query_dt.strftime("%Y-%m-%d") + except Exception: + last_query_str = last_query[:10] if len(last_query) > 10 else last_query + else: + last_query_str = "Never" + + # Determine group + if tenant_status == "GATED_ACCESS": + if last_query and last_query != "Never": + query_time = datetime.fromisoformat(last_query.replace("Z", "+00:00")) + if query_time <= three_month_cutoff: + group = "Gated - No query (3mo)" + elif query_time <= one_month_cutoff: + group = "Gated - Query (1-3mo)" + else: + group = "Gated - Query (1mo)" + else: + group = "Gated - No query (3mo)" + else: + group = f"Other ({tenant_status})" + + print( + f"{idx:<6} {tenant_id:<45} {num_docs:>12,} {num_users:>8} {last_query_str:<12} {group}" + ) + + # Summary stats for top 100 + top_100_docs = sum(t.get("num_documents", 0) for t in top_100) + + print("\n" + "-" * 110) + print(f"Top 100 total documents: {top_100_docs:,}") + print( + f"Percentage of all documents: {(top_100_docs / total_docs * 100) if total_docs > 0 else 0:.2f}%" + ) + + # Additional insights + print("\n" + "=" * 80) + print("ADDITIONAL INSIGHTS") + print("=" * 80) + + # Tenants with no documents + no_docs = [t for t in tenants if t.get("num_documents", 0) == 0] + print( + f"\nTenants with 0 documents: {len(no_docs):,} ({len(no_docs) / len(tenants) * 100:.2f}%)" + ) + + # Tenants with no users + no_users = [t for t in tenants if t.get("num_users", 0) == 0] + print( + f"Tenants with 0 users: {len(no_users):,} ({len(no_users) / len(tenants) * 100:.2f}%)" + ) + + # Document distribution quartiles + doc_counts = sorted([t.get("num_documents", 0) for t in tenants]) + if doc_counts: + print("\nDocument count distribution:") + print(f" Median: {doc_counts[len(doc_counts) // 2]:,}") + print(f" 75th percentile: {doc_counts[int(len(doc_counts) * 0.75)]:,}") + print(f" 90th percentile: {doc_counts[int(len(doc_counts) * 0.90)]:,}") + print(f" 95th percentile: {doc_counts[int(len(doc_counts) * 0.95)]:,}") + print(f" 99th percentile: {doc_counts[int(len(doc_counts) * 0.99)]:,}") + print(f" Max: {doc_counts[-1]:,}") + + return gated_no_query_3_months + + +def find_recent_tenant_data() -> tuple[list[dict[str, Any]] | None, str | None]: + """Find the most recent tenant data file if it's less than 7 days old.""" + current_dir = Path.cwd() + tenant_data_files = list(current_dir.glob("tenant_data_*.json")) + + if not tenant_data_files: + return None, None + + # Sort by modification time, most recent first + tenant_data_files.sort(key=lambda p: p.stat().st_mtime, reverse=True) + most_recent = tenant_data_files[0] + + # Check if file is less than 7 days old + file_age = datetime.now().timestamp() - most_recent.stat().st_mtime + seven_days_in_seconds = 7 * 24 * 60 * 60 + + if file_age < seven_days_in_seconds: + file_age_days = file_age / (24 * 60 * 60) + print( + f"\n✓ Found recent tenant data: {most_recent.name} (age: {file_age_days:.1f} days)" + ) + + with open(most_recent, "r") as f: + tenant_data = json.load(f) + + return tenant_data, str(most_recent) + + return None, None + + +def main() -> None: + try: + # Step 0: Collect control plane data + control_plane_data = collect_control_plane_data() + + # Step 1: Check for recent tenant data (< 7 days old) + tenant_data, cached_file = find_recent_tenant_data() + + if tenant_data: + print(f"Using cached tenant data from: {cached_file}") + print(f"Total tenants in cache: {len(tenant_data)}") + else: + # Step 2a: Find the heavy worker pod + pod_name = find_worker_pod() + + # Step 2b: Collect tenant data + tenant_data = collect_tenant_data(pod_name) + + # Step 2c: Save raw data to file with timestamp + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + output_file = f"tenant_data_{timestamp}.json" + with open(output_file, "w") as f: + json.dump(tenant_data, f, indent=2, default=str) + print(f"\n✓ Raw data saved to: {output_file}") + + # Step 3: Analyze the data and get gated tenants without recent queries + gated_no_query_3_months = analyze_tenants(tenant_data, control_plane_data) + + # Step 4: Export to CSV (sorted by num_documents descending) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + csv_file = f"gated_tenants_no_query_3mo_{timestamp}.csv" + + # Sort by num_documents in descending order + sorted_tenants = sorted( + gated_no_query_3_months, + key=lambda t: t.get("num_documents", 0), + reverse=True, + ) + + with open(csv_file, "w", newline="", encoding="utf-8") as csvfile: + fieldnames = [ + "tenant_id", + "num_documents", + "num_users", + "last_query_time", + "days_since_last_query", + ] + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + writer.writeheader() + + now = datetime.now(timezone.utc) + for tenant in sorted_tenants: + # Calculate days since last query + last_query_time = tenant.get("last_query_time") + if last_query_time: + try: + query_dt = datetime.fromisoformat( + last_query_time.replace("Z", "+00:00") + ) + days_since = str((now - query_dt).days) + except Exception: + days_since = "N/A" + else: + days_since = "Never" + + writer.writerow( + { + "tenant_id": tenant.get("tenant_id", ""), + "num_documents": tenant.get("num_documents", 0), + "num_users": tenant.get("num_users", 0), + "last_query_time": last_query_time or "Never", + "days_since_last_query": days_since, + } + ) + + print(f"\n✓ CSV exported to: {csv_file}") + print( + f" Total gated tenants with no query in last 3 months: {len(gated_no_query_3_months)}" + ) + + except subprocess.CalledProcessError as e: + print(f"Error running command: {e}", file=sys.stderr) + if e.stderr: + print(f"stderr: {e.stderr}", file=sys.stderr) + sys.exit(1) + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/cleanup_tenants.py b/backend/scripts/tenant_cleanup/cleanup_tenants.py new file mode 100644 index 0000000000..c1277db55c --- /dev/null +++ b/backend/scripts/tenant_cleanup/cleanup_tenants.py @@ -0,0 +1,693 @@ +#!/usr/bin/env python3 +""" +Tenant cleanup script that: +1. Deletes all documents from Vespa +2. Drops the data plane PostgreSQL schema +3. Clean up control plane (tenants, subscription table) + +Usage: + python backend/scripts/cleanup_tenant.py [--force] + python backend/scripts/cleanup_tenant.py --csv [--force] + +Arguments: + tenant_id The tenant ID to clean up (required if not using --csv) + --csv PATH Path to CSV file containing tenant IDs to clean up + --force Skip all confirmation prompts (optional) + +Examples: + python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789 + python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789 --force + python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv + python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv --force +""" + +import json +import signal +import subprocess +import sys +from pathlib import Path + +from scripts.tenant_cleanup.cleanup_utils import confirm_step +from scripts.tenant_cleanup.cleanup_utils import execute_control_plane_query +from scripts.tenant_cleanup.cleanup_utils import find_worker_pod +from scripts.tenant_cleanup.cleanup_utils import get_tenant_status +from scripts.tenant_cleanup.cleanup_utils import read_tenant_ids_from_csv + + +# Global variable to track active subprocess +_active_process = None + + +def signal_handler(signum: int, frame: object) -> None: + """Handle termination signals by killing active subprocess.""" + global _active_process + if _active_process: + print( + "\n\n⚠ Received interrupt signal. Terminating active subprocess...", + file=sys.stderr, + ) + try: + _active_process.terminate() + _active_process.wait(timeout=5) + except subprocess.TimeoutExpired: + print(" Subprocess didn't terminate, killing it...", file=sys.stderr) + _active_process.kill() + _active_process.wait() + print(" Subprocess terminated.", file=sys.stderr) + sys.exit(1) + + +def get_tenant_index_name(pod_name: str, tenant_id: str) -> str: + """Get the default index name for the given tenant by running script on pod.""" + print(f"Getting default index name for tenant: {tenant_id}") + + # Get the path to the script + script_dir = Path(__file__).parent + index_name_script = script_dir / "on_pod_scripts" / "get_tenant_index_name.py" + + if not index_name_script.exists(): + raise FileNotFoundError( + f"get_tenant_index_name.py not found at {index_name_script}" + ) + + try: + # Copy script to pod + print(" Copying script to pod...") + subprocess.run( + [ + "kubectl", + "cp", + str(index_name_script), + f"{pod_name}:/tmp/get_tenant_index_name.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + print(" Executing script on pod...") + result = subprocess.run( + [ + "kubectl", + "exec", + pod_name, + "--", + "python", + "/tmp/get_tenant_index_name.py", + tenant_id, + ], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(f" {result.stderr}", end="") + + # Parse JSON result from stdout + result_data = json.loads(result.stdout) + status = result_data.get("status") + + if status == "success": + index_name = result_data.get("index_name") + print(f"✓ Found index name: {index_name}") + return index_name + else: + message = result_data.get("message", "Unknown error") + raise RuntimeError(f"Failed to get index name: {message}") + + except subprocess.CalledProcessError as e: + print( + f"✗ Failed to get index name for tenant {tenant_id}: {e}", file=sys.stderr + ) + if e.stderr: + print(f" Error details: {e.stderr}", file=sys.stderr) + raise + except Exception as e: + print( + f"✗ Failed to get index name for tenant {tenant_id}: {e}", file=sys.stderr + ) + raise + + +def get_tenant_users(pod_name: str, tenant_id: str) -> list[str]: + """Get list of user emails from the tenant's data plane schema. + + Args: + pod_name: The Kubernetes pod name to execute on + tenant_id: The tenant ID to query + + Returns: + List of user email addresses, or empty list if query fails + """ + print(f"Fetching user emails for tenant: {tenant_id}") + + # Get the path to the script + script_dir = Path(__file__).parent + get_users_script = script_dir / "on_pod_scripts" / "get_tenant_users.py" + + if not get_users_script.exists(): + raise FileNotFoundError(f"get_tenant_users.py not found at {get_users_script}") + + try: + # Copy script to pod + print(" Copying script to pod...") + subprocess.run( + [ + "kubectl", + "cp", + str(get_users_script), + f"{pod_name}:/tmp/get_tenant_users.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + print(" Executing script on pod...") + result = subprocess.run( + [ + "kubectl", + "exec", + pod_name, + "--", + "python", + "/tmp/get_tenant_users.py", + tenant_id, + ], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(f" {result.stderr}", end="") + + # Parse JSON result from stdout + result_data = json.loads(result.stdout) + status = result_data.get("status") + + if status == "success": + users = result_data.get("users", []) + if users: + print(f"✓ Found {len(users)} user(s):") + for email in users: + print(f" - {email}") + else: + print(" No users found in tenant") + return users + else: + message = result_data.get("message", "Unknown error") + print(f"⚠ Could not fetch users: {message}") + return [] + + except subprocess.CalledProcessError as e: + print(f"⚠ Failed to get users for tenant {tenant_id}: {e}") + if e.stderr: + print(f" Error details: {e.stderr}") + return [] + except Exception as e: + print(f"⚠ Failed to get users for tenant {tenant_id}: {e}") + return [] + + +def check_documents_deleted(pod_name: str, tenant_id: str) -> None: + """Check if all documents and connector credential pairs have been deleted. + + Raises RuntimeError if any ConnectorCredentialPairs or Documents remain. + """ + print(f"Checking for remaining documents in tenant: {tenant_id}") + + # Get the path to the script + script_dir = Path(__file__).parent + check_script = script_dir / "on_pod_scripts" / "check_documents_deleted.py" + + if not check_script.exists(): + raise FileNotFoundError( + f"check_documents_deleted.py not found at {check_script}" + ) + + try: + # Copy script to pod + print(" Copying script to pod...") + subprocess.run( + [ + "kubectl", + "cp", + str(check_script), + f"{pod_name}:/tmp/check_documents_deleted.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + print(" Executing check on pod...") + result = subprocess.run( + [ + "kubectl", + "exec", + pod_name, + "--", + "python", + "/tmp/check_documents_deleted.py", + tenant_id, + ], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(f" {result.stderr}", end="") + + # Parse JSON result from stdout + result_data = json.loads(result.stdout) + status = result_data.get("status") + + if status == "success": + message = result_data.get("message") + print(f"✓ {message}") + elif status == "not_found": + message = result_data.get("message", "Schema not found") + print(f"⚠ {message}") + else: + message = result_data.get("message", "Unknown error") + cc_count = result_data.get("connector_credential_pair_count", 0) + doc_count = result_data.get("document_count", 0) + error_details = f"{message}" + if cc_count > 0 or doc_count > 0: + error_details += f"\n ConnectorCredentialPairs: {cc_count}\n Documents: {doc_count}" + raise RuntimeError(error_details) + + except subprocess.CalledProcessError as e: + print( + f"✗ Failed to check documents for tenant {tenant_id}: {e}", + file=sys.stderr, + ) + if e.stderr: + print(f" Error details: {e.stderr}", file=sys.stderr) + raise + except Exception as e: + print( + f"✗ Failed to check documents for tenant {tenant_id}: {e}", + file=sys.stderr, + ) + raise + + +def drop_data_plane_schema(pod_name: str, tenant_id: str) -> None: + """Drop the PostgreSQL schema for the given tenant by running script on pod.""" + print(f"Dropping data plane schema for tenant: {tenant_id}") + + # Get the path to the cleanup script + script_dir = Path(__file__).parent + schema_cleanup_script = script_dir / "on_pod_scripts" / "cleanup_tenant_schema.py" + + if not schema_cleanup_script.exists(): + raise FileNotFoundError( + f"cleanup_tenant_schema.py not found at {schema_cleanup_script}" + ) + + try: + # Copy script to pod + print(" Copying script to pod...") + subprocess.run( + [ + "kubectl", + "cp", + str(schema_cleanup_script), + f"{pod_name}:/tmp/cleanup_tenant_schema.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + print(" Executing schema cleanup on pod...") + result = subprocess.run( + [ + "kubectl", + "exec", + pod_name, + "--", + "python", + "/tmp/cleanup_tenant_schema.py", + tenant_id, + ], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(f" {result.stderr}", end="") + + # Parse JSON result from stdout + result_data = json.loads(result.stdout) + status = result_data.get("status") + message = result_data.get("message") + + if status == "success": + print(f"✓ {message}") + elif status == "not_found": + print(f"⚠ {message}") + else: + print(f"✗ {message}", file=sys.stderr) + raise RuntimeError(message) + + except subprocess.CalledProcessError as e: + print(f"✗ Failed to drop schema for tenant {tenant_id}: {e}", file=sys.stderr) + if e.stderr: + print(f" Error details: {e.stderr}", file=sys.stderr) + raise + except Exception as e: + print(f"✗ Failed to drop schema for tenant {tenant_id}: {e}", file=sys.stderr) + raise + + +def cleanup_control_plane(tenant_id: str, force: bool = False) -> None: + """ + Clean up control plane data (tenants table, subscription table, etc.) + + Deletes from tables in this order: + 1. tenant_notification (foreign key to tenant) + 2. tenant_config (foreign key to tenant) + 3. subscription (foreign key to tenant) + 4. tenant (primary table) + """ + print(f"Cleaning up control plane data for tenant: {tenant_id}") + + # Delete in order respecting foreign key constraints + delete_queries = [ + ( + "tenant_notification", + "DELETE FROM tenant_notification WHERE tenant_id = '{tenant_id}';", + ), + ("tenant_config", "DELETE FROM tenant_config WHERE tenant_id = '{tenant_id}';"), + ("subscription", "DELETE FROM subscription WHERE tenant_id = '{tenant_id}';"), + ("tenant", "DELETE FROM tenant WHERE tenant_id = '{tenant_id}';"), + ] + + try: + for table_name, query in delete_queries: + formatted_query = query.format(tenant_id=tenant_id) + print(f" Deleting from {table_name}...") + + if not confirm_step(f"Delete from {table_name}?", force): + print(f" Skipping deletion from {table_name}") + continue + + result = execute_control_plane_query(formatted_query) + + if result.stdout: + # Extract row count from output (e.g., "DELETE 5") + print(f" {result.stdout.strip()}") + + print(f"✓ Successfully cleaned up control plane data for tenant: {tenant_id}") + + except subprocess.CalledProcessError as e: + print( + f"✗ Failed to clean up control plane for tenant {tenant_id}: {e}", + file=sys.stderr, + ) + if e.stderr: + print(f" Error details: {e.stderr}", file=sys.stderr) + raise + + +def cleanup_tenant(tenant_id: str, force: bool = False) -> None: + """ + Main cleanup function that orchestrates all cleanup steps. + + Args: + tenant_id: The tenant ID to clean up + force: If True, skip all confirmation prompts + """ + print(f"Starting cleanup for tenant: {tenant_id}") + + # Check tenant status first + print(f"\n{'=' * 80}") + try: + tenant_status = get_tenant_status(tenant_id) + + # If tenant is not GATED_ACCESS, require explicit confirmation even in force mode + if tenant_status and tenant_status != "GATED_ACCESS": + print( + f"\n⚠️ WARNING: Tenant status is '{tenant_status}', not 'GATED_ACCESS'!" + ) + print( + "This tenant may be active and should not be deleted without careful review." + ) + print(f"{'=' * 80}\n") + + # Always ask for confirmation if not gated, even in force mode + response = input( + "Are you ABSOLUTELY SURE you want to proceed? Type 'yes' to confirm: " + ) + if response.lower() != "yes": + print("Cleanup aborted - tenant is not GATED_ACCESS") + return + elif tenant_status == "GATED_ACCESS": + print("✓ Tenant status is GATED_ACCESS - safe to proceed with cleanup") + elif tenant_status is None: + print("⚠️ WARNING: Could not determine tenant status!") + response = input("Continue anyway? Type 'yes' to confirm: ") + if response.lower() != "yes": + print("Cleanup aborted - could not verify tenant status") + return + except Exception as e: + print(f"⚠️ WARNING: Failed to check tenant status: {e}") + response = input("Continue anyway? Type 'yes' to confirm: ") + if response.lower() != "yes": + print("Cleanup aborted - could not verify tenant status") + return + print(f"{'=' * 80}\n") + + # Find heavy worker pod for Vespa and schema operations + try: + pod_name = find_worker_pod() + except Exception as e: + print(f"✗ Failed to find heavy worker pod: {e}", file=sys.stderr) + print("Cannot proceed with Vespa and schema cleanup") + return + + # Fetch tenant users for informational purposes (non-blocking) + print(f"\n{'=' * 80}") + try: + get_tenant_users(pod_name, tenant_id) + except Exception as e: + print(f"⚠ Could not fetch tenant users: {e}") + print(f"{'=' * 80}\n") + + # Step 1: Make sure all documents are deleted + print(f"\n{'=' * 80}") + print("Step 1/3: Checking for remaining ConnectorCredentialPairs and Documents") + print(f"{'=' * 80}") + try: + check_documents_deleted(pod_name, tenant_id) + except Exception as e: + print(f"✗ Document check failed: {e}", file=sys.stderr) + print( + "\nPlease ensure all ConnectorCredentialPairs and Documents are deleted before running cleanup." + ) + print( + "You may need to mark connectors for deletion and wait for cleanup to complete." + ) + return + print(f"{'=' * 80}\n") + + # Step 2: Drop data plane schema + if confirm_step( + f"Step 2/3: Drop data plane schema '{tenant_id}' (CASCADE - will delete all tables, functions, etc.)", + force, + ): + try: + drop_data_plane_schema(pod_name, tenant_id) + except Exception as e: + print(f"✗ Failed at schema cleanup step: {e}", file=sys.stderr) + if not force: + response = input("Continue with control plane cleanup? (y/n): ") + if response.lower() != "y": + print("Cleanup aborted by user") + return + else: + print("[FORCE MODE] Continuing despite schema cleanup failure") + else: + print("Step 2 skipped by user") + + # Step 3: Clean up control plane + if confirm_step( + "Step 3/3: Delete control plane records (tenant_notification, tenant_config, subscription, tenant)", + force, + ): + try: + cleanup_control_plane(tenant_id, force) + except Exception as e: + print(f"✗ Failed at control plane cleanup step: {e}", file=sys.stderr) + if not force: + print("Control plane cleanup failed") + else: + print("[FORCE MODE] Control plane cleanup failed but continuing") + else: + print("Step 3 skipped by user") + return + + print(f"\n{'=' * 80}") + print(f"✓ Cleanup completed for tenant: {tenant_id}") + print(f"{'=' * 80}") + + +def main() -> None: + # Register signal handlers for graceful shutdown + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + if len(sys.argv) < 2: + print("Usage: python backend/scripts/cleanup_tenant.py [--force]") + print( + " python backend/scripts/cleanup_tenant.py --csv [--force]" + ) + print("\nArguments:") + print( + " tenant_id The tenant ID to clean up (required if not using --csv)" + ) + print(" --csv PATH Path to CSV file containing tenant IDs to clean up") + print(" --force Skip all confirmation prompts (optional)") + print("\nExamples:") + print(" python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789") + print( + " python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789 --force" + ) + print( + " python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv" + ) + print( + " python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv --force" + ) + sys.exit(1) + + # Parse arguments + force = "--force" in sys.argv + tenant_ids = [] + + # Check for CSV mode + if "--csv" in sys.argv: + try: + csv_index = sys.argv.index("--csv") + if csv_index + 1 >= len(sys.argv): + print("Error: --csv flag requires a file path", file=sys.stderr) + sys.exit(1) + + csv_path = sys.argv[csv_index + 1] + tenant_ids = read_tenant_ids_from_csv(csv_path) + + if not tenant_ids: + print("Error: No tenant IDs found in CSV file", file=sys.stderr) + sys.exit(1) + + print(f"Found {len(tenant_ids)} tenant(s) in CSV file: {csv_path}") + + except Exception as e: + print(f"Error reading CSV file: {e}", file=sys.stderr) + sys.exit(1) + else: + # Single tenant mode + tenant_ids = [sys.argv[1]] + + # Initial confirmation (unless --force is used) + if not force: + print(f"\n{'=' * 80}") + print("TENANT CLEANUP - CONFIRMATION REQUIRED") + print(f"{'=' * 80}") + if len(tenant_ids) == 1: + print(f"Tenant ID: {tenant_ids[0]}") + else: + print(f"Number of tenants: {len(tenant_ids)}") + print(f"Tenant IDs: {', '.join(tenant_ids[:5])}") + if len(tenant_ids) > 5: + print(f" ... and {len(tenant_ids) - 5} more") + + print("Index Name: Will be fetched automatically when deleting Vespa documents") + print( + f"Mode: {'FORCE (no confirmations)' if force else 'Interactive (will ask for confirmation at each step)'}" + ) + print("\nThis will:") + print(" 1. Delete ALL Vespa documents for this tenant") + print(" 2. Drop the data plane PostgreSQL schema (CASCADE)") + print(" 3. Clean up control plane data:") + print(" - Delete from tenant_notification table") + print(" - Delete from tenant_config table") + print(" - Delete from subscription table") + print(" - Delete from tenant table") + print(f"\n{'=' * 80}") + print("WARNING: This operation is IRREVERSIBLE!") + print(f"{'=' * 80}\n") + + response = input("Are you sure you want to proceed? Type 'yes' to confirm: ") + + if response.lower() != "yes": + print("Cleanup aborted by user") + sys.exit(0) + else: + if len(tenant_ids) == 1: + print( + f"⚠ FORCE MODE: Running cleanup for {tenant_ids[0]} without confirmations" + ) + else: + print( + f"⚠ FORCE MODE: Running cleanup for {len(tenant_ids)} tenants without confirmations" + ) + + # Run cleanup for each tenant + failed_tenants = [] + successful_tenants = [] + + for idx, tenant_id in enumerate(tenant_ids, 1): + if len(tenant_ids) > 1: + print(f"\n{'=' * 80}") + print(f"Processing tenant {idx}/{len(tenant_ids)}: {tenant_id}") + print(f"{'=' * 80}") + + try: + cleanup_tenant(tenant_id, force) + successful_tenants.append(tenant_id) + except Exception as e: + print(f"✗ Cleanup failed for tenant {tenant_id}: {e}", file=sys.stderr) + failed_tenants.append((tenant_id, str(e))) + + # If not in force mode and there are more tenants, ask if we should continue + if not force and idx < len(tenant_ids): + response = input( + f"\nContinue with remaining {len(tenant_ids) - idx} tenant(s)? (y/n): " + ) + if response.lower() != "y": + print("Cleanup aborted by user") + break + + # Print summary if multiple tenants + if len(tenant_ids) > 1: + print(f"\n{'=' * 80}") + print("CLEANUP SUMMARY") + print(f"{'=' * 80}") + print(f"Total tenants: {len(tenant_ids)}") + print(f"Successful: {len(successful_tenants)}") + print(f"Failed: {len(failed_tenants)}") + + if failed_tenants: + print("\nFailed tenants:") + for tenant_id, error in failed_tenants: + print(f" - {tenant_id}: {error}") + + print(f"{'=' * 80}") + + if failed_tenants: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/cleanup_utils.py b/backend/scripts/tenant_cleanup/cleanup_utils.py new file mode 100644 index 0000000000..91cdc69442 --- /dev/null +++ b/backend/scripts/tenant_cleanup/cleanup_utils.py @@ -0,0 +1,194 @@ +import csv +import os +import random +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + + +@dataclass +class ControlPlaneConfig: + """Configuration for connecting to the control plane database.""" + + db_url: str + bastion_host: str + pem_file_location: str + + +def find_worker_pod() -> str: + """Find a user file processing worker pod using kubectl.""" + print("Finding user file processing worker pod...") + + result = subprocess.run( + ["kubectl", "get", "po"], capture_output=True, text=True, check=True + ) + + # Parse output and find user file processing worker pod + lines = result.stdout.strip().split("\n") + lines = lines[1:] # Skip header + random.shuffle(lines) + for line in lines: + if "celery-worker-user-file-processing" in line and "Running" in line: + pod_name = line.split()[0] + print(f"Found pod: {pod_name}") + return pod_name + + raise RuntimeError("No running user file processing worker pod found") + + +def confirm_step(message: str, force: bool = False) -> bool: + """Ask for confirmation before executing a step. + + Args: + message: The confirmation message to display + force: If True, skip confirmation and return True + + Returns: + True if user confirms or force is True, False otherwise + """ + if force: + print(f"[FORCE MODE] Skipping confirmation: {message}") + return True + + print(f"\n{message}") + response = input("Proceed? (y/n): ") + return response.lower() == "y" + + +def get_control_plane_config() -> ControlPlaneConfig: + """Get control plane database configuration from environment variables. + + Returns: + ControlPlaneConfig with db_url, bastion_host, and pem_file_location + + Raises: + ValueError: If any required environment variable is not set + """ + rds_host = os.environ.get("CONTROL_PLANE_RDS_HOST") + if not rds_host: + raise ValueError("CONTROL_PLANE_RDS_HOST is not set") + + rds_password = os.environ.get("CONTROL_PLANE_RDS_PASSWORD") + if not rds_password: + raise ValueError("CONTROL_PLANE_RDS_PASSWORD is not set") + + bastion_host = os.environ.get("BASTION_HOST") + if not bastion_host: + raise ValueError("BASTION_HOST is not set") + + pem_file_location = os.environ.get("PEM_FILE_LOCATION") + if not pem_file_location: + raise ValueError("PEM_FILE_LOCATION is not set") + + db_url = f"postgresql://postgres:{rds_password}@{rds_host}:5432/control" + + return ControlPlaneConfig( + db_url=db_url, + bastion_host=bastion_host, + pem_file_location=pem_file_location, + ) + + +def execute_control_plane_query( + query: str, tuple_only: bool = False +) -> subprocess.CompletedProcess: + """Execute a SQL query against the control plane database via SSH. + + Args: + query: The SQL query to execute + tuple_only: If True, use psql's tuple-only mode (-t flag) for cleaner output + + Returns: + subprocess.CompletedProcess with the result + + Raises: + subprocess.CalledProcessError: If the command fails + """ + config = get_control_plane_config() + db_url = config.db_url + bastion_host = config.bastion_host + pem_file_location = config.pem_file_location + + # Build psql flags + psql_flags = "-t" if tuple_only else "" + + # Build the SSH command with proper escaping + full_cmd = ( + f"ssh -i {pem_file_location} ec2-user@{bastion_host} " + f'"psql {db_url} {psql_flags} -c \\"{query}\\""' + ) + + result = subprocess.run( + full_cmd, + shell=True, + check=True, + capture_output=True, + text=True, + ) + + return result + + +def get_tenant_status(tenant_id: str) -> str | None: + """ + Get tenant status from control plane database. + + Returns: + Tenant status string (e.g., 'GATED_ACCESS', 'ACTIVE') or None if not found + """ + print(f"Fetching tenant status for tenant: {tenant_id}") + + query = f"SELECT application_status FROM tenant WHERE tenant_id = '{tenant_id}';" + + try: + result = execute_control_plane_query(query, tuple_only=True) + + # Parse the output - psql returns the value with whitespace + status = result.stdout.strip() + + if status: + print(f"✓ Tenant status: {status}") + return status + else: + print("⚠ Tenant not found in control plane") + return None + + except subprocess.CalledProcessError as e: + print( + f"✗ Failed to get tenant status for {tenant_id}: {e}", + file=sys.stderr, + ) + if e.stderr: + print(f" Error details: {e.stderr}", file=sys.stderr) + return None + + +def read_tenant_ids_from_csv(csv_path: str) -> list[str]: + """Read tenant IDs from CSV file. + + Args: + csv_path: Path to CSV file + + Returns: + List of tenant IDs + """ + if not Path(csv_path).exists(): + raise FileNotFoundError(f"CSV file not found: {csv_path}") + + tenant_ids = [] + with open(csv_path, "r", newline="", encoding="utf-8") as csvfile: + reader = csv.DictReader(csvfile) + + # Check if tenant_id column exists + if not reader.fieldnames or "tenant_id" not in reader.fieldnames: + raise ValueError( + f"CSV file must have a 'tenant_id' column. Found columns: {reader.fieldnames}" + ) + + for row in reader: + tenant_id = row.get("tenant_id", "").strip() + if tenant_id: + tenant_ids.append(tenant_id) + + return tenant_ids diff --git a/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py b/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py new file mode 100644 index 0000000000..1d43fdb003 --- /dev/null +++ b/backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py @@ -0,0 +1,465 @@ +#!/usr/bin/env python3 +""" +Mark connectors for deletion script that: +1. Finds all connectors for the specified tenant(s) +2. Cancels any scheduled indexing attempts +3. Marks each connector credential pair as DELETING +4. Triggers the cleanup task + +Usage: + python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py [--force] + python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv [--force] + +Arguments: + tenant_id The tenant ID to process (required if not using --csv) + --csv PATH Path to CSV file containing tenant IDs to process + --force Skip all confirmation prompts (optional) + +Examples: + python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789 + python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789 --force + python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv + python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv --force +""" + +import json +import subprocess +import sys +from pathlib import Path + +from scripts.tenant_cleanup.cleanup_utils import confirm_step +from scripts.tenant_cleanup.cleanup_utils import find_worker_pod +from scripts.tenant_cleanup.cleanup_utils import get_tenant_status +from scripts.tenant_cleanup.cleanup_utils import read_tenant_ids_from_csv + + +def get_tenant_connectors(pod_name: str, tenant_id: str) -> list[dict]: + """Get list of connector credential pairs for the tenant. + + Args: + pod_name: The Kubernetes pod name to execute on + tenant_id: The tenant ID to query + + Returns: + List of connector credential pair dicts with id, connector_id, credential_id, name, status + """ + print(f"Fetching connector credential pairs for tenant: {tenant_id}") + + # Get the path to the script + script_dir = Path(__file__).parent + get_connectors_script = script_dir / "on_pod_scripts" / "get_tenant_connectors.py" + + if not get_connectors_script.exists(): + raise FileNotFoundError( + f"get_tenant_connectors.py not found at {get_connectors_script}" + ) + + try: + # Copy script to pod + print(" Copying script to pod...") + subprocess.run( + [ + "kubectl", + "cp", + str(get_connectors_script), + f"{pod_name}:/tmp/get_tenant_connectors.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + print(" Executing script on pod...") + result = subprocess.run( + [ + "kubectl", + "exec", + pod_name, + "--", + "python", + "/tmp/get_tenant_connectors.py", + tenant_id, + ], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(f" {result.stderr}", end="") + + # Parse JSON result from stdout + result_data = json.loads(result.stdout) + status = result_data.get("status") + + if status == "success": + connectors = result_data.get("connectors", []) + if connectors: + print(f"✓ Found {len(connectors)} connector credential pair(s):") + for cc in connectors: + print( + f" - CC Pair ID: {cc['id']}, Name: {cc['name']}, Status: {cc['status']}" + ) + else: + print(" No connector credential pairs found for tenant") + return connectors + else: + message = result_data.get("message", "Unknown error") + print(f"⚠ Could not fetch connectors: {message}") + return [] + + except subprocess.CalledProcessError as e: + print(f"⚠ Failed to get connectors for tenant {tenant_id}: {e}") + if e.stderr: + print(f" Error details: {e.stderr}") + return [] + except Exception as e: + print(f"⚠ Failed to get connectors for tenant {tenant_id}: {e}") + return [] + + +def mark_connector_for_deletion(pod_name: str, tenant_id: str, cc_pair_id: int) -> None: + """Mark a connector credential pair for deletion. + + Args: + pod_name: The Kubernetes pod name to execute on + tenant_id: The tenant ID + cc_pair_id: The connector credential pair ID to mark for deletion + """ + print(f" Marking CC pair {cc_pair_id} for deletion...") + + # Get the path to the script + script_dir = Path(__file__).parent + mark_deletion_script = ( + script_dir / "on_pod_scripts" / "mark_connector_for_deletion.py" + ) + + if not mark_deletion_script.exists(): + raise FileNotFoundError( + f"mark_connector_for_deletion.py not found at {mark_deletion_script}" + ) + + try: + # Copy script to pod + subprocess.run( + [ + "kubectl", + "cp", + str(mark_deletion_script), + f"{pod_name}:/tmp/mark_connector_for_deletion.py", + ], + check=True, + capture_output=True, + ) + + # Execute script on pod + result = subprocess.run( + [ + "kubectl", + "exec", + pod_name, + "--", + "python", + "/tmp/mark_connector_for_deletion.py", + tenant_id, + str(cc_pair_id), + ], + capture_output=True, + text=True, + check=True, + ) + + # Show progress messages from stderr + if result.stderr: + print(f" {result.stderr}", end="") + + # Parse JSON result from stdout + result_data = json.loads(result.stdout) + status = result_data.get("status") + message = result_data.get("message") + + if status == "success": + print(f" ✓ {message}") + else: + print(f" ✗ {message}", file=sys.stderr) + raise RuntimeError(message) + + except subprocess.CalledProcessError as e: + print( + f" ✗ Failed to mark CC pair {cc_pair_id} for deletion: {e}", + file=sys.stderr, + ) + if e.stderr: + print(f" Error details: {e.stderr}", file=sys.stderr) + raise + except Exception as e: + print( + f" ✗ Failed to mark CC pair {cc_pair_id} for deletion: {e}", + file=sys.stderr, + ) + raise + + +def mark_tenant_connectors_for_deletion(tenant_id: str, force: bool = False) -> None: + """ + Main function to mark all connectors for a tenant for deletion. + + Args: + tenant_id: The tenant ID to process + force: If True, skip all confirmation prompts + """ + print(f"Processing connectors for tenant: {tenant_id}") + + # Check tenant status first + print(f"\n{'=' * 80}") + try: + tenant_status = get_tenant_status(tenant_id) + + # If tenant is not GATED_ACCESS, require explicit confirmation even in force mode + if tenant_status and tenant_status != "GATED_ACCESS": + print( + f"\n⚠️ WARNING: Tenant status is '{tenant_status}', not 'GATED_ACCESS'!" + ) + print( + "This tenant may be active and should not have connectors deleted without careful review." + ) + print(f"{'=' * 80}\n") + + # Always ask for confirmation if not gated, even in force mode + response = input( + "Are you ABSOLUTELY SURE you want to proceed? Type 'yes' to confirm: " + ) + if response.lower() != "yes": + print("Operation aborted - tenant is not GATED_ACCESS") + return + elif tenant_status == "GATED_ACCESS": + print("✓ Tenant status is GATED_ACCESS - safe to proceed") + elif tenant_status is None: + print("⚠️ WARNING: Could not determine tenant status!") + response = input("Continue anyway? Type 'yes' to confirm: ") + if response.lower() != "yes": + print("Operation aborted - could not verify tenant status") + return + except Exception as e: + print(f"⚠️ WARNING: Failed to check tenant status: {e}") + response = input("Continue anyway? Type 'yes' to confirm: ") + if response.lower() != "yes": + print("Operation aborted - could not verify tenant status") + return + print(f"{'=' * 80}\n") + + # Find heavy worker pod for operations + try: + pod_name = find_worker_pod() + except Exception as e: + print(f"✗ Failed to find heavy worker pod: {e}", file=sys.stderr) + print("Cannot proceed with marking connectors for deletion") + return + + # Fetch connectors + print(f"\n{'=' * 80}") + try: + connectors = get_tenant_connectors(pod_name, tenant_id) + except Exception as e: + print(f"✗ Failed to fetch connectors: {e}", file=sys.stderr) + return + print(f"{'=' * 80}\n") + + if not connectors: + print(f"No connectors found for tenant {tenant_id}, nothing to do.") + return + + # Confirm before proceeding + if not confirm_step( + f"Mark {len(connectors)} connector credential pair(s) for deletion?", + force, + ): + print("Operation cancelled by user") + return + + # Mark each connector for deletion + failed_connectors = [] + successful_connectors = [] + + for cc in connectors: + cc_pair_id = cc["id"] + cc_name = cc["name"] + cc_status = cc["status"] + + # Skip if already marked for deletion + if cc_status == "DELETING": + print( + f" Skipping CC pair {cc_pair_id} ({cc_name}) - already marked for deletion" + ) + continue + + try: + mark_connector_for_deletion(pod_name, tenant_id, cc_pair_id) + successful_connectors.append(cc_pair_id) + except Exception as e: + print( + f" ✗ Failed to mark CC pair {cc_pair_id} ({cc_name}) for deletion: {e}", + file=sys.stderr, + ) + failed_connectors.append((cc_pair_id, cc_name, str(e))) + + # Print summary + print(f"\n{'=' * 80}") + print(f"✓ Marked {len(successful_connectors)} connector(s) for deletion") + if failed_connectors: + print(f"✗ Failed to mark {len(failed_connectors)} connector(s):") + for cc_id, cc_name, error in failed_connectors: + print(f" - CC Pair {cc_id} ({cc_name}): {error}") + print(f"{'=' * 80}") + + +def main() -> None: + if len(sys.argv) < 2: + print( + "Usage: python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py [--force]" + ) + print( + " python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv [--force]" + ) + print("\nArguments:") + print( + " tenant_id The tenant ID to process (required if not using --csv)" + ) + print(" --csv PATH Path to CSV file containing tenant IDs to process") + print(" --force Skip all confirmation prompts (optional)") + print("\nExamples:") + print( + " python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789" + ) + print( + " python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789 --force" + ) + print( + " python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv" + ) + print( + " python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv --force" + ) + sys.exit(1) + + # Parse arguments + force = "--force" in sys.argv + tenant_ids = [] + + # Check for CSV mode + if "--csv" in sys.argv: + try: + csv_index = sys.argv.index("--csv") + if csv_index + 1 >= len(sys.argv): + print("Error: --csv flag requires a file path", file=sys.stderr) + sys.exit(1) + + csv_path = sys.argv[csv_index + 1] + tenant_ids = read_tenant_ids_from_csv(csv_path) + + if not tenant_ids: + print("Error: No tenant IDs found in CSV file", file=sys.stderr) + sys.exit(1) + + print(f"Found {len(tenant_ids)} tenant(s) in CSV file: {csv_path}") + + except Exception as e: + print(f"Error reading CSV file: {e}", file=sys.stderr) + sys.exit(1) + else: + # Single tenant mode + tenant_ids = [sys.argv[1]] + + # Initial confirmation (unless --force is used) + if not force: + print(f"\n{'=' * 80}") + print("MARK CONNECTORS FOR DELETION - CONFIRMATION REQUIRED") + print(f"{'=' * 80}") + if len(tenant_ids) == 1: + print(f"Tenant ID: {tenant_ids[0]}") + else: + print(f"Number of tenants: {len(tenant_ids)}") + print(f"Tenant IDs: {', '.join(tenant_ids[:5])}") + if len(tenant_ids) > 5: + print(f" ... and {len(tenant_ids) - 5} more") + + print( + f"Mode: {'FORCE (no confirmations)' if force else 'Interactive (will ask for confirmation at each step)'}" + ) + print("\nThis will:") + print(" 1. Fetch all connector credential pairs for each tenant") + print(" 2. Cancel any scheduled indexing attempts for each connector") + print(" 3. Mark each connector credential pair status as DELETING") + print(" 4. Trigger the connector deletion task") + print(f"\n{'=' * 80}") + print("WARNING: This will mark connectors for deletion!") + print("The actual deletion will be performed by the background celery worker.") + print(f"{'=' * 80}\n") + + response = input("Are you sure you want to proceed? Type 'yes' to confirm: ") + + if response.lower() != "yes": + print("Operation aborted by user") + sys.exit(0) + else: + if len(tenant_ids) == 1: + print( + f"⚠ FORCE MODE: Marking connectors for deletion for {tenant_ids[0]} without confirmations" + ) + else: + print( + f"⚠ FORCE MODE: Marking connectors for deletion for {len(tenant_ids)} tenants without confirmations" + ) + + # Process each tenant + failed_tenants = [] + successful_tenants = [] + + for idx, tenant_id in enumerate(tenant_ids, 1): + if len(tenant_ids) > 1: + print(f"\n{'=' * 80}") + print(f"Processing tenant {idx}/{len(tenant_ids)}: {tenant_id}") + print(f"{'=' * 80}") + + try: + mark_tenant_connectors_for_deletion(tenant_id, force) + successful_tenants.append(tenant_id) + except Exception as e: + print( + f"✗ Failed to process tenant {tenant_id}: {e}", + file=sys.stderr, + ) + failed_tenants.append((tenant_id, str(e))) + + # If not in force mode and there are more tenants, ask if we should continue + if not force and idx < len(tenant_ids): + response = input( + f"\nContinue with remaining {len(tenant_ids) - idx} tenant(s)? (y/n): " + ) + if response.lower() != "y": + print("Operation aborted by user") + break + + # Print summary if multiple tenants + if len(tenant_ids) > 1: + print(f"\n{'=' * 80}") + print("OPERATION SUMMARY") + print(f"{'=' * 80}") + print(f"Total tenants: {len(tenant_ids)}") + print(f"Successful: {len(successful_tenants)}") + print(f"Failed: {len(failed_tenants)}") + + if failed_tenants: + print("\nFailed tenants:") + for tenant_id, error in failed_tenants: + print(f" - {tenant_id}: {error}") + + print(f"{'=' * 80}") + + if failed_tenants: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py b/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py new file mode 100644 index 0000000000..5758f8f592 --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +""" +Script to check for remaining ConnectorCredentialPairs and Documents in a tenant's schema. +Must be run on a pod with access to the data plane PostgreSQL database. + +Usage: + python check_documents_deleted.py + +Output: + JSON object with status, message, and counts of remaining records +""" + +import json +import sys + +from sqlalchemy import func +from sqlalchemy import select + +from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.models import ConnectorCredentialPair +from onyx.db.models import Document + + +def check_documents_deleted(tenant_id: str) -> dict: + """ + Check for remaining ConnectorCredentialPairs and Documents in tenant schema. + + Args: + tenant_id: The tenant ID to query + + Returns: + Dictionary with status and counts of remaining records + """ + try: + print( + f"Checking for remaining documents in tenant: {tenant_id}", + file=sys.stderr, + ) + + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + # Count ConnectorCredentialPairs + cc_count = db_session.scalar( + select(func.count()).select_from(ConnectorCredentialPair) + ) + + # Count Documents + doc_count = db_session.scalar(select(func.count()).select_from(Document)) + + # Handle None values from scalar (should not happen but mypy needs it) + cc_count = cc_count or 0 + doc_count = doc_count or 0 + + # If any records remain, return error status + if cc_count > 0 or doc_count > 0: + return { + "status": "error", + "message": ( + f"Found {cc_count} ConnectorCredentialPair(s) and {doc_count} Document(s) " + "still remaining. All documents must be deleted before cleanup." + ), + "connector_credential_pair_count": cc_count, + "document_count": doc_count, + } + + # All clear + return { + "status": "success", + "message": "No ConnectorCredentialPairs or Documents found - safe to proceed", + "connector_credential_pair_count": 0, + "document_count": 0, + } + + except Exception as e: + error_msg = str(e) + print(f"Error checking documents: {error_msg}", file=sys.stderr) + # Check if it's a schema not found error + if "does not exist" in error_msg: + return { + "status": "not_found", + "message": f"Schema '{tenant_id}' does not exist", + } + return {"status": "error", "message": f"Error checking documents: {error_msg}"} + + +def main() -> None: + if len(sys.argv) != 2: + print( + json.dumps( + { + "status": "error", + "message": "Usage: python check_documents_deleted.py ", + } + ) + ) + sys.exit(1) + + tenant_id = sys.argv[1] + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + result = check_documents_deleted(tenant_id) + print(json.dumps(result)) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py b/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py new file mode 100644 index 0000000000..a9a541282f --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +""" +Script to drop a tenant's PostgreSQL schema. +Designed to be run on a heavy worker pod. + +Usage: + python cleanup_tenant_schema.py +""" + +import json +import sys + +from sqlalchemy import text + +from onyx.db.engine.sql_engine import get_session_with_shared_schema +from onyx.db.engine.sql_engine import SqlEngine + + +def drop_data_plane_schema(tenant_id: str) -> dict[str, str]: + """Drop the PostgreSQL schema for the given tenant.""" + print(f"Dropping data plane schema for tenant: {tenant_id}", file=sys.stderr) + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + try: + with get_session_with_shared_schema() as session: + # First, verify the schema exists + check_schema_query = text( + """ + SELECT nspname + FROM pg_namespace + WHERE nspname = :schema_name + """ + ) + + result = session.execute( + check_schema_query, {"schema_name": tenant_id} + ).fetchone() + + if not result: + print(f"Schema {tenant_id} does not exist", file=sys.stderr) + return { + "status": "not_found", + "message": f"Schema {tenant_id} does not exist", + } + + # Drop the schema with CASCADE to remove all objects within it + drop_schema_query = text(f'DROP SCHEMA IF EXISTS "{tenant_id}" CASCADE') + session.execute(drop_schema_query) + session.commit() + + print(f"Successfully dropped schema: {tenant_id}", file=sys.stderr) + + # Delete the tenant mapping from user_tenant_mapping table + delete_mapping_query = text( + """ + DELETE FROM user_tenant_mapping + WHERE tenant_id = :tenant_id + """ + ) + session.execute(delete_mapping_query, {"tenant_id": tenant_id}) + session.commit() + + print( + f"Successfully deleted tenant mapping for: {tenant_id}", file=sys.stderr + ) + return { + "status": "success", + "message": f"Successfully dropped schema: {tenant_id}", + } + + except Exception as e: + print(f"Failed to drop schema for tenant {tenant_id}: {e}", file=sys.stderr) + return {"status": "error", "message": str(e)} + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python cleanup_tenant_schema.py ", file=sys.stderr) + sys.exit(1) + + tenant_id = sys.argv[1] + + result = drop_data_plane_schema(tenant_id) + + # Output result as JSON to stdout for easy parsing + print(json.dumps(result)) + + # Exit with error code if failed + if result["status"] == "error": + sys.exit(1) diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py new file mode 100644 index 0000000000..34c9a85895 --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +""" +Script to fetch connector credential pairs for a tenant. +Runs on a Kubernetes pod with access to the data plane database. + +Usage: + python get_tenant_connectors.py + +Output: + JSON to stdout with structure: + { + "status": "success" | "error", + "connectors": [ + { + "id": int, + "connector_id": int, + "credential_id": int, + "name": str, + "status": str + }, + ... + ] (if success), + "message": str (if error) + } +""" + +import json +import sys + +from sqlalchemy import select + +from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.models import ConnectorCredentialPair + + +def get_tenant_connectors(tenant_id: str) -> dict: + """Get all connector credential pairs for a tenant. + + Args: + tenant_id: The tenant ID to query + + Returns: + Dict with status and list of connectors or error message + """ + try: + print( + f"Fetching connector credential pairs for tenant: {tenant_id}", + file=sys.stderr, + ) + + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + # Get all connector credential pairs, excluding user files + stmt = select(ConnectorCredentialPair).where( + ConnectorCredentialPair.is_user_file.is_(False) + ) + cc_pairs = db_session.execute(stmt).scalars().all() + + connectors = [ + { + "id": cc.id, + "connector_id": cc.connector_id, + "credential_id": cc.credential_id, + "name": cc.name, + "status": cc.status.value, + } + for cc in cc_pairs + ] + + print( + f"Found {len(connectors)} connector credential pair(s)", + file=sys.stderr, + ) + + return { + "status": "success", + "connectors": connectors, + } + + except Exception as e: + print(f"Error fetching connectors: {e}", file=sys.stderr) + return { + "status": "error", + "message": str(e), + } + + +def main() -> None: + if len(sys.argv) != 2: + print( + json.dumps( + { + "status": "error", + "message": "Usage: python get_tenant_connectors.py ", + } + ) + ) + sys.exit(1) + + tenant_id = sys.argv[1] + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + result = get_tenant_connectors(tenant_id) + print(json.dumps(result)) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py new file mode 100644 index 0000000000..0f9b9d5ace --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +""" +Script to get the default index name for a tenant. +Designed to be run on a heavy worker pod. + +Usage: + python get_tenant_index_name.py +""" + +import json +import sys + +from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.search_settings import get_current_search_settings +from onyx.db.search_settings import get_secondary_search_settings +from onyx.document_index.factory import get_default_document_index + + +def get_tenant_index_name(tenant_id: str) -> dict[str, str]: + """Get the default index name for the given tenant.""" + print(f"Getting default index name for tenant: {tenant_id}", file=sys.stderr) + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + try: + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + search_settings = get_current_search_settings(db_session) + secondary_search_settings = get_secondary_search_settings(db_session) + + document_index = get_default_document_index( + search_settings=search_settings, + secondary_search_settings=secondary_search_settings, + ) + index_name = document_index.index_name + + print(f"Found index name: {index_name}", file=sys.stderr) + return {"status": "success", "index_name": index_name} + + except Exception as e: + print(f"Failed to get index name for tenant {tenant_id}: {e}", file=sys.stderr) + return {"status": "error", "message": str(e)} + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python get_tenant_index_name.py ", file=sys.stderr) + sys.exit(1) + + tenant_id = sys.argv[1] + + result = get_tenant_index_name(tenant_id) + + # Output result as JSON to stdout for easy parsing + print(json.dumps(result)) + + # Exit with error code if failed + if result["status"] == "error": + sys.exit(1) diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py new file mode 100644 index 0000000000..b7a1857bbc --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +""" +Script to fetch user emails from a tenant's data plane schema. +Must be run on a pod with access to the data plane PostgreSQL database. + +Usage: + python get_tenant_users.py + +Output: + JSON object with status and users list +""" + +import json +import sys + +from sqlalchemy import select + +from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.models import User + + +def get_tenant_users(tenant_id: str) -> dict: + """ + Fetch user emails from the tenant's data plane schema. + + Args: + tenant_id: The tenant ID to query + + Returns: + Dictionary with status and users list + """ + try: + print(f"Querying users for tenant: {tenant_id}", file=sys.stderr) + + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + # Query users from the tenant schema + stmt = select(User).order_by(User.email) + result = db_session.execute(stmt) + users = [user.email for user in result.scalars()] + + return {"status": "success", "users": users} + + except Exception as e: + error_msg = str(e) + print(f"Error fetching users: {error_msg}", file=sys.stderr) + # Check if it's a schema not found error + if "does not exist" in error_msg: + return { + "status": "not_found", + "message": f"Schema '{tenant_id}' does not exist", + "users": [], + } + return {"status": "error", "message": error_msg, "users": []} + + +def main() -> None: + if len(sys.argv) != 2: + print( + json.dumps( + { + "status": "error", + "message": "Usage: python get_tenant_users.py ", + } + ) + ) + sys.exit(1) + + tenant_id = sys.argv[1] + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + result = get_tenant_users(tenant_id) + print(json.dumps(result)) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/mark_connector_for_deletion.py b/backend/scripts/tenant_cleanup/on_pod_scripts/mark_connector_for_deletion.py new file mode 100644 index 0000000000..b0784d1460 --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/mark_connector_for_deletion.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +""" +Script to mark a connector credential pair for deletion. +Runs on a Kubernetes pod with access to the data plane database. + +Usage: + python mark_connector_for_deletion.py + +Output: + JSON to stdout with structure: + { + "status": "success" | "error", + "message": str + } +""" + +import json +import sys + +from onyx.background.celery.versioned_apps.client import app as client_app +from onyx.configs.constants import OnyxCeleryPriority +from onyx.configs.constants import OnyxCeleryTask +from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id +from onyx.db.connector_credential_pair import update_connector_credential_pair_from_id +from onyx.db.engine.sql_engine import get_session_with_tenant +from onyx.db.engine.sql_engine import SqlEngine +from onyx.db.enums import ConnectorCredentialPairStatus +from onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair + + +def mark_connector_for_deletion(tenant_id: str, cc_pair_id: int) -> dict: + """Mark a connector credential pair for deletion. + + Args: + tenant_id: The tenant ID + cc_pair_id: The connector credential pair ID + + Returns: + Dict with status and message + """ + try: + print( + f"Marking connector credential pair {cc_pair_id} for deletion", + file=sys.stderr, + ) + + with get_session_with_tenant(tenant_id=tenant_id) as db_session: + # Get the connector credential pair + cc_pair = get_connector_credential_pair_from_id( + db_session=db_session, + cc_pair_id=cc_pair_id, + ) + + if not cc_pair: + return { + "status": "error", + "message": f"Connector credential pair {cc_pair_id} not found", + } + + # Cancel any scheduled indexing attempts + print( + f"Canceling indexing attempts for CC pair {cc_pair_id}", + file=sys.stderr, + ) + cancel_indexing_attempts_for_ccpair( + cc_pair_id=cc_pair.id, + db_session=db_session, + include_secondary_index=True, + ) + + # Mark as deleting + print( + f"Updating CC pair {cc_pair_id} status to DELETING", + file=sys.stderr, + ) + update_connector_credential_pair_from_id( + db_session=db_session, + cc_pair_id=cc_pair.id, + status=ConnectorCredentialPairStatus.DELETING, + ) + + db_session.commit() + + # Trigger the deletion check task + print( + "Triggering connector deletion check task", + file=sys.stderr, + ) + client_app.send_task( + OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION, + priority=OnyxCeleryPriority.HIGH, + kwargs={"tenant_id": tenant_id}, + ) + + return { + "status": "success", + "message": f"Marked connector credential pair {cc_pair_id} for deletion", + } + + except Exception as e: + print( + f"Error marking connector for deletion: {e}", + file=sys.stderr, + ) + return { + "status": "error", + "message": str(e), + } + + +def main() -> None: + if len(sys.argv) != 3: + print( + json.dumps( + { + "status": "error", + "message": "Usage: python mark_connector_for_deletion.py ", + } + ) + ) + sys.exit(1) + + tenant_id = sys.argv[1] + try: + cc_pair_id = int(sys.argv[2]) + except ValueError: + print( + json.dumps( + { + "status": "error", + "message": "cc_pair_id must be an integer", + } + ) + ) + sys.exit(1) + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + result = mark_connector_for_deletion(tenant_id, cc_pair_id) + print(json.dumps(result)) + + +if __name__ == "__main__": + main() diff --git a/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py b/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py new file mode 100644 index 0000000000..67420ca09b --- /dev/null +++ b/backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py @@ -0,0 +1,96 @@ +import json +import sys +from typing import Any + +from sqlalchemy import text +from sqlalchemy.exc import ProgrammingError +from sqlalchemy.orm import Session + +from onyx.db.engine.sql_engine import get_session_with_shared_schema +from onyx.db.engine.sql_engine import SqlEngine + + +def get_tenant_activity_summary(session: Session) -> list[dict[str, Any]]: + """Return a list of dicts, one per tenant, with last query info, doc count, and user count.""" + + # Step 1: fetch all tenant schemas + tenant_schemas = [ + row[0] + for row in session.execute( + text( + """ + SELECT nspname + FROM pg_namespace + WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'public') + AND nspname NOT LIKE 'pg_toast%%' + AND nspname NOT LIKE 'pg_temp%%' + ORDER BY nspname + """ + ) + ) + ] + + print(f"Found {len(tenant_schemas)} tenant schemas", file=sys.stderr) + + summaries = [] + + # Step 2: loop through each tenant schema + for idx, schema in enumerate(tenant_schemas): + if idx % 100 == 0: + print(f"Processing tenant {idx}/{len(tenant_schemas)}", file=sys.stderr) + + try: + # Use a single query to get all data at once + query = text( + f""" + SELECT + :tenant_id AS tenant_id, + ( + SELECT time_sent + FROM "{schema}".chat_message + WHERE message_type = 'USER' + ORDER BY time_sent DESC + LIMIT 1 + ) AS last_query_time, + ( + SELECT message + FROM "{schema}".chat_message + WHERE message_type = 'USER' + ORDER BY time_sent DESC + LIMIT 1 + ) AS last_query_text, + (SELECT COUNT(*) FROM "{schema}".document) AS num_documents, + (SELECT COUNT(*) FROM "{schema}".user) AS num_users + """ + ) + + result = session.execute(query, {"tenant_id": schema}).mappings().first() + + if result: + summaries.append(dict(result)) + + except ProgrammingError as e: + # schema may be missing a table + print(f"Error processing schema {schema}: {e}", file=sys.stderr) + session.rollback() + continue + except Exception as e: + print(f"Unexpected error processing schema {schema}: {e}", file=sys.stderr) + session.rollback() + continue + + return summaries + + +def main() -> None: + + SqlEngine.init_engine(pool_size=5, max_overflow=2) + + with get_session_with_shared_schema() as session: + summaries = get_tenant_activity_summary(session) + + print(json.dumps(summaries, indent=2, default=str)) + + +if __name__ == "__main__": + main()