up
This commit is contained in:
StellaOps Bot
@@ -0,0 +1,210 @@
|
||||
using Dapper;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using StellaOps.Messaging.Abstractions;
|
||||
|
||||
namespace StellaOps.Messaging.Transport.Postgres;
|
||||
|
||||
/// <summary>
|
||||
/// PostgreSQL implementation of <see cref="IIdempotencyStore"/>.
|
||||
/// Uses INSERT ... ON CONFLICT DO NOTHING for atomic claiming.
|
||||
/// </summary>
|
||||
public sealed class PostgresIdempotencyStore : IIdempotencyStore
|
||||
{
|
||||
private readonly PostgresConnectionFactory _connectionFactory;
|
||||
private readonly string _name;
|
||||
private readonly ILogger<PostgresIdempotencyStore>? _logger;
|
||||
private readonly TimeProvider _timeProvider;
|
||||
private bool _tableInitialized;
|
||||
|
||||
public PostgresIdempotencyStore(
|
||||
PostgresConnectionFactory connectionFactory,
|
||||
string name,
|
||||
ILogger<PostgresIdempotencyStore>? logger = null,
|
||||
TimeProvider? timeProvider = null)
|
||||
{
|
||||
_connectionFactory = connectionFactory ?? throw new ArgumentNullException(nameof(connectionFactory));
|
||||
_name = name ?? throw new ArgumentNullException(nameof(name));
|
||||
_logger = logger;
|
||||
_timeProvider = timeProvider ?? TimeProvider.System;
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public string ProviderName => "postgres";
|
||||
|
||||
private string TableName => $"{_connectionFactory.Schema}.idempotency_{_name.ToLowerInvariant().Replace("-", "_")}";
|
||||
|
||||
/// <inheritdoc />
|
||||
public async ValueTask<IdempotencyResult> TryClaimAsync(
|
||||
string key,
|
||||
string value,
|
||||
TimeSpan window,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
ArgumentNullException.ThrowIfNull(key);
|
||||
ArgumentNullException.ThrowIfNull(value);
|
||||
|
||||
await EnsureTableExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
await using var conn = await _connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var now = _timeProvider.GetUtcNow();
|
||||
var expiresAt = now.Add(window);
|
||||
|
||||
// Clean up expired entries first
|
||||
var cleanupSql = $@"DELETE FROM {TableName} WHERE expires_at < @Now";
|
||||
await conn.ExecuteAsync(new CommandDefinition(cleanupSql, new { Now = now.UtcDateTime }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
|
||||
// Try to insert
|
||||
var sql = $@"
|
||||
INSERT INTO {TableName} (key, value, expires_at)
|
||||
VALUES (@Key, @Value, @ExpiresAt)
|
||||
ON CONFLICT (key) DO NOTHING
|
||||
RETURNING TRUE";
|
||||
|
||||
var result = await conn.ExecuteScalarAsync<bool?>(
|
||||
new CommandDefinition(sql, new { Key = key, Value = value, ExpiresAt = expiresAt.UtcDateTime }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
|
||||
if (result == true)
|
||||
{
|
||||
return IdempotencyResult.Claimed();
|
||||
}
|
||||
|
||||
// Key already exists, get existing value
|
||||
var existingSql = $@"SELECT value FROM {TableName} WHERE key = @Key";
|
||||
var existingValue = await conn.ExecuteScalarAsync<string?>(
|
||||
new CommandDefinition(existingSql, new { Key = key }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
|
||||
return IdempotencyResult.Duplicate(existingValue ?? string.Empty);
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public async ValueTask<bool> ExistsAsync(string key, CancellationToken cancellationToken = default)
|
||||
{
|
||||
ArgumentNullException.ThrowIfNull(key);
|
||||
|
||||
await EnsureTableExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
await using var conn = await _connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var now = _timeProvider.GetUtcNow();
|
||||
var sql = $@"SELECT EXISTS(SELECT 1 FROM {TableName} WHERE key = @Key AND expires_at > @Now)";
|
||||
|
||||
return await conn.ExecuteScalarAsync<bool>(
|
||||
new CommandDefinition(sql, new { Key = key, Now = now.UtcDateTime }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public async ValueTask<string?> GetAsync(string key, CancellationToken cancellationToken = default)
|
||||
{
|
||||
ArgumentNullException.ThrowIfNull(key);
|
||||
|
||||
await EnsureTableExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
await using var conn = await _connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var now = _timeProvider.GetUtcNow();
|
||||
var sql = $@"SELECT value FROM {TableName} WHERE key = @Key AND expires_at > @Now";
|
||||
|
||||
return await conn.ExecuteScalarAsync<string?>(
|
||||
new CommandDefinition(sql, new { Key = key, Now = now.UtcDateTime }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public async ValueTask<bool> ReleaseAsync(string key, CancellationToken cancellationToken = default)
|
||||
{
|
||||
ArgumentNullException.ThrowIfNull(key);
|
||||
|
||||
await EnsureTableExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
await using var conn = await _connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var sql = $@"DELETE FROM {TableName} WHERE key = @Key";
|
||||
var deleted = await conn.ExecuteAsync(
|
||||
new CommandDefinition(sql, new { Key = key }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
|
||||
return deleted > 0;
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public async ValueTask<bool> ExtendAsync(
|
||||
string key,
|
||||
TimeSpan extension,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
ArgumentNullException.ThrowIfNull(key);
|
||||
|
||||
await EnsureTableExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
await using var conn = await _connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var sql = $@"
|
||||
UPDATE {TableName}
|
||||
SET expires_at = expires_at + @Extension
|
||||
WHERE key = @Key";
|
||||
|
||||
var updated = await conn.ExecuteAsync(
|
||||
new CommandDefinition(sql, new { Key = key, Extension = extension }, cancellationToken: cancellationToken))
|
||||
.ConfigureAwait(false);
|
||||
|
||||
return updated > 0;
|
||||
}
|
||||
|
||||
private async ValueTask EnsureTableExistsAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
if (_tableInitialized) return;
|
||||
|
||||
await using var conn = await _connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var safeName = _name.ToLowerInvariant().Replace("-", "_");
|
||||
var sql = $@"
|
||||
CREATE TABLE IF NOT EXISTS {TableName} (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL,
|
||||
expires_at TIMESTAMPTZ NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_{safeName}_expires ON {TableName} (expires_at);";
|
||||
|
||||
await conn.ExecuteAsync(new CommandDefinition(sql, cancellationToken: cancellationToken)).ConfigureAwait(false);
|
||||
_tableInitialized = true;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Factory for creating PostgreSQL idempotency store instances.
|
||||
/// </summary>
|
||||
public sealed class PostgresIdempotencyStoreFactory : IIdempotencyStoreFactory
|
||||
{
|
||||
private readonly PostgresConnectionFactory _connectionFactory;
|
||||
private readonly ILoggerFactory? _loggerFactory;
|
||||
private readonly TimeProvider _timeProvider;
|
||||
|
||||
public PostgresIdempotencyStoreFactory(
|
||||
PostgresConnectionFactory connectionFactory,
|
||||
ILoggerFactory? loggerFactory = null,
|
||||
TimeProvider? timeProvider = null)
|
||||
{
|
||||
_connectionFactory = connectionFactory ?? throw new ArgumentNullException(nameof(connectionFactory));
|
||||
_loggerFactory = loggerFactory;
|
||||
_timeProvider = timeProvider ?? TimeProvider.System;
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public string ProviderName => "postgres";
|
||||
|
||||
/// <inheritdoc />
|
||||
public IIdempotencyStore Create(string name)
|
||||
{
|
||||
ArgumentNullException.ThrowIfNull(name);
|
||||
return new PostgresIdempotencyStore(
|
||||
_connectionFactory,
|
||||
name,
|
||||
_loggerFactory?.CreateLogger<PostgresIdempotencyStore>(),
|
||||
_timeProvider);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user