sea-dev 1.0.0

package/scripts/python/export-org-data.py

@@ -0,0 +1,693 @@
+#!/usr/bin/env python3
+"""
+Export organization data with cascading dependencies.
+
+This script generates a SQL dump of all data for a specific organization,
+including cascading foreign key dependencies, with organization_id remapped.
+
+Usage:
+    python export-org-data.py <source_org_id> <target_org_id> [--db-url <url>]
+
+Example:
+    python export-org-data.py rsvnsseb46qohsct2n6lh foyjflwcbvryzqjnb8e4p
+"""
+
+import sys
+import argparse
+import psycopg2
+from typing import Dict, List, Set, Tuple
+from collections import defaultdict
+
+
+# Define the schema with tables, their organization_id columns, and foreign keys
+SCHEMA = {
+    # Tables with organization_id
+    "organization": {
+        "pk": "id",
+        "org_id_col": "id",  # Special case: this IS the org table
+        "fks": {}
+    },
+    "org_settings": {
+        "pk": "organization_id",
+        "org_id_col": "organization_id",
+        "fks": {}
+    },
+    "member": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "user_id": "user.id",
+        }
+    },
+    "invitation": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "inviter_id": "user.id",
+        }
+    },
+    "workflow": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {}
+    },
+    "deal": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {}
+    },
+    "form": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "workflow_id": "workflow.id",
+            "conversation_id": "conversation.id",
+        }
+    },
+    "submission": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "form_id": "form.id",
+            "deal_id": "deal.id",
+            "conversation_id": "conversation.id",
+        }
+    },
+    "conversation": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "active_template_id": "agent_template.id",
+        }
+    },
+    "api_key": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "member_id": "member.id",
+            "target_id": "submission.id",  # Conditional based on target_type
+        }
+    },
+    "field_correction": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "submission_item_id": "submission_item.id",
+        }
+    },
+    "report": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "conversation_id": "conversation.id",
+            "deal_id": "deal.id",
+            "form_id": "form.id",
+        }
+    },
+    # Skip agent tables - schema mismatch
+    # "agent_template": {
+    #     "pk": "id",
+    #     "org_id_col": "organization_id",
+    #     "fks": {}
+    # },
+    # "agent_run": {
+    #     "pk": "id",
+    #     "org_id_col": "organization_id",
+    #     "fks": {
+    #         "conversation_id": "conversation.id",
+    #         "template_id": "agent_template.id",
+    #         "template_version_id": "agent_template_version.id",
+    #     }
+    # },
+    "pipeline": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {}
+    },
+    "pipeline_run": {
+        "pk": "id",
+        "org_id_col": "organization_id",
+        "fks": {
+            "pipeline_id": "pipeline.id",
+        }
+    },
+
+    # Tables without organization_id but with FKs to org tables
+    "form_group": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "form_id": "form.id",
+        }
+    },
+    "form_field": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "form_group_id": "form_group.id",
+        }
+    },
+    "form_derived": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "form_group_id": "form_group.id",
+        }
+    },
+    "submission_group": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "submission_id": "submission.id",
+            "form_group_id": "form_group.id",
+        }
+    },
+    "submission_item": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "submission_group_id": "submission_group.id",
+            "form_field_id": "form_field.id",
+        }
+    },
+    "submission_derived": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "submission_group_id": "submission_group.id",
+            "form_derived_id": "form_derived.id",
+        }
+    },
+    "submission_item_version": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "submission_item_id": "submission_item.id",
+            "source_id": None,  # Polymorphic - could be document_citation, message, or agent_transformation
+        }
+    },
+    "document": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "deal_id": "deal.id",
+        }
+    },
+    "document_extraction": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "document_id": "document.id",
+            "submission_id": "submission.id",
+        }
+    },
+    "document_citation": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "document_extraction_id": "document_extraction.id",
+        }
+    },
+    # Skip document_chunks - not needed for org export
+    # "document_chunks": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "document_id": "document.id",
+    #         "submission_id": "submission.id",
+    #         "parent_id": "document_chunks.id",
+    #     }
+    # },
+    "message": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "conversation_id": "conversation.id",
+        }
+    },
+    "context": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "conversation_id": "conversation.id",
+        }
+    },
+    "tool": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "conversation_id": "conversation.id",
+        }
+    },
+    "conversation_resource": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "conversation_id": "conversation.id",
+        }
+    },
+    # Skip agent tables - schema mismatch
+    # "template_form": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "template_id": "agent_template.id",
+    #         "form_id": "form.id",
+    #     }
+    # },
+    # "agent_template_step": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "template_id": "agent_template.id",
+    #         "parent_step_id": "agent_template_step.id",
+    #         "output_form_id": "form.id",
+    #     }
+    # },
+    # "agent_template_version": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "template_id": "agent_template.id",
+    #     }
+    # },
+    # "agent_run_step": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "run_id": "agent_run.id",
+    #         "parent_step_id": "agent_run_step.id",
+    #         "output_submission_id": "submission.id",
+    #     }
+    # },
+    # "agent_tool_call": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "run_step_id": "agent_run_step.id",
+    #     }
+    # },
+    # "agent_transformation": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "run_step_id": "agent_run_step.id",
+    #     }
+    # },
+    # "agent_artifact": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "conversation_id": "conversation.id",
+    #         "run_id": "agent_run.id",
+    #         "run_step_id": "agent_run_step.id",
+    #     }
+    # },
+    # "agent_run_submission": {
+    #     "pk": "id",
+    #     "org_id_col": None,
+    #     "fks": {
+    #         "run_id": "agent_run.id",
+    #         "submission_id": "submission.id",
+    #     }
+    # },
+    "pipeline_form": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "pipeline_id": "pipeline.id",
+            "form_id": "form.id",
+        }
+    },
+    "pipeline_step": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "pipeline_id": "pipeline.id",
+            "parent_step_id": "pipeline_step.id",
+            "output_form_id": "form.id",
+        }
+    },
+    "pipeline_resource": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "pipeline_id": "pipeline.id",
+        }
+    },
+    "pipeline_run_step": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "run_id": "pipeline_run.id",
+            "parent_run_step_id": "pipeline_run_step.id",
+            "output_submission_id": "submission.id",
+        }
+    },
+    "tool_call": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "tool_id": "tool.id",
+        }
+    },
+    "pipeline_run_step_tool_call": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "run_step_id": "pipeline_run_step.id",
+            "tool_call_id": "tool_call.id",
+        }
+    },
+    "artifact": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {}
+    },
+    "pipeline_run_artifact": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "run_id": "pipeline_run.id",
+            "artifact_id": "artifact.id",
+            "run_step_id": "pipeline_run_step.id",
+        }
+    },
+    "pipeline_run_submission": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {
+            "run_id": "pipeline_run.id",
+            "submission_id": "submission.id",
+        }
+    },
+    "data_transformation": {
+        "pk": "id",
+        "org_id_col": None,
+        "fks": {}
+    },
+}
+
+
+class OrgDataExporter:
+    def __init__(self, conn, source_org_id: str, target_org_id: str):
+        self.conn = conn
+        self.source_org_id = source_org_id
+        self.target_org_id = target_org_id
+        self.collected_ids: Dict[str, Set[str]] = defaultdict(set)
+        self.sql_statements: List[str] = []
+
+    def get_table_columns(self, table_name: str) -> List[str]:
+        """Get all column names for a table."""
+        cursor = self.conn.cursor()
+        cursor.execute("""
+            SELECT column_name
+            FROM information_schema.columns
+            WHERE table_name = %s
+            ORDER BY ordinal_position
+        """, (table_name,))
+        return [row[0] for row in cursor.fetchall()]
+
+    def get_column_types(self, table_name: str) -> Dict[str, str]:
+        """Get column types for a table."""
+        cursor = self.conn.cursor()
+        cursor.execute("""
+            SELECT column_name, data_type
+            FROM information_schema.columns
+            WHERE table_name = %s
+        """, (table_name,))
+        return {row[0]: row[1] for row in cursor.fetchall()}
+
+    def fetch_rows(self, table_name: str, where_clause: str, params: tuple) -> List[Dict]:
+        """Fetch rows from a table with a WHERE clause."""
+        cursor = self.conn.cursor()
+        columns = self.get_table_columns(table_name)
+
+        # Quote column names to handle reserved keywords
+        quoted_columns = [f'"{col}"' for col in columns]
+        query = f"SELECT {', '.join(quoted_columns)} FROM {table_name} WHERE {where_clause}"
+        cursor.execute(query, params)
+
+        rows = []
+        for row in cursor.fetchall():
+            rows.append(dict(zip(columns, row)))
+        return rows
+
+    def collect_table_data(self, table_name: str, initial: bool = False):
+        """Collect data for a table based on already collected IDs."""
+        config = SCHEMA.get(table_name)
+        if not config:
+            return
+
+        pk = config["pk"]
+        org_id_col = config["org_id_col"]
+
+        # Determine what IDs to fetch
+        if initial and org_id_col:
+            # Initial tables: fetch by organization_id
+            if org_id_col == "id":
+                # Special case: organization table
+                rows = self.fetch_rows(table_name, f"{org_id_col} = %s", (self.source_org_id,))
+            else:
+                rows = self.fetch_rows(table_name, f"{org_id_col} = %s", (self.source_org_id,))
+        else:
+            # Dependent tables: fetch by foreign keys
+            fks = config["fks"]
+            conditions = []
+            params = []
+
+            for fk_col, fk_ref in fks.items():
+                if not fk_ref:
+                    continue
+
+                ref_table = fk_ref.split(".")[0]
+                if ref_table in self.collected_ids and self.collected_ids[ref_table]:
+                    conditions.append(f"{fk_col} = ANY(%s)")
+                    params.append(list(self.collected_ids[ref_table]))
+
+            if not conditions:
+                return
+
+            where_clause = " OR ".join(conditions)
+            rows = self.fetch_rows(table_name, where_clause, tuple(params))
+
+        # Track new IDs and generate INSERT statements
+        new_ids = set()
+        for row in rows:
+            row_id = row[pk]
+            if row_id in self.collected_ids[table_name]:
+                continue
+
+            new_ids.add(row_id)
+            self.collected_ids[table_name].add(row_id)
+
+            # Generate INSERT statement
+            self.generate_insert(table_name, row, config)
+
+        return len(new_ids) > 0
+
+    def generate_insert(self, table_name: str, row: Dict, config: Dict):
+        """Generate an INSERT statement for a row."""
+        import json
+        from datetime import datetime, date
+
+        # Don't remap organization_id - keep everything with source org ID
+        # The target_org_id parameter is ignored
+
+        # Get column types for this table
+        column_types = self.get_column_types(table_name)
+
+        # Build INSERT statement with proper escaping using dollar quoting for complex values
+        columns = [col for col in row.keys() if row[col] is not None]
+        values = [row[col] for col in columns]
+
+        # Escape values properly
+        escaped_values = []
+        for col, val in zip(columns, values):
+            col_type = column_types.get(col, "")
+
+            if val is None:
+                escaped_values.append("NULL")
+            elif isinstance(val, bool):
+                # Must check bool before int since bool is subclass of int
+                if col_type == "jsonb":
+                    escaped_values.append(f"'{str(val).lower()}'::jsonb")
+                else:
+                    escaped_values.append("TRUE" if val else "FALSE")
+            elif isinstance(val, (int, float)):
+                if col_type == "jsonb":
+                    escaped_values.append(f"'{val}'::jsonb")
+                else:
+                    escaped_values.append(str(val))
+            elif isinstance(val, (datetime, date)):
+                # Handle datetime and date objects
+                escaped_values.append(f"'{val.isoformat()}'")
+            elif isinstance(val, str):
+                # Check if column is jsonb type
+                if col_type == "jsonb":
+                    # String value in jsonb column needs to be JSON-encoded
+                    json_str = json.dumps(val).replace("'", "''")
+                    escaped_values.append(f"'{json_str}'::jsonb")
+                else:
+                    # Regular string
+                    escaped_val = val.replace("'", "''")
+                    escaped_values.append(f"'{escaped_val}'")
+            elif isinstance(val, list):
+                # Check if column is ARRAY or jsonb
+                if col_type == "ARRAY":
+                    # PostgreSQL array - use ARRAY[] syntax
+                    if not val:
+                        escaped_values.append("ARRAY[]::text[]")
+                    else:
+                        array_items = []
+                        for item in val:
+                            if isinstance(item, str):
+                                array_items.append(f"'{item.replace(chr(39), chr(39)+chr(39))}'")
+                            else:
+                                array_items.append(f"'{str(item)}'")
+                        escaped_values.append(f"ARRAY[{','.join(array_items)}]")
+                else:
+                    # jsonb
+                    json_str = json.dumps(val).replace("'", "''")
+                    escaped_values.append(f"'{json_str}'::jsonb")
+            elif isinstance(val, dict):
+                # For jsonb objects
+                json_str = json.dumps(val).replace("'", "''")
+                escaped_values.append(f"'{json_str}'::jsonb")
+            else:
+                # Fallback: convert to string
+                str_val = str(val).replace("'", "''")
+                escaped_values.append(f"'{str_val}'")
+
+        # Quote column names to handle reserved keywords
+        quoted_columns = [f'"{col}"' for col in columns]
+        sql = f"INSERT INTO {table_name} ({', '.join(quoted_columns)}) VALUES ({', '.join(escaped_values)}) ON CONFLICT DO NOTHING;"
+        self.sql_statements.append(sql)
+
+    def export(self):
+        """Main export logic."""
+        print(f"-- Exporting data for organization {self.source_org_id}", file=sys.stderr)
+        print(f"-- Remapping to organization {self.target_org_id}", file=sys.stderr)
+
+        # Phase 1: Collect all data from tables with organization_id
+        print("-- Phase 1: Collecting org-scoped tables", file=sys.stderr)
+        org_tables = [name for name, cfg in SCHEMA.items() if cfg["org_id_col"]]
+        for table in org_tables:
+            print(f"--   Collecting {table}...", file=sys.stderr)
+            self.collect_table_data(table, initial=True)
+
+        # Phase 2: Cascade through dependent tables
+        print("-- Phase 2: Cascading through dependencies", file=sys.stderr)
+        max_iterations = 50
+        iteration = 0
+
+        while iteration < max_iterations:
+            iteration += 1
+            print(f"--   Iteration {iteration}", file=sys.stderr)
+
+            had_changes = False
+            for table_name in SCHEMA.keys():
+                if SCHEMA[table_name]["org_id_col"]:
+                    # Skip org tables, already collected
+                    continue
+
+                changed = self.collect_table_data(table_name)
+                if changed:
+                    had_changes = True
+                    print(f"--     Found new rows in {table_name}", file=sys.stderr)
+
+            if not had_changes:
+                print("--   No more changes, done!", file=sys.stderr)
+                break
+
+        # Output SQL
+        print("-- SQL Dump generated by export-org-data.py")
+        print(f"-- Source org: {self.source_org_id}")
+        print(f"-- Target org: {self.target_org_id}")
+        print()
+
+        # Group statements by table to maintain insertion order
+        table_statements = defaultdict(list)
+        for stmt in self.sql_statements:
+            # Extract table name from INSERT statement
+            table_name = stmt.split("INSERT INTO ")[1].split(" ")[0]
+            table_statements[table_name].append(stmt)
+
+        # Define insertion order respecting foreign key dependencies
+        insertion_order = [
+            "organization", "org_settings",
+            "user", "account", "session", "impersonation",
+            "member", "invitation", "api_key",
+            "workflow", "deal",
+            "pipeline", "pipeline_step", "pipeline_form", "pipeline_resource",
+            "conversation", "message", "context", "tool", "conversation_resource",
+            "form", "form_group", "form_field", "form_derived",
+            "submission", "submission_group", "submission_item", "submission_item_version", "submission_derived",
+            "document", "document_extraction", "document_citation",
+            "pipeline_run", "pipeline_run_step", "tool_call", "pipeline_run_step_tool_call",
+            "artifact", "pipeline_run_artifact", "pipeline_run_submission",
+            "data_transformation", "field_correction", "report",
+        ]
+
+        # Output statements in order with a transaction per table
+        for table in insertion_order:
+            if table in table_statements:
+                print(f"\n-- Table: {table}")
+                print("BEGIN;")
+                for stmt in table_statements[table]:
+                    print(stmt)
+                print("COMMIT;")
+
+        # Output any remaining tables not in the order list
+        for table, stmts in table_statements.items():
+            if table not in insertion_order:
+                print(f"\n-- Table: {table}")
+                print("BEGIN;")
+                for stmt in stmts:
+                    print(stmt)
+                print("COMMIT;")
+
+        # Summary
+        print(f"\n-- Summary:", file=sys.stderr)
+        for table, ids in sorted(self.collected_ids.items()):
+            if ids:
+                print(f"--   {table}: {len(ids)} rows", file=sys.stderr)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Export organization data with cascading dependencies"
+    )
+    parser.add_argument("source_org_id", help="Source organization ID")
+    parser.add_argument("target_org_id", help="Target organization ID (for remapping)")
+    parser.add_argument(
+        "--db-url",
+        help="PostgreSQL connection URL",
+        default=None
+    )
+
+    args = parser.parse_args()
+
+    # Get DB URL from args or environment
+    db_url = args.db_url
+    if not db_url:
+        import os
+        db_url = os.environ.get("DATABASE_URL")
+
+    if not db_url:
+        print("Error: No database URL provided. Use --db-url or set DATABASE_URL", file=sys.stderr)
+        sys.exit(1)
+
+    # Connect to database
+    print(f"-- Connecting to database...", file=sys.stderr)
+    conn = psycopg2.connect(db_url)
+
+    try:
+        exporter = OrgDataExporter(conn, args.source_org_id, args.target_org_id)
+        exporter.export()
+    finally:
+        conn.close()
+
+
+if __name__ == "__main__":
+    main()