Fix runs paging regression (#158)

microsoft · Nov 26, 2024 · cb56a32 · cb56a32
1 parent fc27d78
commit cb56a32
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
diff --git a/server/ControlPlane/Database/Migrations/4.cs b/server/ControlPlane/Database/Migrations/4.cs
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace Tyger.ControlPlane.Database.Migrations;
+
+public class Migrator4 : Migrator
+{
+    public override async Task Apply(Npgsql.NpgsqlDataSource dataSource, ILogger logger, CancellationToken cancellationToken)
+    {
+        using (var cmd = dataSource.CreateCommand(
+            """
+            CREATE INDEX IF NOT EXISTS idx_runs_created_at_id_status
+            ON runs (created_at, id)
+            INCLUDE (status)
+            """))
+        {
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+        }
+
+        using (var cmd = dataSource.CreateCommand(
+            """
+            DROP INDEX IF EXISTS idx_runs_created_at
+            """))
+        {
+            await cmd.ExecuteNonQueryAsync(cancellationToken);
+        }
+    }
+}
diff --git a/server/ControlPlane/Database/Migrations/DatabaseVersions.cs b/server/ControlPlane/Database/Migrations/DatabaseVersions.cs
@@ -27,6 +27,10 @@ public enum DatabaseVersion
     [Description("Making run management more scalable")]
     [MinimumSupportedVersion]
     RunScalability = 3,
+
+    [Migrator(typeof(Migrator4))]
+    [Description("Adjusting run indexes")]
+    AddRunsIndexForPaging = 4,
 }
 
 public sealed class DatabaseVersions : BackgroundService, IHealthCheck

diff --git a/server/ControlPlane/Database/Repository.cs b/server/ControlPlane/Database/Repository.cs
@@ -785,13 +785,14 @@ FROM runs
                 try
                 {
                     var fields = JsonSerializer.Deserialize<long[]>(Encoding.ASCII.GetString(Base32.ZBase32.Decode(continuationToken)), _serializerOptions);
-                    if (fields is { Length: 1 or 2 })
+                    if (fields is { Length: 2 })
                     {
-                        var id = fields[^1];
-                        sb.AppendLine($"{(hasPredicate ? "AND" : "WHERE")} id < ${++paramNumber}");
+                        var createdAt = new DateTimeOffset(fields[0], TimeSpan.Zero);
+                        var id = fields[1];
+                        sb.AppendLine($"{(hasPredicate ? "AND" : "WHERE")} (created_at, id) < (${++paramNumber}, ${++paramNumber})");
+                        parameters.Add(new() { Value = createdAt, NpgsqlDbType = NpgsqlDbType.TimestampTz });
                         parameters.Add(new() { Value = id, NpgsqlDbType = NpgsqlDbType.Bigint });
                         valid = true;
-                        hasPredicate = true;
                     }
                 }
                 catch (Exception e) when (e is JsonException or FormatException)
@@ -838,7 +839,7 @@ FROM runs
             {
                 results.RemoveAt(limit);
                 var (run, final) = results[^1];
-                string newToken = Base32.ZBase32.Encode(Encoding.ASCII.GetBytes(JsonSerializer.Serialize(new[] { run.Id!.Value }, _serializerOptions)));
+                string newToken = Base32.ZBase32.Encode(Encoding.ASCII.GetBytes(JsonSerializer.Serialize(new[] { run.CreatedAt!.Value.UtcTicks, run.Id!.Value }, _serializerOptions)));
                 return (results, newToken);
             }