local-deep-research/scripts/pre_commit/check_datetime_timezone.py

#!/usr/bin/env python3
"""Pre-commit hook to ensure all datetime columns use UtcDateTime for SQLite compatibility."""

import ast
import re
import sys
from pathlib import Path
from typing import List, Tuple


def check_datetime_columns(file_path: Path) -> List[Tuple[int, str, str]]:
    """Check a Python file for DateTime columns that should use UtcDateTime.

    Returns a list of (line_number, line_content, error_message) tuples for violations.
    """
    violations = []

    try:
        with open(file_path, "r") as f:
            content = f.read()
            lines = content.split("\n")
    except Exception as e:
        print(f"Error reading {file_path}: {e}", file=sys.stderr)
        return violations

    # Check if file imports UtcDateTime (if it uses any DateTime columns)
    has_utc_datetime_import = (
        "from sqlalchemy_utc import UtcDateTime" in content
        or "from sqlalchemy_utc import utcnow, UtcDateTime" in content
    )

    # Parse the AST to find Column definitions with DateTime
    try:
        tree = ast.parse(content)
    except SyntaxError:
        # Not valid Python, skip
        return violations

    for node in ast.walk(tree):
        if isinstance(node, ast.Call):
            # Check if this is a Column call
            if isinstance(node.func, ast.Name) and node.func.id == "Column":
                # Check if first argument is DateTime
                if node.args and isinstance(node.args[0], ast.Call):
                    datetime_call = node.args[0]
                    if (
                        isinstance(datetime_call.func, ast.Name)
                        and datetime_call.func.id == "DateTime"
                    ):
                        # This should be UtcDateTime instead
                        line_num = node.lineno
                        if 0 <= line_num - 1 < len(lines):
                            violations.append(
                                (
                                    line_num,
                                    lines[line_num - 1].strip(),
                                    "Use UtcDateTime instead of DateTime for SQLite compatibility",
                                )
                            )
                    elif (
                        isinstance(datetime_call.func, ast.Name)
                        and datetime_call.func.id == "UtcDateTime"
                    ):
                        # This is correct, but check if import exists
                        if not has_utc_datetime_import:
                            line_num = node.lineno
                            if 0 <= line_num - 1 < len(lines):
                                violations.append(
                                    (
                                        line_num,
                                        lines[line_num - 1].strip(),
                                        "Missing import: from sqlalchemy_utc import UtcDateTime",
                                    )
                                )

    # Also check for func.now() usage which should be utcnow()
    for i, line in enumerate(lines, 1):
        if "func.now()" in line and "Column" in line:
            violations.append(
                (
                    i,
                    line.strip(),
                    "Use utcnow() instead of func.now() for timezone-aware defaults",
                )
            )
        # Check for datetime.utcnow or datetime.now(UTC) in defaults
        if re.search(
            r"default\s*=\s*(lambda:\s*)?datetime\.(utcnow|now)", line
        ):
            violations.append(
                (
                    i,
                    line.strip(),
                    "Use utcnow() from sqlalchemy_utc instead of datetime functions for defaults",
                )
            )

    return violations


def main():
    """Main entry point for the pre-commit hook."""
    files_to_check = sys.argv[1:]

    if not files_to_check:
        print("No files to check")
        return 0

    all_violations = []

    for file_path_str in files_to_check:
        file_path = Path(file_path_str)

        # Only check Python files in database/models directories
        if file_path.suffix == ".py" and (
            "database/models" in str(file_path) or "models" in file_path.parts
        ):
            violations = check_datetime_columns(file_path)
            if violations:
                all_violations.append((file_path, violations))

    if all_violations:
        print("\n❌ DateTime column issues found:\n")
        for file_path, violations in all_violations:
            print(f"  {file_path}:")
            for line_num, line_content, error_msg in violations:
                print(f"    Line {line_num}: {error_msg}")
                print(f"      > {line_content}")
        print(
            "\n  Fix: Use UtcDateTime from sqlalchemy_utc for all datetime columns"
        )
        print("  Example: ")
        print("    from sqlalchemy_utc import UtcDateTime, utcnow")
        print("    Column(UtcDateTime, default=utcnow(), ...)\n")
        return 1

    return 0


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