update_company_ids.py

import psycopg2
from psycopg2.extras import execute_values
import uuid
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

def connect_to_db():
    """Create database connection"""
    return psycopg2.connect(
        os.getenv('DATABASE_URL')
    )

def update_company_ids(conn):
    """Update existing companies with UUIDs and update related contacts"""
    try:
        cursor = conn.cursor()
        
        # First, add a new UUID column
        print("Adding UUID column...")
        add_column_query = """
        ALTER TABLE "companyProfilesData" 
        ADD COLUMN IF NOT EXISTS company_id UUID;
        """
        cursor.execute(add_column_query)
        
        # Generate and set UUIDs for all existing records
        print("Generating UUIDs for existing companies...")
        update_uuid_query = """
        UPDATE "companyProfilesData"
        SET company_id = gen_random_uuid()
        WHERE company_id IS NULL;
        """
        cursor.execute(update_uuid_query)
        
        # Add a new UUID column to contacts table for company reference
        print("Adding company_id column to contacts table...")
        add_contact_column_query = """
        ALTER TABLE contacts 
        ADD COLUMN IF NOT EXISTS company_id UUID;
        """
        cursor.execute(add_contact_column_query)
        
        # Update contacts with corresponding company_ids
        print("Updating contacts with company references...")
        update_contacts_query = """
        UPDATE contacts c
        SET company_id = cp.company_id
        FROM "companyProfilesData" cp
        WHERE c.domain = cp.domain;
        """
        cursor.execute(update_contacts_query)
        
        # Make company_id the new primary key
        print("Updating primary key...")
        update_pk_queries = [
            # Drop the existing primary key constraint
            """ALTER TABLE "companyProfilesData" DROP CONSTRAINT IF EXISTS "companyProfilesData_pkey";""",
            
            # Make company_id NOT NULL
            """ALTER TABLE "companyProfilesData" ALTER COLUMN company_id SET NOT NULL;""",
            
            # Add new primary key constraint
            """ALTER TABLE "companyProfilesData" ADD PRIMARY KEY (company_id);""",
            
            # Add index on domain
            """CREATE INDEX IF NOT EXISTS idx_company_domain ON "companyProfilesData"(domain);"""
        ]
        
        for query in update_pk_queries:
            cursor.execute(query)
        
        # Update foreign key in contacts table
        print("Updating foreign key constraints...")
        update_fk_queries = [
            # Drop the existing foreign key if it exists
            """ALTER TABLE contacts DROP CONSTRAINT IF EXISTS contacts_company_id_fkey;""",
            
            # Add new foreign key constraint
            """ALTER TABLE contacts 
               ADD CONSTRAINT contacts_company_id_fkey 
               FOREIGN KEY (company_id) 
               REFERENCES "companyProfilesData"(company_id);"""
        ]
        
        for query in update_fk_queries:
            cursor.execute(query)
        
        # Commit all changes
        conn.commit()
        print("Successfully updated company IDs and related references")
        
    except Exception as e:
        conn.rollback()
        print(f"Error updating company IDs: {str(e)}")
        raise
    finally:
        cursor.close()

def main():
    """Main function to update company IDs"""
    conn = connect_to_db()
    try:
        update_company_ids(conn)
    finally:
        conn.close()

if __name__ == "__main__":
    main()