auditrum 0.2.0__py3-none-any.whl

auditrum/cli.py

@@ -0,0 +1,237 @@
+from typing import Callable, Optional, Tuple
+
+import structlog
+import typer
+from dotenv import load_dotenv
+from psycopg import connect
+from rich import print as rich_print
+from settings import PgAuditSettings
+
+from auditrum.revert import generate_revert_sql_from_log
+from auditrum.schema import (
+    generate_auditlog_partitions_sql,
+    generate_auditlog_table_sql,
+)
+from auditrum.triggers import generate_trigger_sql
+
+load_dotenv()
+
+# Helper function to get DB DSN from arguments or environment
+def get_db_dsn(
+    dsn: Optional[str] = None,
+    host: Optional[str] = None,
+    port: Optional[int] = None,
+    user: Optional[str] = None,
+    password: Optional[str] = None,
+    dbname: Optional[str] = None,
+) -> Optional[str]:
+    if dsn:
+        return dsn
+    if all([host, port, user, password, dbname]):
+        return f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
+    return settings.db_dsn
+
+
+app = typer.Typer()
+
+
+@app.callback()
+def main(
+    verbose: Optional[int] = typer.Option(
+        1, "--verbose", "-v", help="Verbosity level (0=quiet, 1=default, 2=debug)"
+    ),
+):
+    level = structlog.INFO
+    if verbose == 0:
+        level = structlog.NOTSET
+    elif verbose == 2:
+        level = structlog.DEBUG
+
+    structlog.configure(
+        wrapper_class=structlog.make_filtering_bound_logger(level),
+        processors=[
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.dev.ConsoleRenderer(),
+        ],
+    )
+
+
+log = structlog.get_logger()
+
+settings = PgAuditSettings()
+
+
+@app.command()
+def generate_trigger(
+    table: str,
+    audit_table: str = settings.audit_table,
+    dry_run: bool = typer.Option(False, "--dry-run", help="Print SQL instead of executing"),
+    output: Optional[str] = typer.Option(
+        None, "--output", "-o", help="Write SQL to file instead of printing"
+    ),
+    dsn: Optional[str] = typer.Option(None, "--db-dsn"),
+    host: Optional[str] = typer.Option(None, "--db-host"),
+    port: Optional[int] = typer.Option(None, "--db-port"),
+    user: Optional[str] = typer.Option(None, "--db-user"),
+    password: Optional[str] = typer.Option(None, "--db-password"),
+    dbname: Optional[str] = typer.Option(None, "--db-name"),
+):
+    log.info("Generating trigger SQL")
+    sql = generate_trigger_sql(table_name=table, audit_table=audit_table)
+    db_dsn = get_db_dsn(dsn, host, port, user, password, dbname)
+    execute_or_print_sql(lambda conn: sql, (), dry_run, output, db_dsn=db_dsn)
+
+
+@app.command()
+def init_schema(
+    audit_table: str = settings.audit_table,
+    dry_run: bool = typer.Option(False, "--dry-run", help="Print SQL instead of executing"),
+    output: Optional[str] = typer.Option(
+        None, "--output", "-o", help="Write SQL to file instead of printing"
+    ),
+    dsn: Optional[str] = typer.Option(None, "--db-dsn"),
+    host: Optional[str] = typer.Option(None, "--db-host"),
+    port: Optional[int] = typer.Option(None, "--db-port"),
+    user: Optional[str] = typer.Option(None, "--db-user"),
+    password: Optional[str] = typer.Option(None, "--db-password"),
+    dbname: Optional[str] = typer.Option(None, "--db-name"),
+):
+    log.info("Generating auditlog schema SQL")
+    sql = generate_auditlog_table_sql(audit_table)
+    db_dsn = get_db_dsn(dsn, host, port, user, password, dbname)
+    execute_or_print_sql(lambda conn: sql, (), dry_run, output, db_dsn=db_dsn)
+
+
+@app.command()
+def create_partitions(
+    audit_table: str = settings.audit_table,
+    months: int = 3,
+    dry_run: bool = typer.Option(False, "--dry-run", help="Print SQL instead of executing"),
+    output: Optional[str] = typer.Option(
+        None, "--output", "-o", help="Write SQL to file instead of printing"
+    ),
+    dsn: Optional[str] = typer.Option(None, "--db-dsn"),
+    host: Optional[str] = typer.Option(None, "--db-host"),
+    port: Optional[int] = typer.Option(None, "--db-port"),
+    user: Optional[str] = typer.Option(None, "--db-user"),
+    password: Optional[str] = typer.Option(None, "--db-password"),
+    dbname: Optional[str] = typer.Option(None, "--db-name"),
+):
+    log.info("Generating partition SQL", months=months)
+    sql = generate_auditlog_partitions_sql(audit_table, months)
+    db_dsn = get_db_dsn(dsn, host, port, user, password, dbname)
+    execute_or_print_sql(lambda conn: sql, (), dry_run, output, db_dsn=db_dsn)
+
+
+@app.command()
+def revert(
+    table: str,
+    record_id: str,
+    log_id: int,
+    audit_table: str = settings.audit_table,
+    dry_run: bool = typer.Option(False, "--dry-run", help="Print SQL instead of executing"),
+    output: Optional[str] = typer.Option(
+        None, "--output", "-o", help="Write SQL to file instead of printing"
+    ),
+    dsn: Optional[str] = typer.Option(None, "--db-dsn"),
+    host: Optional[str] = typer.Option(None, "--db-host"),
+    port: Optional[int] = typer.Option(None, "--db-port"),
+    user: Optional[str] = typer.Option(None, "--db-user"),
+    password: Optional[str] = typer.Option(None, "--db-password"),
+    dbname: Optional[str] = typer.Option(None, "--db-name"),
+):
+    log.info("Generating revert SQL", table=table, record_id=record_id, log_id=log_id)
+    db_dsn = get_db_dsn(dsn, host, port, user, password, dbname)
+    execute_or_print_sql(
+        sql_fn=generate_revert_sql_from_log,
+        sql_args=(audit_table, table, record_id, log_id),
+        dry_run=dry_run,
+        output=output,
+        db_dsn=db_dsn,
+    )
+
+
+@app.command()
+def status(
+    audit_table: str = settings.audit_table,
+    dsn: Optional[str] = typer.Option(None, "--db-dsn"),
+    host: Optional[str] = typer.Option(None, "--db-host"),
+    port: Optional[int] = typer.Option(None, "--db-port"),
+    user: Optional[str] = typer.Option(None, "--db-user"),
+    password: Optional[str] = typer.Option(None, "--db-password"),
+    dbname: Optional[str] = typer.Option(None, "--db-name"),
+):
+    log.info("Checking audit trigger and table status")
+
+    db_dsn = get_db_dsn(dsn, host, port, user, password, dbname)
+
+    try:
+        with connect(dsn=db_dsn) as conn, conn.cursor() as cur:
+            # Check triggers
+            cur.execute("""
+                SELECT event_object_table AS "table", trigger_name
+                FROM information_schema.triggers
+                WHERE action_statement ILIKE '%audit_trigger_fn%'
+            """)
+            triggers = cur.fetchall()
+
+            # Check partitions
+            cur.execute(f"""
+                SELECT inhrelid::regclass AS partition
+                FROM pg_inherits
+                WHERE inhparent = '{audit_table}'::regclass
+            """)
+            partitions = cur.fetchall()
+
+            # Check auditlog size and row count
+            cur.execute(f"SELECT COUNT(*) FROM {audit_table}")
+            row_count = cur.fetchone()[0]
+
+        log.info("Audit triggers found", count=len(triggers))
+        for table, trig in triggers:
+            rich_print(f"[yellow]- {table}: {trig}[/yellow]")
+
+        log.info("Auditlog partitions", count=len(partitions))
+        for p in partitions:
+            rich_print(f"[blue]- {p[0]}[/blue]")
+
+        log.info("Auditlog row count", rows=row_count)
+
+    except Exception as e:
+        log.error("Failed to check audit status", error=str(e))
+
+
+def display_sql(sql: str, output: Optional[str] = None):
+    if output:
+        with open(output, "w") as f:
+            f.write(sql + "\n")
+        log.info("SQL written to file", path=output)
+    else:
+        rich_print(f"[bold green]{sql}[/bold green]")
+
+
+def execute_or_print_sql(
+    sql_fn: Callable[..., str],
+    sql_args: Tuple,
+    dry_run: bool,
+    output: Optional[str] = None,
+    db_dsn: Optional[str] = None,
+):
+    try:
+        with connect(dsn=db_dsn) as conn:
+            sql = sql_fn(conn, *sql_args)
+
+            if dry_run:
+                display_sql(sql, output)
+                return
+
+            with conn.cursor() as cur:
+                cur.execute(sql)
+            conn.commit()
+            log.info("SQL successfully executed")
+    except Exception as e:
+        log.error("Failed to execute SQL", error=str(e))
+
+
+if __name__ == "__main__":
+    app()

auditrum/context.py

@@ -0,0 +1,131 @@
+from contextlib import contextmanager
+from contextvars import ContextVar
+from functools import wraps
+from typing import Any, Callable, Dict
+
+
+class AuditContext:
+    def __init__(self):
+        self._data: ContextVar[Dict[str, Any]] = ContextVar("audit_ctx_data")
+        self._reason_stack: ContextVar[list[str]] = ContextVar("audit_ctx_reason")
+
+    def _ensure_data(self) -> dict:
+        try:
+            return self._data.get()
+        except LookupError:
+            value = {}
+            self._data.set(value)
+            return value
+
+    def _ensure_reason_stack(self) -> list[str]:
+        try:
+            return self._reason_stack.get()
+        except LookupError:
+            value = []
+            self._reason_stack.set(value)
+            return value
+
+    def set(self, key: str, value: Any):
+        ctx = self._ensure_data().copy()
+        ctx[key] = value
+        self._data.set(ctx)
+
+    def get(self, key: str) -> Any:
+        return self._ensure_data().get(key)
+
+    def push_change_reason(self, reason: str):
+        stack = self._ensure_reason_stack().copy()
+        stack.append(reason)
+        self._reason_stack.set(stack)
+
+    def get_change_reason(self) -> str:
+        return " -> ".join(self._ensure_reason_stack())
+
+    def build_sql(self) -> str:
+        ctx = self._ensure_data().copy()
+        reason = self.get_change_reason()
+        if reason:
+            ctx["change_reason"] = reason
+
+        lines = []
+        for k, v in ctx.items():
+            key = f"session.myapp_{k}"
+            val = str(v).replace("'", "''")
+            lines.append(f"SET {key} = '{val}';")
+        return "\n".join(lines)
+
+    @contextmanager
+    def use(self, reset: bool = False, **kwargs):
+        from django.db import connection
+
+        original_data = self._ensure_data().copy()
+        ctx = original_data.copy()
+        ctx.update(kwargs)
+        self._data.set(ctx)
+
+        # SET session variables
+        sql = self.build_sql()
+        with connection.cursor() as cursor:
+            for stmt in sql.strip().split(";"):
+                if stmt.strip():
+                    cursor.execute(stmt)
+
+        try:
+            yield
+        finally:
+            self._data.set(original_data)
+            if reset:
+                with connection.cursor() as cursor:
+                    for k in ctx.keys():
+                        cursor.execute(f"RESET session.myapp_{k}")
+
+    @contextmanager
+    def use_change_reason(self, reason: str, reset: bool = False):
+        from django.db import connection
+
+        original_stack = self._ensure_reason_stack().copy()
+        stack = original_stack.copy()
+        stack.append(reason)
+        self._reason_stack.set(stack)
+
+        # SET session.myapp_change_reason
+        sql = self.build_sql()
+        with connection.cursor() as cursor:
+            for stmt in sql.strip().split(";"):
+                if stmt.strip():
+                    cursor.execute(stmt)
+
+        try:
+            yield
+        finally:
+            self._reason_stack.set(original_stack)
+            if reset:
+                with connection.cursor() as cursor:
+                    cursor.execute("RESET session.myapp_change_reason")
+
+
+audit_context = AuditContext()
+
+
+def with_context(**kwargs):
+    def decorator(func: Callable):
+        @wraps(func)
+        def wrapper(*args, **fkwargs):
+            with audit_context.use(**kwargs):
+                return func(*args, **fkwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def with_change_reason(reason: str):
+    def decorator(func: Callable):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with audit_context.use_change_reason(reason):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

auditrum/integrations/django/admin.py

@@ -0,0 +1,34 @@
+from django.contrib import admin
+from django.utils.html import format_html
+
+from auditrum.integrations.django.models import AuditLog
+
+
+@admin.register(AuditLog)
+class AuditLogAdmin(admin.ModelAdmin):
+    list_display = (
+        "changed_at",
+        "operation",
+        "table_name",
+        "object_id",
+        "user_id",
+        "source",
+        "request_id",
+        "linked_object",
+    )
+    list_filter = ("operation", "table_name", "source")
+    search_fields = ("object_id", "request_id", "change_reason")
+    readonly_fields = [f.name for f in AuditLog._meta.fields]
+    ordering = ("-changed_at",)
+
+    def linked_object(self, obj):
+        if obj.content_object:
+            return format_html(
+                '<a href="{}">{}</a>',
+                obj.content_object.get_absolute_url(),
+                obj.content_object,
+            )
+        return "-"
+
+    linked_object.allow_tags = True
+    linked_object.short_description = "Linked Object"

auditrum/integrations/django/apps.py

@@ -0,0 +1,23 @@
+from django.apps import AppConfig
+from django.db.models.signals import post_migrate
+
+
+class PgAuditIntegrationConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "auditrum.integrations.django"
+    label = "auditrum_django"
+    verbose_name = "PostgreSQL Audit (auditrum)"
+
+    def ready(self):
+        from .hooks import register_post_migrate_hook
+        from django.apps import apps
+        from importlib import import_module
+
+        post_migrate.connect(register_post_migrate_hook, sender=self)
+
+        # Automatically discover and import audit.py from all installed apps
+        for app_config in apps.get_app_configs():
+            try:
+                import_module(f"{app_config.name}.audit")
+            except ModuleNotFoundError:
+                continue

auditrum/integrations/django/audit.py

@@ -0,0 +1,27 @@
+from typing import Any, Callable, Dict, Type
+
+from django.db.models import Model
+
+registry: Dict[Type[Model], Dict[str, Any]] = {}
+
+
+def register(
+    model_cls: Type[Model],
+    *,
+    track_only=None,
+    exclude_fields=None,
+    log_conditions=None,
+    meta_fields=None
+) -> None:
+    if track_only and exclude_fields:
+        raise ValueError(
+            f"Cannot define both `track_only` and `exclude_fields` on model {model_cls.__name__}"
+        )
+    registry[model_cls] = {
+        "table_name": model_cls._meta.db_table,
+        "fields": [f.name for f in model_cls._meta.get_fields() if f.concrete],
+        "track_only": track_only,
+        "exclude_fields": exclude_fields,
+        "log_conditions": log_conditions,
+        "meta_fields": meta_fields,
+    }

auditrum/integrations/django/hooks.py

@@ -0,0 +1,29 @@
+import structlog
+from django.db import connection
+
+from auditrum.integrations.django.audit import registry
+from auditrum.integrations.django.settings import audit_settings
+from auditrum.triggers import generate_trigger_sql
+
+log = structlog.get_logger()
+
+
+def register_post_migrate_hook(sender, **kwargs):
+    with connection.cursor() as cursor:
+        for _model_cls, config in registry.items():
+            table_name = config["table_name"]
+            track_only = config["track_only"]
+            exclude_fields = config["exclude_fields"]
+            log_conditions = config["log_conditions"]
+            meta_fields = config["meta_fields"]
+
+            sql = generate_trigger_sql(
+                table_name=table_name,
+                track_only=track_only,
+                exclude_fields=exclude_fields,
+                log_conditions=log_conditions,
+                meta_fields=meta_fields,
+                audit_table=audit_settings.table_name,
+            )
+            log.info("Applying audit trigger", table=table_name)
+            cursor.execute(sql)

auditrum/integrations/django/management/commands/audit_add_partitions.py

@@ -0,0 +1,40 @@
+
+
+import calendar
+from datetime import date
+from django.core.management.base import BaseCommand
+from django.conf import settings
+from auditrum.schema import generate_auditlog_partitions_sql
+import psycopg
+
+
+class Command(BaseCommand):
+    help = "Create partitions for the auditlog table N months ahead"
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--months",
+            type=int,
+            default=1,
+            help="Number of months ahead to create partitions for (default: 1)",
+        )
+
+    def handle(self, *args, **options):
+        months = options["months"]
+        sql = generate_auditlog_partitions_sql(months_ahead=months)
+        self.stdout.write(f"Executing SQL to create partitions for {months} month(s)...")
+
+        from psycopg.conninfo import make_conninfo
+
+        db_config = settings.DATABASES["default"]
+        db_dsn = make_conninfo(
+            host=db_config.get("HOST") or "localhost",
+            port=db_config.get("PORT") or "5432",
+            dbname=db_config["NAME"],
+            user=db_config["USER"],
+            password=db_config["PASSWORD"],
+        )
+
+        with psycopg.connect(db_dsn) as conn, conn.cursor() as cur:
+            cur.execute(sql)
+        self.stdout.write(self.style.SUCCESS("Partitions created successfully."))

auditrum/integrations/django/middleware.py

@@ -0,0 +1,25 @@
+import uuid
+
+from django.utils.deprecation import MiddlewareMixin
+
+from auditrum.context import audit_context as ctx
+
+
+class RequestIDMiddleware(MiddlewareMixin):
+    def process_request(self, request):
+        request.request_id = str(uuid.uuid4())
+        ctx.set("request_id", request.request_id)
+
+
+class AuditrumMiddleware(MiddlewareMixin):
+    def process_request(self, request):
+        user = getattr(request, "user", None)
+        session = getattr(request, "session", None)
+
+        ctx.set("user_id", getattr(user, "id", None))
+        ctx.set("username", getattr(user, "username", None))
+        ctx.set("client_ip", request.META.get("REMOTE_ADDR"))
+        ctx.set("user_agent", request.META.get("HTTP_USER_AGENT"))
+        ctx.set("session_key", getattr(session, "session_key", None))
+        ctx.set("request_id", str(uuid.uuid4()))
+        ctx.set("source", "http")

auditrum/integrations/django/migrations/0001_initial.py

@@ -0,0 +1,24 @@
+from django.db import migrations
+
+from auditrum.integrations.django.settings import audit_settings
+from auditrum.schema import (
+    generate_auditlog_table_sql,
+    generate_auditlog_partitions_sql,
+    generate_jsonb_diff_function_sql,
+)
+
+
+class Migration(migrations.Migration):
+    initial = True
+    dependencies = []
+
+    operations = [
+        migrations.RunSQL(
+            sql="\n\n".join([
+                generate_auditlog_table_sql(audit_settings.table_name),
+                generate_jsonb_diff_function_sql(),
+                generate_auditlog_partitions_sql(audit_settings.table_name, months_ahead=3),
+            ]),
+            reverse_sql=f"DROP TABLE IF EXISTS {audit_settings.table_name} CASCADE;",
+        ),
+    ]

auditrum/integrations/django/mixins.py

@@ -0,0 +1,46 @@
+from apps.audit.models import AuditLog
+from django.contrib.contenttypes.models import ContentType
+from django.core.paginator import EmptyPage, PageNotAnInteger, Paginator
+from django.shortcuts import render
+from django.urls import path
+
+
+class AuditHistoryMixin:
+    def get_urls(self):
+        urls = super().get_urls()
+        custom_urls = [
+            path(
+                "<path:object_id>/history/",
+                self.admin_site.admin_view(self.object_history_view),
+                name=f"{self.model._meta.app_label}_{self.model._meta.model_name}_audit_history",
+            ),
+        ]
+        return custom_urls + urls
+
+    def object_history_view(self, request, object_id):
+        obj = self.get_object(request, object_id)
+        content_type = ContentType.objects.get_for_model(self.model)
+        logs = AuditLog.objects.filter(
+            content_type=content_type, object_id=str(object_id)
+        ).order_by("-changed_at")
+
+        paginator = Paginator(logs, 20)
+        page_number = request.GET.get("page", 1)
+
+        try:
+            page_obj = paginator.page(page_number)
+        except PageNotAnInteger:
+            page_obj = paginator.page(1)
+        except EmptyPage:
+            page_obj = paginator.page(paginator.num_pages)
+
+        context = dict(
+            self.admin_site.each_context(request),
+            title=f"History: {obj}",
+            object=obj,
+            opts=self.model._meta,
+            audit_logs=page_obj.object_list,
+            page_obj=page_obj,
+            is_paginated=True,
+        )
+        return render(request, "audit/object_history.html", context)

auditrum/integrations/django/models.py

@@ -0,0 +1,31 @@
+from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.models import ContentType
+from django.db import models
+
+from auditrum.integrations.django.settings import audit_settings
+
+
+class AuditLog(models.Model):
+    operation = models.CharField(max_length=16)
+    changed_at = models.DateTimeField(auto_now_add=True)
+    content_type = models.ForeignKey(ContentType, null=True, blank=True, on_delete=models.SET_NULL)
+    object_id = models.CharField(max_length=255)
+    content_object = GenericForeignKey("content_type", "object_id")
+    table_name = models.CharField(max_length=255)
+    user_id = models.IntegerField(null=True, blank=True)
+    old_data = models.JSONField(null=True, blank=True)
+    new_data = models.JSONField(null=True, blank=True)
+    diff = models.JSONField(null=True, blank=True)
+    meta = models.JSONField(null=True, blank=True)
+    request_id = models.CharField(max_length=255, null=True, blank=True)
+    change_reason = models.TextField(null=True, blank=True)
+    source = models.CharField(max_length=255, null=True, blank=True)
+
+    class Meta:
+        db_table = audit_settings.table_name
+        verbose_name = "Audit Log"
+        verbose_name_plural = "Audit Logs"
+        managed = False
+
+    def __str__(self):
+        return f"{self.operation} on {self.table_name} ({self.object_id})"

auditrum/integrations/django/settings.py

@@ -0,0 +1,14 @@
+from django.conf import settings
+
+
+class AuditSettings:
+    @property
+    def table_name(self) -> str:
+        return getattr(settings, "PGAUDIT_TABLE_NAME", "auditlog")
+
+    @property
+    def enabled(self) -> bool:
+        return getattr(settings, "PGAUDIT_ENABLED", True)
+
+
+audit_settings = AuditSettings()