Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix lyrics analyzer rate limit #64

Merged
merged 6 commits into from
Jan 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/App/Extensions/CosmosDbServiceActivityExtensions.cs
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ public static class CosmosDbServiceActivityExtensions
/// <param name="activitySource">The <see cref="ActivitySource"/>.</param>
/// <param name="userId">The ID of the user.</param>
/// <returns>The started <see cref="Activity"/>.</returns>
public static Activity? StartDbGetLyricsAnalyzerUserRateLimitActivity(this ActivitySource activitySource, ulong userId) => StartDbGetLyricsAnalyzerUserRateLimitActivity(activitySource, userId, null);
public static Activity? StartDbGetLyricsAnalyzerUserRateLimitActivity(this ActivitySource activitySource, string userId) => StartDbGetLyricsAnalyzerUserRateLimitActivity(activitySource, userId, null);

/// <summary>
/// Starts an activity for getting a lyrics analyzer rate limit for a user from the database.
Expand All @@ -95,14 +95,14 @@ public static class CosmosDbServiceActivityExtensions
/// <param name="userId">The ID of the user.</param>
/// <param name="parentActivityId">The parent activity ID.</param>
/// <returns>The started <see cref="Activity"/>.</returns>
public static Activity? StartDbGetLyricsAnalyzerUserRateLimitActivity(this ActivitySource activitySource, ulong userId, string? parentActivityId)
public static Activity? StartDbGetLyricsAnalyzerUserRateLimitActivity(this ActivitySource activitySource, string userId, string? parentActivityId)
{
return activitySource.StartActivity(
name: "Database:GetLyricsAnalyzerUserRateLimitAsync",
kind: ActivityKind.Internal,
tags: new ActivityTagsCollection
{
{ "userId", userId.ToString() }
{ "userId", userId }
},
parentId: parentActivityId
);
Expand Down
25 changes: 25 additions & 0 deletions src/App/Logging/CosmosDb/CosmosDbServiceLogging.cs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,31 @@ namespace MuzakBot.App.Logging.CosmosDb;
/// </summary>
internal static partial class CosmosDbServiceLogging
{
/// <summary>
/// Logs the initialization of the database and ensures that it exists.
/// </summary>
/// <param name="logger">The <see cref="ILogger"/> instance.</param>
/// <param name="databaseName">The name of the database.</param>
[LoggerMessage(
EventName = "CosmosDbService.Initialize.EnsureDbExists",
Level = LogLevel.Information,
Message = "Ensuring that the database, '{databaseName}', exists."
)]
public static partial void LogInitializeEnsureDbExists(this ILogger logger, string databaseName);

/// <summary>
/// Logs the initialization of a container and ensures that it exists.
/// </summary>
/// <param name="logger">The <see cref="ILogger"/> instance.</param>
/// <param name="containerName">The name of the container.</param>
/// <param name="databaseName">The name of the database.</param>
[LoggerMessage(
EventName = "CosmosDbService.Initialize.EnsureContainerExists",
Level = LogLevel.Information,
Message = "Ensuring that the container, '{containerName}', exists in '{databaseName}'."
)]
public static partial void LogInitializeEnsureContainerExists(this ILogger logger, string containerName, string databaseName);

/// <summary>
/// Logs the start of an add or update operation to the database.
/// </summary>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,15 @@ public LyricsAnalyzerUserRateLimit()
/// Initializes a new instance of the <see cref="LyricsAnalyzerUserRateLimit"/> class.
/// </summary>
/// <param name="userId">The user's ID.</param>
public LyricsAnalyzerUserRateLimit(ulong userId): this(userId, 0, DateTimeOffset.UtcNow)
public LyricsAnalyzerUserRateLimit(string userId): this(userId, 0, DateTimeOffset.UtcNow)
{}

/// <summary>
/// Initializes a new instance of the <see cref="LyricsAnalyzerUserRateLimit"/> class.
/// </summary>
/// <param name="userId">The user's ID.</param>
/// <param name="currentRequestCount">The current request count.</param>
public LyricsAnalyzerUserRateLimit(ulong userId, int currentRequestCount): this(userId, currentRequestCount, DateTimeOffset.UtcNow)
public LyricsAnalyzerUserRateLimit(string userId, int currentRequestCount): this(userId, currentRequestCount, DateTimeOffset.UtcNow)
{}

/// <summary>
Expand All @@ -33,7 +33,7 @@ public LyricsAnalyzerUserRateLimit(ulong userId, int currentRequestCount): this(
/// <param name="userId">The user's ID.</param>
/// <param name="currentRequestCount">The current request count.</param>
/// <param name="lastRequestTime">The last request time.</param>
public LyricsAnalyzerUserRateLimit(ulong userId, int currentRequestCount, DateTimeOffset lastRequestTime)
public LyricsAnalyzerUserRateLimit(string userId, int currentRequestCount, DateTimeOffset lastRequestTime)
{
Id = Guid.NewGuid().ToString();
PartitionKey = "user-item";
Expand All @@ -47,7 +47,7 @@ public LyricsAnalyzerUserRateLimit(ulong userId, int currentRequestCount, DateTi
/// The user's ID.
/// </summary>
[JsonPropertyName("userId")]
public ulong UserId { get; set; }
public string UserId { get; set; } = null!;

/// <summary>
/// The current request count.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ public interface ILyricsAnalyzerUserRateLimit
{
string Id { get; set; }
string PartitionKey { get; set; }
ulong UserId { get; set; }
string UserId { get; set; }
int CurrentRequestCount { get; set; }
DateTimeOffset LastRequestTime { get; set; }

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ await FollowupAsync(
if (lyricsAnalyzerConfig.RateLimitEnabled)
{
_logger.LogInformation("Getting current rate limit for user '{UserId}' from database.", Context.User.Id);
lyricsAnalyzerUserRateLimit = await _cosmosDbService.GetLyricsAnalyzerUserRateLimitAsync(Context.User.Id, activity?.Id);
lyricsAnalyzerUserRateLimit = await _cosmosDbService.GetLyricsAnalyzerUserRateLimitAsync(Context.User.Id.ToString(), activity?.Id);

_logger.LogInformation("Current rate limit for user '{UserId}' is {CurrentRequestCount}/{MaxRequests}.", Context.User.Id, lyricsAnalyzerUserRateLimit.CurrentRequestCount, lyricsAnalyzerConfig.RateLimitMaxRequests);

Expand Down
4 changes: 4 additions & 0 deletions src/App/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -82,4 +82,8 @@

using var host = builder.Build();

// Initialize the database and containers for the Cosmos DB service
// before running the host.
await host.Services.GetRequiredService<ICosmosDbService>().InitializeDatabaseAsync();

await host.RunAsync();
39 changes: 39 additions & 0 deletions src/App/Services/CosmosDbService/CosmosDbService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using MuzakBot.App.Logging.CosmosDb;

namespace MuzakBot.App.Services;

Expand All @@ -28,6 +29,44 @@ public CosmosDbService(ILogger<CosmosDbService> logger, IOptions<CosmosDbService
_cosmosDbClient = new CosmosClient(_options.ConnectionString);
}

/// <summary>
/// Initializes the database and containers for MuzakBot.
/// </summary>
/// <returns></returns>
public async Task InitializeDatabaseAsync()
{
// Initialize the database and containers for 'lyrics-analyzer',
// if they don't exist.
_logger.LogInitializeEnsureDbExists("lyrics-analyzer");
Database lyricsAnalyzerDb = await _cosmosDbClient.CreateDatabaseIfNotExistsAsync(
id: "lyrics-analyzer"
);

_logger.LogInitializeEnsureContainerExists("song-lyrics", "lyrics-analyzer");
await lyricsAnalyzerDb.CreateContainerIfNotExistsAsync(
id: "song-lyrics",
partitionKeyPath: "/partitionKey"
);

_logger.LogInitializeEnsureContainerExists("command-configs", "lyrics-analyzer");
await lyricsAnalyzerDb.CreateContainerIfNotExistsAsync(
id: "command-configs",
partitionKeyPath: "/partitionKey"
);

_logger.LogInitializeEnsureContainerExists("prompt-styles", "lyrics-analyzer");
await lyricsAnalyzerDb.CreateContainerIfNotExistsAsync(
id: "prompt-styles",
partitionKeyPath: "/partitionKey"
);

_logger.LogInitializeEnsureContainerExists("rate-limit", "lyrics-analyzer");
await lyricsAnalyzerDb.CreateContainerIfNotExistsAsync(
id: "rate-limit",
partitionKeyPath: "/partitionKey"
);
}

/// <inheritdoc cref="IDisposable.Dispose"/>
public void Dispose()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,15 @@ public partial class CosmosDbService
/// </summary>
/// <param name="userId">The ID of the user.</param>
/// <returns>The retrieved lyrics analyzer rate limit for the user.</returns>
public async Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(ulong userId) => await GetLyricsAnalyzerUserRateLimitAsync(userId, null);
public async Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(string userId) => await GetLyricsAnalyzerUserRateLimitAsync(userId, null);

/// <summary>
/// Gets the lyrics analyzer rate limit for a user from the database.
/// </summary>
/// <param name="userId">The ID of the user.</param>
/// <param name="parentActivityId">The parent activity ID.</param>
/// <returns>The retrieved lyrics analyzer rate limit for the user.</returns>
public async Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(ulong userId, string? parentActivityId)
public async Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(string userId, string? parentActivityId)
{
using var activity = _activitySource.StartDbGetLyricsAnalyzerUserRateLimitActivity(
userId: userId,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,16 @@ namespace MuzakBot.App.Services;

public interface ICosmosDbService : IDisposable
{
Task InitializeDatabaseAsync();

Task<LyricsAnalyzerConfig> GetLyricsAnalyzerConfigAsync();
Task<LyricsAnalyzerConfig> GetLyricsAnalyzerConfigAsync(string? parentActivityId);

Task AddOrUpdateLyricsAnalyzerConfigAsync(LyricsAnalyzerConfig lyricsAnalyzerConfig);
Task AddOrUpdateLyricsAnalyzerConfigAsync(LyricsAnalyzerConfig lyricsAnalyzerConfig, string? parentActivityId);

Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(ulong userId);
Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(ulong userId, string? parentActivityId);
Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(string userId);
Task<LyricsAnalyzerUserRateLimit> GetLyricsAnalyzerUserRateLimitAsync(string userId, string? parentActivityId);

Task AddOrUpdateLyricsAnalyzerUserRateLimitAsync(LyricsAnalyzerUserRateLimit lyricsAnalyzerUserRateLimit);
Task AddOrUpdateLyricsAnalyzerUserRateLimitAsync(LyricsAnalyzerUserRateLimit lyricsAnalyzerUserRateLimit, string? parentActivityId);
Expand Down