Create tables on startup (#2309)

src/ApiService/ApiService/Program.cs

@@ -7,6 +7,7 @@ using System.Linq;
 global
 using Async = System.Threading.Tasks;
 using System.Text.Json;
+using ApiService.OneFuzzLib.Orm;
 using Azure.Core.Serialization;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Middleware;
@@ -109,12 +110,39 @@ public class Program {
             .AddSingleton<ILogSinks, LogSinks>()
             .AddHttpClient()
             .AddMemoryCache();
-        }
+        })
-        )
         .Build();
 
+        await SetupStorage(
+            host.Services.GetRequiredService<IStorage>(),
+            host.Services.GetRequiredService<IServiceConfig>());
+
         await host.RunAsync();
     }
 
+    public static async Async.Task SetupStorage(IStorage storage, IServiceConfig serviceConfig) {
+        // Creates the tables for each implementor of IOrm<T>
 
+        // locate all IOrm<T> instances and collect the Ts
+        var toCreate = new List<Type>();
+        var types = typeof(Program).Assembly.GetTypes();
+        foreach (var type in types) {
+            if (type.IsAbstract) {
+                continue;
+            }
+
+            foreach (var iface in type.GetInterfaces()) {
+                if (iface.IsGenericType && iface.GetGenericTypeDefinition() == typeof(IOrm<>)) {
+                    toCreate.Add(iface.GenericTypeArguments.Single());
+                    break;
+                }
+            }
+        }
+
+        var storageAccount = serviceConfig.OneFuzzFuncStorage;
+        if (storageAccount is not null) {
+            var tableClient = await storage.GetTableServiceClientForAccount(storageAccount);
+            await Async.Task.WhenAll(toCreate.Select(t => tableClient.CreateTableIfNotExistsAsync(serviceConfig.OneFuzzStoragePrefix + t.Name)));
+        }
+    }
 }

src/ApiService/ApiService/onefuzzlib/orm/Orm.cs

@@ -27,7 +27,7 @@ namespace ApiService.OneFuzzLib.Orm {
     }
 
 
-    public class Orm<T> : IOrm<T> where T : EntityBase {
+    public abstract class Orm<T> : IOrm<T> where T : EntityBase {
 #pragma warning disable CA1051 // permit visible instance fields
         protected readonly EntityConverter _entityConverter;
         protected readonly IOnefuzzContext _context;
@@ -109,7 +109,6 @@ namespace ApiService.OneFuzzLib.Orm {
 
             var account = accountId ?? _context.ServiceConfiguration.OneFuzzFuncStorage ?? throw new ArgumentNullException(nameof(accountId));
             var tableClient = await _context.Storage.GetTableServiceClientForAccount(account);
-            await tableClient.CreateTableIfNotExistsAsync(tableName);
             return tableClient.GetTableClient(tableName);
         }
 
@@ -146,7 +145,7 @@ namespace ApiService.OneFuzzLib.Orm {
     }
 
 
-    public class StatefulOrm<T, TState, TSelf> : Orm<T>, IStatefulOrm<T, TState> where T : StatefulEntityBase<TState> where TState : Enum {
+    public abstract class StatefulOrm<T, TState, TSelf> : Orm<T>, IStatefulOrm<T, TState> where T : StatefulEntityBase<TState> where TState : Enum {
         static Lazy<Func<object>>? _partitionKeyGetter;
         static Lazy<Func<object>>? _rowKeyGetter;
         static ConcurrentDictionary<string, Func<T, Async.Task<T>>?> _stateFuncs = new ConcurrentDictionary<string, Func<T, Async.Task<T>>?>();

src/ApiService/IntegrationTests/_FunctionTestBase.cs

@@ -1,6 +1,7 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Threading.Tasks;
 using ApiService.OneFuzzLib.Orm;
 using Azure.Data.Tables;
 using Azure.Storage.Blobs;
@@ -8,8 +9,11 @@ using IntegrationTests.Fakes;
 using Microsoft.Azure.Functions.Worker.Http;
 using Microsoft.OneFuzz.Service;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
+using Xunit;
 using Xunit.Abstractions;
 
+using Task = System.Threading.Tasks.Task;
+
 namespace IntegrationTests;
 
 // FunctionTestBase contains shared implementations for running
@@ -22,7 +26,7 @@ namespace IntegrationTests;
 // - one for Azure Storage (marked with [Trait("Category", "Live")])
 //
 // See AgentEventsTests for an example.
-public abstract class FunctionTestBase : IDisposable {
+public abstract class FunctionTestBase : IAsyncLifetime {
     private readonly IStorage _storage;
 
     // each test will use a different prefix for storage (tables, blobs) so they don't interfere
@@ -52,6 +56,18 @@ public abstract class FunctionTestBase : IDisposable {
         _blobClient = _storage.GetBlobServiceClientForAccount("").Result; // for test implementations this is always sync
     }
 
+    public async Task InitializeAsync() {
+        await Program.SetupStorage(Context.Storage, Context.ServiceConfiguration);
+    }
+
+    public async Task DisposeAsync() {
+        // clean up any tables & blobs that this test created
+        // these Get methods are always sync for test impls
+        await (
+            CleanupTables(_storage.GetTableServiceClientForAccount("").Result),
+            CleanupBlobs(_storage.GetBlobServiceClientForAccount("").Result));
+    }
+
     protected static string BodyAsString(HttpResponseData data) {
         data.Body.Seek(0, SeekOrigin.Begin);
         using var sr = new StreamReader(data.Body);
@@ -61,38 +77,32 @@ public abstract class FunctionTestBase : IDisposable {
     protected static T BodyAs<T>(HttpResponseData data)
         => EntityConverter.FromJsonString<T>(BodyAsString(data)) ?? throw new Exception($"unable to deserialize body as {typeof(T)}");
 
-    public void Dispose() {
+    private async Task CleanupBlobs(BlobServiceClient blobClient)
-        GC.SuppressFinalize(this);
+    => await Task.WhenAll(
+        await blobClient
+            .GetBlobContainersAsync(prefix: _storagePrefix)
+            .Where(c => c.IsDeleted != true)
+            .Select(async container => {
+                try {
+                    await blobClient.DeleteBlobContainerAsync(container.Name);
+                    Logger.Info($"cleaned up container {container.Name}");
+                } catch (Exception ex) {
+                    // swallow any exceptions: this is a best-effort attempt to cleanup
+                    Logger.Exception(ex, "error deleting container at end of test");
+                }
+            }).ToListAsync());
 
-        // clean up any tables & blobs that this test created
+    private async Task CleanupTables(TableServiceClient tableClient)
-        // these Get methods are always sync for test impls
+        => await Task.WhenAll(
-        CleanupTables(_storage.GetTableServiceClientForAccount("").Result);
+            await tableClient
-        CleanupBlobs(_storage.GetBlobServiceClientForAccount("").Result);
+                .QueryAsync(filter: Query.StartsWith("TableName", _storagePrefix))
-    }
+                .Select(async table => {
-
+                    try {
-    private void CleanupBlobs(BlobServiceClient blobClient) {
+                        await tableClient.DeleteTableAsync(table.Name);
-        var containersToDelete = blobClient.GetBlobContainers(prefix: _storagePrefix);
+                        Logger.Info($"cleaned up table {table.Name}");
-        foreach (var container in containersToDelete.Where(c => c.IsDeleted != true)) {
+                    } catch (Exception ex) {
-            try {
+                        // swallow any exceptions: this is a best-effort attempt to cleanup
-                blobClient.DeleteBlobContainer(container.Name);
+                        Logger.Exception(ex, "error deleting table at end of test");
-                Logger.Info($"cleaned up container {container.Name}");
+                    }
-            } catch (Exception ex) {
+                }).ToListAsync());
-                // swallow any exceptions: this is a best-effort attempt to cleanup
-                Logger.Exception(ex, "error deleting container at end of test");
-            }
-        }
-    }
-
-    private void CleanupTables(TableServiceClient tableClient) {
-        var tablesToDelete = tableClient.Query(filter: Query.StartsWith("TableName", _storagePrefix));
-        foreach (var table in tablesToDelete) {
-            try {
-                tableClient.DeleteTable(table.Name);
-                Logger.Info($"cleaned up table {table.Name}");
-            } catch (Exception ex) {
-                // swallow any exceptions: this is a best-effort attempt to cleanup
-                Logger.Exception(ex, "error deleting table at end of test");
-            }
-        }
-    }
 }