stones/scripts/utils/resolve_all_ens.py

#!/usr/bin/env python3
"""
Resolve ENS Names and Contact Information for All Contacts

This script resolves ENS names and additional contact information for all contacts
in the database that have Ethereum addresses. It uses the existing ENS resolver utility
to fetch ENS names and text records containing social profiles and contact information.

Usage:
    python resolve_all_ens.py [--batch-size BATCH_SIZE] [--delay DELAY]
"""

import os
import sys
import logging
import time
import argparse
from typing import Dict, Any, List, Optional, Tuple
from web3 import Web3
from dotenv import load_dotenv

# Add parent directory to path to import utils
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.db_connector import DatabaseConnector
from utils.ens_resolver import ENSResolver
from utils.logger import setup_logger

# Load environment variables
load_dotenv()

# Setup logging
logger = setup_logger("all_ens_resolver")

class AllContactsENSResolver:
    """Resolver for ENS names and contact information for all contacts"""
    
    def __init__(self):
        """Initialize the resolver"""
        # Initialize database
        self.db = DatabaseConnector()
        
        # Initialize Web3 and ENS resolver
        alchemy_api_key = os.getenv("ALCHEMY_API_KEY")
        if not alchemy_api_key:
            raise ValueError("ALCHEMY_API_KEY not found in environment variables")
            
        self.web3 = Web3(Web3.HTTPProvider(f"https://eth-mainnet.g.alchemy.com/v2/{alchemy_api_key}"))
        self.ens_resolver = ENSResolver(self.web3)
        
        # Register data source
        self.data_source_id = self.register_data_source()
    
    def register_data_source(self) -> str:
        """Register the ENS data source in the database"""
        return self.db.upsert_data_source(
            name="ENS Resolver",
            source_type="blockchain",
            description="ENS names and profile information resolved from Ethereum addresses"
        )
    
    def get_contacts_without_ens(self) -> List[Dict[str, Any]]:
        """Get all contacts that have an Ethereum address but no ENS name"""
        query = """
        SELECT id, "ethereumAddress", name
        FROM "Contact"
        WHERE "ethereumAddress" IS NOT NULL
        AND "ensName" IS NULL
        """
        
        result = self.db.execute_query(query)
        logger.info(f"Found {len(result)} contacts without ENS names")
        return result
    
    def get_all_contacts_with_eth_address(self) -> List[Dict[str, Any]]:
        """Get all contacts that have an Ethereum address"""
        query = """
        SELECT id, "ethereumAddress", "ensName", name, twitter, discord, telegram, email, farcaster
        FROM "Contact"
        WHERE "ethereumAddress" IS NOT NULL
        """
        
        result = self.db.execute_query(query)
        logger.info(f"Found {len(result)} contacts with Ethereum addresses")
        return result
    
    def process_contact(self, contact: Dict[str, Any]) -> Tuple[bool, bool]:
        """
        Process a single contact to resolve ENS name and contact info
        
        Args:
            contact: Contact data from the database
            
        Returns:
            Tuple of (ens_updated, info_updated) booleans
        """
        contact_id = contact["id"]
        address = contact["ethereumAddress"]
        current_ens = contact.get("ensName")
        
        ens_updated = False
        info_updated = False
        
        # Skip if no address
        if not address:
            return ens_updated, info_updated
        
        # Resolve ENS name if not already set
        if not current_ens:
            ens_name = self.ens_resolver.get_ens_name(address)
            if ens_name:
                # Update contact with ENS name
                self.db.update_contact(contact_id, {"ensName": ens_name})
                logger.info(f"Updated contact {contact_id} with ENS name: {ens_name}")
                current_ens = ens_name
                ens_updated = True
        
        # Get contact info from ENS text records if we have an ENS name
        if current_ens:
            # Update profile from ENS
            self.ens_resolver.update_contact_from_ens(contact_id, current_ens)
            info_updated = True
            
            # Link to data source
            self.db.link_contact_to_data_source(contact_id, self.data_source_id)
        
        return ens_updated, info_updated
    
    def run(self, batch_size: int = 50, delay_seconds: float = 0.5, resolve_all: bool = False):
        """
        Run the resolver for contacts
        
        Args:
            batch_size: Number of contacts to process in each batch
            delay_seconds: Delay between processing contacts
            resolve_all: Whether to process all contacts or just those without ENS names
            
        Returns:
            Tuple of (ens_updated_count, info_updated_count)
        """
        # Create a scraping job
        job_id = self.db.create_scraping_job("ENS Resolver", "running")
        logger.info(f"Created scraping job with ID: {job_id}")
        
        try:
            if resolve_all:
                contacts = self.get_all_contacts_with_eth_address()
            else:
                contacts = self.get_contacts_without_ens()
            
            if not contacts:
                logger.info("No contacts found to process")
                self.db.update_scraping_job(job_id, "completed")
                return 0, 0
            
            ens_updated_count = 0
            info_updated_count = 0
            
            # Process in batches to avoid rate limiting
            for i in range(0, len(contacts), batch_size):
                batch = contacts[i:i+batch_size]
                logger.info(f"Processing batch {i//batch_size + 1}/{(len(contacts) + batch_size - 1)//batch_size}")
                
                for contact in batch:
                    ens_updated, info_updated = self.process_contact(contact)
                    
                    if ens_updated:
                        ens_updated_count += 1
                    if info_updated:
                        info_updated_count += 1
                    
                    # Add a small delay to avoid rate limiting
                    time.sleep(delay_seconds)
                
                # Update the scraping job
                self.db.update_scraping_job(
                    job_id, 
                    "running", 
                    records_processed=len(batch),
                    records_updated=sum(1 for c in batch if self.process_contact(c)[0] or self.process_contact(c)[1])
                )
            
            # Complete the scraping job
            self.db.update_scraping_job(
                job_id, 
                "completed", 
                records_processed=len(contacts),
                records_added=ens_updated_count,
                records_updated=info_updated_count
            )
            
            logger.info(f"Updated ENS names for {ens_updated_count} contacts and contact info for {info_updated_count} contacts out of {len(contacts)} processed")
            return ens_updated_count, info_updated_count
            
        except Exception as e:
            # Update the scraping job with error
            self.db.update_scraping_job(job_id, "failed", error_message=str(e))
            logger.exception(f"Error resolving ENS names: {e}")
            raise

def main():
    """Main function"""
    try:
        parser = argparse.ArgumentParser(description="Resolve ENS names and contact information for all contacts")
        parser.add_argument("--all", action="store_true", help="Process all contacts with Ethereum addresses, not just those without ENS names")
        parser.add_argument("--batch-size", type=int, default=50, help="Number of contacts to process in each batch")
        parser.add_argument("--delay", type=float, default=0.5, help="Delay in seconds between processing contacts")
        
        args = parser.parse_args()
        
        resolver = AllContactsENSResolver()
        ens_count, info_count = resolver.run(
            batch_size=args.batch_size,
            delay_seconds=args.delay,
            resolve_all=args.all
        )
        
        logger.info(f"ENS resolution completed successfully. Updated {ens_count} ENS names and {info_count} contact info records.")
        return 0
    except Exception as e:
        logger.exception(f"Error running ENS resolver: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())