work-club-manager/backend/WorkClub.Infrastructure/Data/Interceptors/TenantDbConnectionInterceptor.cs

using System.Data;
using System.Data.Common;
using Microsoft.AspNetCore.Http;
using Microsoft.EntityFrameworkCore.Diagnostics;
using Microsoft.Extensions.Logging;
using Npgsql;

namespace WorkClub.Infrastructure.Data.Interceptors;

public class TenantDbConnectionInterceptor : DbCommandInterceptor
{
    private readonly IHttpContextAccessor _httpContextAccessor;
    private readonly ILogger<TenantDbConnectionInterceptor> _logger;

    public TenantDbConnectionInterceptor(
        IHttpContextAccessor httpContextAccessor,
        ILogger<TenantDbConnectionInterceptor> logger)
    {
        _httpContextAccessor = httpContextAccessor;
        _logger = logger;
    }

    public override ValueTask<InterceptionResult<DbDataReader>> ReaderExecutingAsync(
        DbCommand command,
        CommandEventData eventData,
        InterceptionResult<DbDataReader> result,
        CancellationToken cancellationToken = default)
    {
        SetTenantContext(command);
        return base.ReaderExecutingAsync(command, eventData, result, cancellationToken);
    }

    public override InterceptionResult<DbDataReader> ReaderExecuting(
        DbCommand command,
        CommandEventData eventData,
        InterceptionResult<DbDataReader> result)
    {
        SetTenantContext(command);
        return base.ReaderExecuting(command, eventData, result);
    }

    public override ValueTask<InterceptionResult<int>> NonQueryExecutingAsync(
        DbCommand command,
        CommandEventData eventData,
        InterceptionResult<int> result,
        CancellationToken cancellationToken = default)
    {
        SetTenantContext(command);
        return base.NonQueryExecutingAsync(command, eventData, result, cancellationToken);
    }

    public override InterceptionResult<int> NonQueryExecuting(
        DbCommand command,
        CommandEventData eventData,
        InterceptionResult<int> result)
    {
        SetTenantContext(command);
        return base.NonQueryExecuting(command, eventData, result);
    }

    private void SetTenantContext(DbCommand command)
    {
        var tenantId = _httpContextAccessor.HttpContext?.Items["TenantId"] as string;
        
        if (!string.IsNullOrWhiteSpace(tenantId) && command.Connection is NpgsqlConnection conn)
        {
            try
            {
                // If no transaction exists AND connection is open, start one
                // SET LOCAL only persists within a transaction, so we need explicit transaction scope
                if (command.Transaction == null && conn.State == ConnectionState.Open)
                {
                    command.Transaction = (NpgsqlTransaction)conn.BeginTransaction();
                    _logger.LogInformation("SetTenantContext: Started transaction for tenant context");
                }
                
                // Execute SET LOCAL as a separate command on the same connection
                // This ensures the session variable is properly set before the main query executes
                using var setLocalCmd = conn.CreateCommand();
                setLocalCmd.CommandText = $"SET LOCAL app.current_tenant_id = '{tenantId}'";
                
                // Use the same transaction if one exists (which it should now)
                if (command.Transaction is NpgsqlTransaction transaction)
                {
                    setLocalCmd.Transaction = transaction;
                }
                
                // Execute synchronously (safe in interceptor context)
                setLocalCmd.ExecuteNonQuery();
                _logger.LogInformation("SetTenantContext: Executed SET LOCAL for tenant {TenantId}", tenantId);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "SetTenantContext: Failed to execute SET LOCAL for tenant {TenantId}", tenantId);
                throw;
            }
        }
        else
        {
            _logger.LogWarning("SetTenantContext: No tenant context available (tenantId={TenantId})", tenantId ?? "null");
        }
    }
}