"""Analyze Auth0 identities by cross-referencing with Neo4j and export results to JSON.""" import argparse import json import logging import sys from datetime import datetime from pathlib import Path from typing import Any from src.neo4j.cypher_runner import CypherRunner from src.readers import read_auth0_users # Configure logging at module level logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def format_file_size(size_bytes: int) -> str: """Format bytes as human-readable string (KB or MB).""" if size_bytes < 1024 * 1024: return f'{size_bytes / 1024:.1f} KB' return f'{size_bytes / (1024 * 1024):.1f} MB' def list_job_files(directory: str = 'data/input') -> list[tuple[Path, datetime, int]]: """Find job_*.json files with modification date and size. Returns list of (path, mtime, size) tuples sorted by date (newest first). """ input_dir = Path(directory) if not input_dir.exists(): return [] files: list[tuple[Path, datetime, int]] = [] for f in input_dir.glob('job_*.json'): stat = f.stat() files.append((f, datetime.fromtimestamp(stat.st_mtime), stat.st_size)) return sorted(files, key=lambda x: x[1], reverse=True) def prompt_file_selection(files: list[tuple[Path, datetime, int]]) -> Path: """Display numbered menu of files and return selected path.""" print('\nAvailable input files in data/input/:') for i, (path, mtime, size) in enumerate(files, 1): date_str = mtime.strftime('%Y-%m-%d %H:%M:%S') print(f' {i}) {path.name} ({date_str}, {format_file_size(size)})') while True: choice = input(f'\nSelect file (1-{len(files)}): ').strip() if choice.isdigit() and 1 <= int(choice) <= len(files): return files[int(choice) - 1][0] print(f'Invalid choice. Enter a number between 1 and {len(files)}.') def batch_process_identities( input_file: str, output_file: str, batch_size: int = 500, max_batches: int | None = None, include_profiles: bool = False, ) -> None: """ Analyze Auth0 identities in batches by querying Neo4j and export results to JSON. Combines Auth0 user data with Neo4j identity and profile data, preserving all fields. Args: input_file: Path to Auth0 export JSON file output_file: Path to output JSON file batch_size: Number of identities to process per batch max_batches: Maximum number of batches to process (None for all) include_profiles: Whether to include profile data in output (default: False) """ # Load Cypher query cypher_file = Path(__file__).parent / 'cypher' / 'get_identities_batch.cypher' query = cypher_file.read_text() # Read Auth0 users logger.info(f'Reading Auth0 data from {input_file}') auth0_users = list(read_auth0_users(input_file)) total_users = len(auth0_users) logger.info(f'Found {total_users} Auth0 users to process') # Process in batches results: list[dict[str, Any]] = [] found_count = 0 not_found_count = 0 with CypherRunner() as runner: batches_processed = 0 for i in range(0, total_users, batch_size): batch = auth0_users[i : i + batch_size] batch_num = i // batch_size + 1 total_batches = (total_users + batch_size - 1) // batch_size # Check if we've hit the max batch limit if max_batches and batches_processed >= max_batches: logger.info(f'Reached max batch limit of {max_batches}, stopping') break logger.info(f'Processing batch {batch_num}/{total_batches} ({i + 1}-{min(i + batch_size, total_users)}/{total_users})') # Collect all identity IDs for this batch identity_ids = [user.identity_id for user in batch] # Debug: log first few IDs to verify format if batch_num == 1: logger.info(f'Sample identity_ids from Auth0 data: {identity_ids[:3]}') try: # Query Neo4j with all IDs at once logger.debug(f'Querying {len(identity_ids)} identities in single batch query') query_results = runner.run_query( query, {'identityIds': identity_ids}, ) # Debug: log what came back from Neo4j if batch_num == 1: logger.info(f'Query returned {len(query_results)} results') if query_results: sample = query_results[0] logger.info(f'Sample result keys: {sample.keys()}') logger.info(f'Sample identityId from result: {sample.get("identityId")}') # Create lookup map: identity_id -> Neo4j identity data results_map: dict[str, dict[str, Any] | None] = {} for record in query_results: identity_id = record.get('identityId') identity = record.get('identity') if identity_id: results_map[identity_id] = identity # Process each user with their result for idx, user in enumerate(batch, 1): identity_id = user.identity_id # Get the Neo4j identity data (or None if not found) neo4j_identity = results_map.get(identity_id) # Build flattened result if neo4j_identity: result: dict[str, Any] = {k: v for k, v in neo4j_identity.items() if k != 'profiles' or include_profiles} result['found'] = True result['originalAuth0Id'] = user.id else: result = { 'id': identity_id, 'found': False, 'originalAuth0Id': user.id, 'email': user.email, 'name': user.name, } results.append(result) # Update counts and log if neo4j_identity: found_count += 1 profile_count = len(neo4j_identity.get('profiles', [])) logger.info(f' [{idx}/{len(batch)}] {user.id} (identity_id: {identity_id}): FOUND ({profile_count} profiles)') else: not_found_count += 1 logger.info(f' [{idx}/{len(batch)}] {user.id} (identity_id: {identity_id}): NOT FOUND') except Exception as e: logger.error(f'Error processing batch: {e}') # Add error results for all users in batch for user in batch: results.append( { 'id': user.identity_id, 'found': False, 'originalAuth0Id': user.id, 'email': user.email, 'name': user.name, 'error': str(e), } ) not_found_count += 1 batches_processed += 1 # Write results to JSON logger.info(f'Writing results to {output_file}') Path(output_file).parent.mkdir(parents=True, exist_ok=True) with open(output_file, 'w', encoding='utf-8') as f: json.dump(results, f, indent=2, default=str) # Print summary total_processed = len(results) logger.info('=' * 60) logger.info('Processing Complete!') logger.info(f'Total processed: {total_processed} (out of {total_users} total users)') logger.info(f'Found in Neo4j: {found_count}') logger.info(f'Not found: {not_found_count}') logger.info(f'Results written to: {output_file}') logger.info('=' * 60) def main() -> None: """Main entry point.""" parser = argparse.ArgumentParser(description='Analyze Auth0 identities by cross-referencing with Neo4j and export to JSON') parser.add_argument( '--input', help='Path to Auth0 export JSON file (interactive prompt if omitted)', ) parser.add_argument( '--output', help='Path to output JSON file (default: data/input/neo4j_identities_TIMESTAMP.json)', ) parser.add_argument( '--batch-size', type=int, default=500, help='Number of identities to process per batch (default: 500)', ) parser.add_argument( '--max-batches', type=int, help='Maximum number of batches to process (useful for test runs, e.g., 5)', ) parser.add_argument( '--exclude-profiles', action='store_true', help='Exclude profile data from output to reduce file size', ) args = parser.parse_args() # Interactive file selection if --input not provided if not args.input: job_files = list_job_files() if not job_files: logger.error('No job_*.json files found in data/input/') sys.exit(1) args.input = str(prompt_file_selection(job_files)) # Set default output path with timestamp if not args.output: timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') args.output = f'data/input/neo4j_identities_{timestamp}.json' batch_process_identities(args.input, args.output, args.batch_size, args.max_batches, include_profiles=not args.exclude_profiles) if __name__ == '__main__': main()