Migrating notifications (#2188)

* Migrating notification

* add dotnet enpoint setting in the config

* format

* fix unit test

* format

* build fix

* fix notifictions function definition

* fix deserilization of requests
refactor secretdata
finish transfering Notifiction objects

* format

src/ApiService/ApiService/Functions/Notifications.cs

@@ -0,0 +1,84 @@
+using System.Net;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Azure.Functions.Worker.Http;
+
+namespace Microsoft.OneFuzz.Service.Functions;
+
+public class Notifications {
+    private readonly ILogTracer _log;
+    private readonly IEndpointAuthorization _auth;
+    private readonly IOnefuzzContext _context;
+
+    public Notifications(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) {
+        _log = log;
+        _auth = auth;
+        _context = context;
+    }
+
+    private async Async.Task<HttpResponseData> Get(HttpRequestData req) {
+        _log.Info("Notification search");
+        var request = await RequestHandling.ParseUri<NotificationSearch>(req);
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, "notification search");
+        }
+
+        var entries = request.OkV switch { { Container: null } => _context.NotificationOperations.SearchAll(), { Container: var c } => _context.NotificationOperations.SearchByRowKeys(c.Select(x => x.ContainerName))
+        };
+        var response = req.CreateResponse(HttpStatusCode.OK);
+        await response.WriteAsJsonAsync(entries);
+        return response;
+    }
+
+
+    private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
+        _log.Info("adding notification hook");
+        var request = await RequestHandling.ParseRequest<NotificationCreate>(req);
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, "notification create");
+        }
+
+        var notificationRequest = request.OkV;
+
+        var entry = await _context.NotificationOperations.Create(notificationRequest.Container, notificationRequest.Config,
+            notificationRequest.ReplaceExisting);
+
+        if (!entry.IsOk) {
+            return await _context.RequestHandling.NotOk(req, entry.ErrorV, context: "notification create");
+        }
+
+        var response = req.CreateResponse(HttpStatusCode.OK);
+        await response.WriteAsJsonAsync(entry.OkV);
+        return response;
+    }
+
+    private async Async.Task<HttpResponseData> Delete(HttpRequestData req) {
+        var request = await RequestHandling.ParseRequest<NotificationGet>(req);
+
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, context: "notification delete");
+        }
+        var entries = await _context.NotificationOperations.SearchByPartitionKeys(new[] { $"{request.OkV.NotificationId}" }).ToListAsync();
+
+        if (entries.Count == 0) {
+            return await _context.RequestHandling.NotOk(req, new Error(ErrorCode.INVALID_REQUEST, new[] { "unable to find notification" }), context: "notification delete");
+        }
+
+        if (entries.Count > 1) {
+            return await _context.RequestHandling.NotOk(req, new Error(ErrorCode.INVALID_REQUEST, new[] { "error identifying Notification" }), context: "notification delete");
+        }
+        var response = req.CreateResponse(HttpStatusCode.OK);
+        await response.WriteAsJsonAsync(entries[0]);
+        return response;
+    }
+
+
+    [Function("Notifications")]
+    public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req) {
+        return _auth.CallIfUser(req, r => r.Method switch {
+            "GET" => Get(r),
+            "POST" => Post(r),
+            "DELETE" => Delete(r),
+            _ => throw new InvalidOperationException("Unsupported HTTP method"),
+        });
+    }
+}

src/ApiService/ApiService/OneFuzzTypes/Enums.cs

@@ -308,3 +308,14 @@ public enum NodeDisposalStrategy {
     ScaleIn,
     Decomission
 }
+
+
+public enum GithubIssueState {
+    Open,
+    Closed
+}
+
+public enum GithubIssueSearchMatch {
+    Title,
+    Body
+}

src/ApiService/ApiService/OneFuzzTypes/Model.cs

@@ -1,4 +1,5 @@
-using System.Text.Json;
+using System.Reflection;
+using System.Text.Json;
 using System.Text.Json.Serialization;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 using Endpoint = System.String;
@@ -402,6 +403,9 @@ public record Scaleset(
 [JsonConverter(typeof(ContainerConverter))]
 public record Container(string ContainerName) {
     public string ContainerName { get; } = ContainerName.All(c => char.IsLetterOrDigit(c) || c == '-') ? ContainerName : throw new ArgumentException("Container name must have only numbers, letters or dashes");
+    public override string ToString() {
+        return ContainerName;
+    }
 }
 
 public class ContainerConverter : JsonConverter<Container> {
@@ -416,8 +420,8 @@ public class ContainerConverter : JsonConverter<Container> {
 }
 
 public record Notification(
-    Container Container,
-    Guid NotificationId,
+    [PartitionKey] Guid NotificationId,
+    [RowKey] Container Container,
     NotificationTemplate Config
 ) : EntityBase();
 
@@ -469,17 +473,98 @@ public record RegressionReport(
     CrashTestResult? OriginalCrashTestResult
 ) : IReport;
 
-public record NotificationTemplate(
-    AdoTemplate? AdoTemplate,
-    TeamsTemplate? TeamsTemplate,
-    GithubIssuesTemplate? GithubIssuesTemplate
+
+[JsonConverter(typeof(NotificationTemplateConverter))]
+#pragma warning disable CA1715
+public interface NotificationTemplate {
+#pragma warning restore CA1715
+}
+
+
+public class NotificationTemplateConverter : JsonConverter<NotificationTemplate> {
+    public override NotificationTemplate? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) {
+        using var templateJson = JsonDocument.ParseValue(ref reader);
+        try {
+            return templateJson.Deserialize<AdoTemplate>(options);
+        } catch (JsonException) {
+
+        }
+
+        try {
+            return templateJson.Deserialize<TeamsTemplate>(options);
+        } catch (JsonException) {
+        }
+
+        try {
+            return templateJson.Deserialize<GithubIssuesTemplate>(options);
+        } catch (JsonException) {
+        }
+        throw new JsonException("Unsupported notification template");
+    }
+
+    public override void Write(Utf8JsonWriter writer, NotificationTemplate value, JsonSerializerOptions options) {
+        if (value is AdoTemplate adoTemplate) {
+            JsonSerializer.Serialize(writer, adoTemplate, options);
+        } else if (value is TeamsTemplate teamsTemplate) {
+            JsonSerializer.Serialize(writer, teamsTemplate, options);
+        } else if (value is GithubIssuesTemplate githubIssuesTemplate) {
+            JsonSerializer.Serialize(writer, githubIssuesTemplate, options);
+        } else {
+            throw new JsonException("Unsupported notification template");
+        }
+
+    }
+}
+
+
+public record ADODuplicateTemplate(
+    List<string> Increment,
+    string? Comment,
+    Dictionary<string, string> SetState,
+    Dictionary<string, string> AdoFields
 );
 
-public record AdoTemplate();
+public record AdoTemplate(
+    Uri BaseUrl,
+    SecretData<string> AuthToken,
+    string Project,
+    string Type,
+    List<string> UniqueFields,
+    string? Comment,
+    Dictionary<string, string> AdoFields,
+    ADODuplicateTemplate OnDuplicate
+    ) : NotificationTemplate;
 
-public record TeamsTemplate();
+public record TeamsTemplate(SecretData<string> Url) : NotificationTemplate;
 
-public record GithubIssuesTemplate();
+
+public record GithubAuth(string User, string PersonalAccessToken);
+
+public record GithubIssueSearch(
+    string? Author,
+    GithubIssueState? State,
+    List<GithubIssueSearchMatch> FieldMatch,
+    [property: JsonPropertyName("string")] String str
+);
+
+public record GithubIssueDuplicate(
+    string? Comment,
+    List<string> Labels,
+    bool Reopen
+);
+
+
+public record GithubIssuesTemplate(
+    SecretData<GithubAuth> Auth,
+    string Organization,
+    string Repository,
+    string Title,
+    string Body,
+    GithubIssueSearch UniqueSearch,
+    List<string> Assignees,
+    List<string> Labels,
+    GithubIssueDuplicate OnDuplicate
+    ) : NotificationTemplate;
 
 public record Repro(
     [PartitionKey][RowKey] Guid VmId,
@@ -553,25 +638,57 @@ public record Vm(
     public string Name { get; } = Name.Length > 40 ? throw new ArgumentOutOfRangeException("VM name too long") : Name;
 };
 
+[JsonConverter(typeof(ISecretConverterFactory))]
+public interface ISecret<T> { }
 
-public record SecretAddress(Uri Url);
-
-
-/// This class allows us to store some data that are intended to be secret
-/// The secret field stores either the raw data or the address of that data
-/// This class allows us to maintain backward compatibility with existing
-/// NotificationTemplate classes
-public record SecretData<T>(T Secret) {
-    public override string ToString() {
-        if (Secret is SecretAddress) {
-            if (Secret is null) {
-                return string.Empty;
-            } else {
-                return Secret.ToString()!;
-            }
-        } else
-            return "[REDACTED]";
+public class ISecretConverterFactory : JsonConverterFactory {
+    public override bool CanConvert(Type typeToConvert) {
+        return typeToConvert.IsGenericType && typeToConvert.Name == typeof(ISecret<string>).Name;
     }
+
+    public override JsonConverter? CreateConverter(Type typeToConvert, JsonSerializerOptions options) {
+        var innerType = typeToConvert.GetGenericArguments().First();
+        return (JsonConverter)Activator.CreateInstance(
+            typeof(ISecretConverter<>).MakeGenericType(innerType),
+            BindingFlags.Instance | BindingFlags.Public,
+            binder: null,
+            args: Array.Empty<object?>(),
+            culture: null)!;
+    }
+}
+
+public class ISecretConverter<T> : JsonConverter<ISecret<T>> {
+    public override ISecret<T> Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) {
+
+        using var secretJson = JsonDocument.ParseValue(ref reader);
+
+        if (secretJson.RootElement.ValueKind == JsonValueKind.String) {
+            return (ISecret<T>)new SecretValue<string>(secretJson.RootElement.GetString()!);
+        }
+
+        if (secretJson.RootElement.TryGetProperty("url", out var secretUrl)) {
+            return new SecretAddress<T>(new Uri(secretUrl.GetString()!));
+        }
+
+        return new SecretValue<T>(secretJson.Deserialize<T>(options)!);
+    }
+
+    public override void Write(Utf8JsonWriter writer, ISecret<T> value, JsonSerializerOptions options) {
+        if (value is SecretAddress<T> secretAddress) {
+            JsonSerializer.Serialize(writer, secretAddress, options);
+        } else if (value is SecretValue<T> secretValue) {
+            JsonSerializer.Serialize(writer, secretValue.Value, options);
+        }
+    }
+}
+
+
+
+public record SecretValue<T>(T Value) : ISecret<T>;
+
+public record SecretAddress<T>(Uri Url) : ISecret<T>;
+
+public record SecretData<T>(ISecret<T> Secret) {
 }
 
 public record JobConfig(

src/ApiService/ApiService/OneFuzzTypes/Requests.cs

@@ -113,6 +113,20 @@ public record ContainerDelete(
     IDictionary<string, string>? Metadata = null
 ) : BaseRequest;
 
+public record NotificationCreate(
+    Container Container,
+    bool ReplaceExisting,
+    NotificationTemplate Config
+) : BaseRequest;
+
+public record NotificationSearch(
+    List<Container>? Container
+) : BaseRequest;
+
+public record NotificationGet(
+    Guid NotificationId
+) : BaseRequest;
+
 public record JobGet(
     Guid JobId
 );

src/ApiService/ApiService/OneFuzzTypes/Validated.cs

@@ -48,7 +48,7 @@ public abstract class ValidatedStringConverter<T> : JsonConverter<T> where T : V
 
 [JsonConverter(typeof(Converter))]
 public record PoolName : ValidatedString {
-    private PoolName(string value) : base(value) {
+    public PoolName(string value) : base(value) {
         // Debug.Assert(Check.IsAlnumDash(value));
     }
 

src/ApiService/ApiService/TestHooks.cs

@@ -79,7 +79,7 @@ public class TestHooks {
             select new KeyValuePair<string, string>(Uri.UnescapeDataString(cs.Substring(0, i)), Uri.UnescapeDataString(cs.Substring(i + 1)));
 
         var qs = new Dictionary<string, string>(q);
-        var d = await _secretOps.GetSecretStringValue(new SecretData<string>(qs["SecretName"]));
+        var d = await _secretOps.GetSecretStringValue(new SecretData<string>(new SecretValue<string>(qs["SecretName"])));
 
         var resp = req.CreateResponse(HttpStatusCode.OK);
         await resp.WriteAsJsonAsync(d);

src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs

@@ -8,6 +8,7 @@ public interface INotificationOperations : IOrm<Notification> {
     Async.Task NewFiles(Container container, string filename, bool failTaskOnTransientError);
     IAsyncEnumerable<Notification> GetNotifications(Container container);
     IAsyncEnumerable<(Task, IEnumerable<string>)> GetQueueTasks();
+    Async.Task<OneFuzzResult<Notification>> Create(Container container, NotificationTemplate config, bool replaceExisting);
 }
 
 public class NotificationOperations : Orm<Notification>, INotificationOperations {
@@ -33,20 +34,20 @@ public class NotificationOperations : Orm<Notification>, INotificationOperations
 
             done.Add(notification.Config);
 
-            if (notification.Config.TeamsTemplate != null) {
-                NotifyTeams(notification.Config.TeamsTemplate, container, filename, reportOrRegression!);
+            if (notification.Config is TeamsTemplate teamsTemplate) {
+                NotifyTeams(teamsTemplate, container, filename, reportOrRegression!);
             }
 
             if (reportOrRegression == null) {
                 continue;
             }
 
-            if (notification.Config.AdoTemplate != null) {
-                NotifyAdo(notification.Config.AdoTemplate, container, filename, reportOrRegression, failTaskOnTransientError);
+            if (notification.Config is AdoTemplate adoTemplate) {
+                NotifyAdo(adoTemplate, container, filename, reportOrRegression, failTaskOnTransientError);
             }
 
-            if (notification.Config.GithubIssuesTemplate != null) {
-                GithubIssue(notification.Config.GithubIssuesTemplate, container, filename, reportOrRegression);
+            if (notification.Config is GithubIssuesTemplate githubIssuesTemplate) {
+                GithubIssue(githubIssuesTemplate, container, filename, reportOrRegression);
             }
         }
 
@@ -86,6 +87,26 @@ public class NotificationOperations : Orm<Notification>, INotificationOperations
             .Where(taskTuple => taskTuple.Item2 != null)!;
     }
 
+    public async Async.Task<OneFuzzResult<Notification>> Create(Container container, NotificationTemplate config, bool replaceExisting) {
+        if (await _context.Containers.FindContainer(container, StorageType.Corpus) == null) {
+            return OneFuzzResult<Notification>.Error(ErrorCode.INVALID_REQUEST, errors: new[] { "invalid container" });
+        }
+
+        if (replaceExisting) {
+            var existing = this.SearchByRowKeys(new[] { container.ContainerName });
+            await foreach (var existingEntry in existing) {
+                _logTracer.Info($"replacing existing notification: {existingEntry.NotificationId} - {container}");
+                await this.Delete(existingEntry);
+            }
+        }
+
+        var entry = new Notification(Guid.NewGuid(), container, config);
+        await this.Insert(entry);
+        _logTracer.Info($"created notification.  notification_id:{entry.NotificationId} container:{entry.Container}");
+
+        return OneFuzzResult<Notification>.Ok(entry);
+    }
+
     public async Async.Task<Task?> GetRegressionReportTask(RegressionReport report) {
         if (report.CrashTestResult.CrashReport != null) {
             return await _context.TaskOperations.GetByJobIdAndTaskId(report.CrashTestResult.CrashReport.JobId, report.CrashTestResult.CrashReport.TaskId);

src/ApiService/ApiService/onefuzzlib/Request.cs

@@ -1,5 +1,9 @@
 using System.Net;
+using System.Text.Json;
+using System.Text.Json.Nodes;
+using Faithlife.Utility;
 using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
 namespace Microsoft.OneFuzz.Service;
 
@@ -29,7 +33,7 @@ public class RequestHandling : IRequestHandling {
     public static async Async.Task<OneFuzzResult<T>> ParseRequest<T>(HttpRequestData req) {
         Exception? exception = null;
         try {
-            var t = await req.ReadFromJsonAsync<T>();
+            var t = await JsonSerializer.DeserializeAsync<T>(req.Body, EntityConverter.GetJsonSerializerOptions());
             if (t != null) {
                 return OneFuzzResult<T>.Ok(t);
             }
@@ -46,6 +50,28 @@ public class RequestHandling : IRequestHandling {
         );
     }
 
+    public static async Async.Task<OneFuzzResult<T>> ParseUri<T>(HttpRequestData req) {
+        var query = System.Web.HttpUtility.ParseQueryString(req.Url.Query);
+        var doc = new JsonObject();
+        foreach (var key in query.AllKeys.WhereNotNull()) {
+            doc[key] = JsonValue.Create(query[key]);
+        }
+
+        try {
+            var result = doc.Deserialize<T>(EntityConverter.GetJsonSerializerOptions());
+            return result switch {
+                null => OneFuzzResult<T>.Error(
+                    ErrorCode.INVALID_REQUEST,
+                    $"Failed to deserialize message into type: {typeof(T)} - {await req.ReadAsStringAsync()}"
+                ),
+                var r => OneFuzzResult<T>.Ok(r),
+            };
+
+        } catch (JsonException exception) {
+            return OneFuzzResult<T>.Error(ConvertError(exception));
+        }
+    }
+
     public static Error ConvertError(Exception exception) {
         return new Error(
             ErrorCode.INVALID_REQUEST,

src/ApiService/ApiService/onefuzzlib/Secrets.cs

@@ -7,7 +7,8 @@ namespace Microsoft.OneFuzz.Service;
 
 public interface ISecretsOperations {
     public (Uri, string) ParseSecretUrl(Uri secretsUrl);
-    public Task<SecretData<SecretAddress>?> SaveToKeyvault<T>(SecretData<T> secretData);
+    public Task<SecretAddress<T>> SaveToKeyvault<T>(SecretData<T> secretData);
+
     public Task<string?> GetSecretStringValue<T>(SecretData<T> data);
 
     public Task<KeyVaultSecret> StoreInKeyvault(Uri keyvaultUrl, string secretName, string secretValue);
@@ -34,33 +35,30 @@ public class SecretsOperations : ISecretsOperations {
         return (new Uri(vaultUrl), secretName);
     }
 
-    public async Task<SecretData<SecretAddress>?> SaveToKeyvault<T>(SecretData<T> secretData) {
-        if (secretData == null || secretData.Secret is null)
-            return null;
+    public async Task<SecretAddress<T>> SaveToKeyvault<T>(SecretData<T> secretData) {
 
-        if (secretData.Secret is SecretAddress) {
-            return secretData as SecretData<SecretAddress>;
-        } else {
+        if (secretData.Secret is SecretAddress<T> secretAddress) {
+            return secretAddress;
+        } else if (secretData.Secret is SecretValue<T> sValue) {
             var secretName = Guid.NewGuid();
             string secretValue;
-            if (secretData.Secret is string) {
-                secretValue = (secretData.Secret as string)!.Trim();
+            if (sValue.Value is string secretString) {
+                secretValue = secretString.Trim();
             } else {
-                secretValue = JsonSerializer.Serialize(secretData.Secret, EntityConverter.GetJsonSerializerOptions());
+                secretValue = JsonSerializer.Serialize(sValue.Value, EntityConverter.GetJsonSerializerOptions());
             }
 
             var kv = await StoreInKeyvault(GetKeyvaultAddress(), secretName.ToString(), secretValue);
-            return new SecretData<SecretAddress>(new SecretAddress(kv.Id));
+            return new SecretAddress<T>(kv.Id);
         }
+
+        throw new Exception("Invalid secret value");
     }
 
     public async Task<string?> GetSecretStringValue<T>(SecretData<T> data) {
-        if (data.Secret is null) {
-            return null;
-        }
 
-        if (data.Secret is SecretAddress) {
-            var secret = await GetSecret((data.Secret as SecretAddress)!.Url);
+        if (data.Secret is SecretAddress<T> secretAddress) {
+            var secret = await GetSecret(secretAddress.Url);
             return secret.Value;
         } else {
             return data.Secret.ToString();
@@ -101,11 +99,8 @@ public class SecretsOperations : ISecretsOperations {
     }
 
     public async Task<DeleteSecretOperation?> DeleteRemoteSecretData<T>(SecretData<T> data) {
-        if (data.Secret is SecretAddress) {
-            if (data.Secret is not null)
-                return await DeleteSecret((data.Secret as SecretAddress)!.Url);
-            else
-                return null;
+        if (data.Secret is SecretAddress<T> secretAddress) {
+            return await DeleteSecret(secretAddress.Url);
         } else {
             return null;
         }

src/ApiService/ApiService/onefuzzlib/TaskOperations.cs

@@ -110,7 +110,7 @@ public class TaskOperations : StatefulOrm<Task, TaskState, TaskOperations>, ITas
     }
 
     private async Async.Task MarkDependantsFailed(Task task, List<Task>? taskInJob = null) {
-        taskInJob ??= await SearchByPartitionKey(task.JobId.ToString()).ToListAsync();
+        taskInJob ??= await SearchByPartitionKeys(new[] { task.JobId.ToString() }).ToListAsync();
 
         foreach (var t in taskInJob) {
             if (t.Config.PrereqTasks != null) {

src/ApiService/ApiService/onefuzzlib/orm/CustomConverterFactory.cs

@@ -165,10 +165,10 @@ public sealed class PolymorphicConverterFactory : JsonConverterFactory {
     }
 
     public override JsonConverter? CreateConverter(Type typeToConvert, JsonSerializerOptions options) {
-        var (field, attribute) = typeToConvert.GetProperties()
-                .Select(p => (p.Name, p.GetCustomAttribute<TypeDiscrimnatorAttribute>()))
-                .Where(p => p.Item2 != null)
-                .First();
+        var (field, attribute) = typeToConvert
+            .GetProperties()
+            .Select(p => (p.Name, p.GetCustomAttribute<TypeDiscrimnatorAttribute>()))
+            .First(p => p.Item2 != null);
 
 
         return (JsonConverter)Activator.CreateInstance(

src/ApiService/ApiService/onefuzzlib/orm/EntityConverter.cs

@@ -232,9 +232,8 @@ public class EntityConverter {
                 return Guid.Parse(entity.GetString(ef.kind.ToString()));
             else if (ef.type == typeof(int))
                 return int.Parse(entity.GetString(ef.kind.ToString()));
-            else if (ef.type == typeof(PoolName))
-                // TODO: this should be able to be generic over any ValidatedString
-                return PoolName.Parse(entity.GetString(ef.kind.ToString()));
+            else if (ef.type.IsClass)
+                return ef.type.GetConstructor(new[] { typeof(string) })!.Invoke(new[] { entity.GetString(ef.kind.ToString()) });
             else {
                 throw new Exception($"invalid partition or row key type of {info.type} property {name}: {ef.type}");
             }

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

@@ -17,8 +17,8 @@ namespace ApiService.OneFuzzLib.Orm {
         Task<ResultVoid<(int, string)>> Delete(T entity);
 
         IAsyncEnumerable<T> SearchAll();
-        IAsyncEnumerable<T> SearchByPartitionKey(string partitionKey);
-        IAsyncEnumerable<T> SearchByRowKey(string rowKey);
+        IAsyncEnumerable<T> SearchByPartitionKeys(IEnumerable<string> partitionKeys);
+        IAsyncEnumerable<T> SearchByRowKeys(IEnumerable<string> rowKeys);
         IAsyncEnumerable<T> SearchByTimeRange(DateTimeOffset min, DateTimeOffset max);
 
         // Allow using tuple to search.
@@ -123,11 +123,11 @@ namespace ApiService.OneFuzzLib.Orm {
         public IAsyncEnumerable<T> SearchAll()
             => QueryAsync(null);
 
-        public IAsyncEnumerable<T> SearchByPartitionKey(string partitionKey)
-            => QueryAsync(Query.PartitionKey(partitionKey));
+        public IAsyncEnumerable<T> SearchByPartitionKeys(IEnumerable<string> partitionKeys)
+            => QueryAsync(Query.PartitionKeys(partitionKeys));
 
-        public IAsyncEnumerable<T> SearchByRowKey(string rowKey)
-            => QueryAsync(Query.RowKey(rowKey));
+        public IAsyncEnumerable<T> SearchByRowKeys(IEnumerable<string> rowKeys)
+            => QueryAsync(Query.RowKeys(rowKeys));
 
         public IAsyncEnumerable<T> SearchByTimeRange(DateTimeOffset min, DateTimeOffset max) {
             return QueryAsync(Query.TimeRange(min, max));

src/ApiService/Tests/OrmModelsTest.cs

@@ -41,6 +41,18 @@ namespace Tests {
             );
         }
 
+        public static Gen<ISecret<T>> ISecret<T>() {
+            if (typeof(T) == typeof(string)) {
+                return Arb.Generate<string>().Select(s => (ISecret<T>)new SecretValue<string>(s));
+            }
+
+            if (typeof(T) == typeof(GithubAuth)) {
+                return Arb.Generate<GithubAuth>().Select(s => (ISecret<T>)new SecretValue<GithubAuth>(s));
+            } else {
+                throw new Exception($"Unsupported secret type {typeof(T)}");
+            }
+        }
+
         public static Gen<Version> Version() {
             //OneFuzz version uses 3 number version
             return Arb.Generate<Tuple<UInt16, UInt16, UInt16>>().Select(
@@ -270,7 +282,7 @@ namespace Tests {
                         Id: arg.Item4,
                         EventTime: arg.Item5
                     )
-            ); ;
+            );
         }
 
         public static Gen<Report> Report() {
@@ -341,6 +353,13 @@ namespace Tests {
             );
         }
 
+        public static Gen<NotificationTemplate> NotificationTemplate() {
+            return Gen.OneOf(new[] {
+                Arb.Generate<AdoTemplate>().Select(e => e as NotificationTemplate),
+                Arb.Generate<TeamsTemplate>().Select(e => e as NotificationTemplate),
+                Arb.Generate<GithubIssuesTemplate>().Select(e => e as NotificationTemplate)
+            });
+        }
 
         public static Gen<Notification> Notification() {
             return Arb.Generate<Tuple<Container, Guid, NotificationTemplate>>().Select(
@@ -443,6 +462,11 @@ namespace Tests {
             return Arb.From(OrmGenerators.Container());
         }
 
+
+        public static Arbitrary<NotificationTemplate> NotificationTemplate() {
+            return Arb.From(OrmGenerators.NotificationTemplate());
+        }
+
         public static Arbitrary<Notification> Notification() {
             return Arb.From(OrmGenerators.Notification());
         }
@@ -454,6 +478,12 @@ namespace Tests {
         public static Arbitrary<Job> Job() {
             return Arb.From(OrmGenerators.Job());
         }
+
+        public static Arbitrary<ISecret<T>> ISecret<T>() {
+            return Arb.From(OrmGenerators.ISecret<T>());
+        }
+
+
     }