PostgresExtensions.Rows.cs

using System.Diagnostics.CodeAnalysis;
using System.Linq.Expressions;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Query;
using Microsoft.EntityFrameworkCore.Query.SqlExpressions;
using Npgsql.EntityFrameworkCore.PostgreSQL.Infrastructure.Internal;
using Npgsql.EntityFrameworkCore.PostgreSQL.Query.Internal;

namespace PostgresExtensions;

public enum RowLock
{
    Update,
    NoKeyUpdate,
    Share,
    KeyShare
}

public static partial class PostgresEntityFrameworkExtensions
{
    public static DbContextOptionsBuilder UseRowHints(this DbContextOptionsBuilder b)
    {
        b.ReplaceService<IQuerySqlGeneratorFactory, CustomQuerySqlGeneratorFactory>();
        return b;
    }

    /// <summary>
    /// <example>
    /// using var tr = db.Database.BeginTransaction();<br/>
    /// db.Cities.For(RowLock.Update).Where( ... ) // add "FOR xxxx" hint to the query<br/>
    /// await tr.CommitAsync();<br/>
    /// </example>
    /// </summary>
    public static IQueryable<T> For<T>(this DbSet<T> dataSet, RowLock lockMode) where T : class
    {
        var lockModeString = lockMode switch
                             {
                                 RowLock.Update      => "UPDATE",
                                 RowLock.NoKeyUpdate => "NO KEY UPDATE",
                                 RowLock.Share       => "SHARE",
                                 RowLock.KeyShare    => "KEY SHARE",
                                 _                   => throw new ArgumentOutOfRangeException(nameof(lockMode), lockMode, null)
                             };

        return dataSet.TagWith($"FOR:{lockModeString}");
    }
}

#region CustomQuerySqlGeneratorFactory

[SuppressMessage("Usage", "EF1001:Internal EF Core API usage.")]
public class CustomQuerySqlGeneratorFactory : IQuerySqlGeneratorFactory
{
    readonly QuerySqlGeneratorDependencies dependencies;
    readonly INpgsqlSingletonOptions       npgsqlSingletonOptions;

    public CustomQuerySqlGeneratorFactory(QuerySqlGeneratorDependencies dependencies, INpgsqlSingletonOptions npgsqlSingletonOptions)
    {
        this.dependencies           = dependencies;
        this.npgsqlSingletonOptions = npgsqlSingletonOptions;
    }

    public virtual QuerySqlGenerator Create() => new CustomQuerySqlGenerator(dependencies,
                                                                             npgsqlSingletonOptions.ReverseNullOrderingEnabled,
                                                                             npgsqlSingletonOptions.PostgresVersion);
}

#endregion

#region CustomQuerySqlGenerator

[SuppressMessage("Usage", "EF1001:Internal EF Core API usage.")]
public class CustomQuerySqlGenerator : NpgsqlQuerySqlGenerator
{
    public CustomQuerySqlGenerator(QuerySqlGeneratorDependencies dependencies, bool reverseNullOrderingEnabled, Version postgresVersion) :
        base(dependencies, reverseNullOrderingEnabled, postgresVersion)
    {
    }

    protected override Expression VisitSelect(SelectExpression selectExpression)
    {
        var r = base.VisitSelect(selectExpression);

        var lockTableTag = selectExpression.Tags?
                                           .Where(t => t.StartsWith("FOR:") && !t.Contains("--"))
                                           .Select(t => t.Substring(4))
                                           .FirstOrDefault();
        if (!string.IsNullOrWhiteSpace(lockTableTag))
            Sql.Append($" FOR {lockTableTag} ");

        return r;
    }
}

#endregion