public class MyDbContext : IdentityDbContext<...>
{
    public DbSet<EntityA> EntitiesA { get; set; }
    public DbSet<EntityB> EntitiesB { get; set; }
}

public static class EntitiesA
{
    private static Func<MyDbContext, ...> _queryEntitiesA = EF.CompileAsyncQuery((MyDbContext context, ...) => context...);

    public static async Task<EntityA> QueryEntitiesA(this DbSet<EntityA> dbSet, ...)
    {
        var context = dbSet.GetContext();
        return await _queryEntitiesA(...)
    }
}

public static class Common
{
    private static Dictionary<Type, FieldInfo> _fields = new Dictionary<Type, FieldInfo>();
    public static MyDbContext GetContext<TEntity>(this DbSet<TEntity> dbSet) where TEntity : class
    {
        var type = typeof(TEntity);
        if (!Common._fields.ContainsKey(type))
        {
            Common._fields.Add(type, typeof(InternalDbSet<TEntity>).GetField("_context", BindingFlags.NonPublic | BindingFlags.Instance));
        }
        return Common._fields[type].GetValue(dbSet) as MyDbContext;
    }
}