annator-atom / backend /scripts /migrate_db_sessions.py
techprotrade's picture
Full stack ATOM backend + AIMONEYFLOW clients (port 7860) (part 5)
f84a02d verified
Raw
History Blame Contribute Delete
7.13 kB
#!/usr/bin/env python3
"""
Database Session Migration Script
Identifies and helps migrate manual database session management to the
context manager pattern.
Usage:
python scripts/migrate_db_sessions.py # Identify files to migrate
python scripts/migrate_db_sessions.py --fix # Auto-fix simple cases
"""
import ast
import os
from pathlib import Path
import re
from typing import List, Tuple
def find_manual_session_patterns(file_path: str) -> List[Tuple[int, str, str]]:
"""
Find manual session management patterns in a Python file.
Returns list of (line_number, pattern_type, matched_line)
"""
patterns = [
(r'SessionLocal\(\)', 'Direct SessionLocal() call'),
(r'with\s+SessionLocal\(\)', 'Manual with SessionLocal()'),
(r'db\s*=\s*SessionLocal\(\)', 'Variable assignment'),
(r'\.close\(\)', 'Manual close() call'),
(r'\.commit\(\)', 'Manual commit() call'),
]
findings = []
try:
with open(file_path, 'r', encoding='utf-8') as f:
lines = f.readlines()
for line_num, line in enumerate(lines, 1):
for pattern, description in patterns:
if re.search(pattern, line):
findings.append((line_num, description, line.strip()))
break # Only report first match per line
except Exception as e:
print(f"Error reading {file_path}: {e}")
return findings
def scan_directory(directory: str, exclude_dirs: List[str] = None) -> dict:
"""
Scan directory for Python files with manual session management.
"""
if exclude_dirs is None:
exclude_dirs = ['venv', '__pycache__', '.pytest_cache',
'node_modules', '.git', 'migrations', 'alembic']
results = {
'files_with_manual_sessions': [],
'total_files_scanned': 0,
'files_with_issues': {}
}
for root, dirs, files in os.walk(directory):
# Remove excluded directories
dirs[:] = [d for d in dirs if d not in exclude_dirs]
for file in files:
if file.endswith('.py'):
file_path = os.path.join(root, file)
results['total_files_scanned'] += 1
findings = find_manual_session_patterns(file_path)
if findings:
results['files_with_manual_sessions'].append(file_path)
results['files_with_issues'][file_path] = findings
return results
def categorize_by_priority(results: dict) -> dict:
"""
Categorize files by migration priority.
"""
high_priority = []
medium_priority = []
low_priority = []
for file_path, issues in results['files_with_issues'].items():
line_count = len(issues)
# High priority: Service layer files with multiple issues
if any(path in file_path for path in ['service', 'services', 'core']):
if line_count >= 3:
high_priority.append((file_path, line_count))
else:
medium_priority.append((file_path, line_count))
# Medium priority: API routes, integrations
elif any(path in file_path for path in ['api', 'integrations']):
if line_count >= 3:
medium_priority.append((file_path, line_count))
else:
low_priority.append((file_path, line_count))
# Low priority: Scripts, tests, tools
else:
low_priority.append((file_path, line_count))
# Sort by issue count (descending)
high_priority.sort(key=lambda x: x[1], reverse=True)
medium_priority.sort(key=lambda x: x[1], reverse=True)
low_priority.sort(key=lambda x: x[1], reverse=True)
return {
'high': high_priority,
'medium': medium_priority,
'low': low_priority
}
def main():
"""Main entry point."""
import argparse
parser = argparse.ArgumentParser(description='Migrate database session management')
parser.add_argument('--directory', default='.', help='Directory to scan')
parser.add_argument('--fix', action='store_true', help='Auto-fix simple cases')
parser.add_argument('--output', help='Output file for results')
args = parser.parse_args()
print("=" * 80)
print("Database Session Migration Scanner")
print("=" * 80)
print()
print(f"Scanning directory: {args.directory}")
print()
results = scan_directory(args.directory)
print(f"Files scanned: {results['total_files_scanned']}")
print(f"Files with manual sessions: {len(results['files_with_manual_sessions'])}")
print()
if not results['files_with_manual_sessions']:
print("✅ No files with manual session management found!")
return
# Categorize by priority
categorized = categorize_by_priority(results)
# Print results
print("Priority Classification:")
print("-" * 80)
if categorized['high']:
print(f"\n🔴 HIGH PRIORITY ({len(categorized['high'])} files):")
print(" Service layer files with multiple manual session patterns")
for file_path, count in categorized['high'][:10]:
rel_path = os.path.relpath(file_path, args.directory)
print(f" - {rel_path} ({count} issues)")
if categorized['medium']:
print(f"\n🟡 MEDIUM PRIORITY ({len(categorized['medium'])} files):")
print(" API routes and integrations")
for file_path, count in categorized['medium'][:10]:
rel_path = os.path.relpath(file_path, args.directory)
print(f" - {rel_path} ({count} issues)")
if categorized['low']:
print(f"\n🟢 LOW PRIORITY ({len(categorized['low'])} files):")
print(" Scripts, tests, and tools")
for file_path, count in categorized['low'][:10]:
rel_path = os.path.relpath(file_path, args.directory)
print(f" - {rel_path} ({count} issues)")
print()
print("=" * 80)
print(f"Total: {len(results['files_with_manual_sessions'])} files need migration")
print("=" * 80)
# Write output file if requested
if args.output:
with open(args.output, 'w') as f:
f.write("# Database Session Migration Report\n\n")
f.write(f"Total files: {len(results['files_with_manual_sessions'])}\n\n")
f.write("## High Priority\n\n")
for file_path, count in categorized['high']:
rel_path = os.path.relpath(file_path, args.directory)
f.write(f"- {rel_path} ({count} issues)\n")
f.write("\n## Medium Priority\n\n")
for file_path, count in categorized['medium']:
rel_path = os.path.relpath(file_path, args.directory)
f.write(f"- {rel_path} ({count} issues)\n")
f.write("\n## Low Priority\n\n")
for file_path, count in categorized['low']:
rel_path = os.path.relpath(file_path, args.directory)
f.write(f"- {rel_path} ({count} issues)\n")
print(f"\nReport written to: {args.output}")
if __name__ == '__main__':
main()