From 52a3db2dd829dcdde6d650eb2bd78a7d7abc4251 Mon Sep 17 00:00:00 2001 From: Andriy Svyryd Date: Thu, 13 Jun 2024 21:47:57 -0700 Subject: [PATCH] Use ARM to create containers for Cosmos tests. --- .../Storage/Internal/CosmosDatabaseCreator.cs | 21 +- .../F1CosmosFixture.cs | 16 ++ .../ReloadTest.cs | 70 +++-- .../TestUtilities/CosmosTestStore.cs | 270 +++++++++++++----- 4 files changed, 262 insertions(+), 115 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseCreator.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseCreator.cs index 9ea755ce142..7b6382ce10e 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseCreator.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseCreator.cs @@ -212,24 +212,11 @@ public virtual Task CanConnectAsync(CancellationToken cancellationToken = => throw new NotSupportedException(CosmosStrings.CanConnectNotSupported); /// - /// Returns the store name of the property that is used to store the partition key. - /// - /// The entity type to get the partition key property name for. - /// The name of the partition key property. - [Obsolete("Use GetPartitionKeyStoreNames")] - private static string GetPartitionKeyStoreName(IEntityType entityType) - { - var name = entityType.GetPartitionKeyPropertyName(); - return name != null - ? entityType.FindProperty(name)!.GetJsonPropertyName() - : CosmosClientWrapper.DefaultPartitionKey; - } - - /// - /// Returns the store names of the properties that is used to store the partition keys. + /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to + /// the same compatibility standards as public APIs. It may be changed or removed without notice in + /// any release. You should only use it directly in your code with extreme caution and knowing that + /// doing so can result in application failures when updating to a new Entity Framework Core release. /// - /// The entity type to get the partition key property names for. - /// The names of the partition key property. private static IReadOnlyList GetPartitionKeyStoreNames(IEntityType entityType) { var properties = entityType.GetPartitionKeyProperties(); diff --git a/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs b/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs index 804a3eda4ac..9cf273f93bc 100644 --- a/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs +++ b/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs @@ -16,6 +16,22 @@ protected override ITestStoreFactory TestStoreFactory public override TestHelpers TestHelpers => CosmosTestHelpers.Instance; + public override async Task ReseedAsync() + { + await base.ReseedAsync(); + + using var context = CreateContext(); + try + { + await context.Teams.SingleAsync(t => t.Id == Team.Ferrari); + } + catch (Exception) + { + // Recreating the containers without using CosmosClient causes cached metadata in CosmosClient to be out of sync + // and causes the first query to fail. This is a workaround for that. + } + } + public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(w => w.Ignore(CosmosEventId.NoPartitionKeyDefined)); diff --git a/test/EFCore.Cosmos.FunctionalTests/ReloadTest.cs b/test/EFCore.Cosmos.FunctionalTests/ReloadTest.cs index 1a21d1de5be..33094e6724b 100644 --- a/test/EFCore.Cosmos.FunctionalTests/ReloadTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/ReloadTest.cs @@ -1,24 +1,34 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. -using Azure.Core; using Newtonsoft.Json.Linq; namespace Microsoft.EntityFrameworkCore; #nullable disable -public class ReloadTest +public class ReloadTest : IClassFixture { - public static IEnumerable IsAsyncData = new object[][] { [false], [true] }; + public static IEnumerable IsAsyncData = [[false], [true]]; + + private void AssertSql(params string[] expected) + => Fixture.TestSqlLoggerFactory.AssertBaseline(expected); + + protected void ClearLog() + => Fixture.TestSqlLoggerFactory.Clear(); + + protected CosmosReloadTestFixture Fixture { get; } + + public ReloadTest(CosmosReloadTestFixture fixture) + { + Fixture = fixture; + ClearLog(); + } [ConditionalFact] public async Task Entity_reference_can_be_reloaded() { - await using var testDatabase = await CosmosTestStore.CreateInitializedAsync("ReloadTest"); - - using var context = new ReloadTestContext(testDatabase); - await context.Database.EnsureCreatedAsync(); + using var context = CreateContext(); var entry = await context.AddAsync(new Item { Id = 1337 }); @@ -33,35 +43,33 @@ public async Task Entity_reference_can_be_reloaded() Assert.Null(itemJson["unmapped"]); } - public class ReloadTestContext(CosmosTestStore testStore) : DbContext + protected ReloadTestContext CreateContext() + => Fixture.CreateContext(); + + public class CosmosReloadTestFixture : SharedStoreFixtureBase { - private readonly string _connectionUri = testStore.ConnectionUri; - private readonly string _authToken = testStore.AuthToken; - private readonly string _name = testStore.Name; - private readonly TokenCredential _tokenCredential = testStore.TokenCredential; + protected override string StoreName + => nameof(ReloadTest); - protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder) - { - if (TestEnvironment.UseTokenCredential) - { - optionsBuilder.UseCosmos( - _connectionUri, - _tokenCredential, - _name, - b => b.ApplyConfiguration()); - } - else - { - optionsBuilder.UseCosmos( - _connectionUri, - _authToken, - _name, - b => b.ApplyConfiguration()); - } - } + protected override bool UsePooling + => false; + + protected override ITestStoreFactory TestStoreFactory + => CosmosTestStoreFactory.Instance; + public TestSqlLoggerFactory TestSqlLoggerFactory + => (TestSqlLoggerFactory)ServiceProvider.GetRequiredService(); + } + + public class ReloadTestContext(DbContextOptions dbContextOptions) : DbContext(dbContextOptions) + { protected override void OnModelCreating(ModelBuilder modelBuilder) { + modelBuilder.Entity( + b => + { + b.HasPartitionKey(e => e.Id); + }); } public DbSet Items { get; set; } diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs index 4b91646ddae..d49b0ed01f4 100644 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs @@ -3,6 +3,7 @@ using System.Net; using System.Net.Sockets; +using System.Threading; using Azure; using Azure.Core; using Azure.ResourceManager; @@ -34,7 +35,7 @@ public static async Task CreateInitializedAsync( Action? extensionConfiguration = null) { var testStore = Create(name, extensionConfiguration); - await testStore.InitializeAsync(null, (Func?)null); + await testStore.InitializeAsync(null, (Func?)null).ConfigureAwait(false); return testStore; } @@ -95,7 +96,7 @@ public static async ValueTask IsConnectionAvailableAsync() { if (_connectionAvailable == null) { - _connectionAvailable = await TryConnectAsync(); + _connectionAvailable = await TryConnectAsync().ConfigureAwait(false); } return _connectionAvailable.Value; @@ -106,7 +107,7 @@ private static async Task TryConnectAsync() CosmosTestStore? testStore = null; try { - testStore = await CreateInitializedAsync("NonExistent"); + testStore = await CreateInitializedAsync("NonExistent").ConfigureAwait(false); return true; } @@ -132,7 +133,7 @@ private static async Task TryConnectAsync() { if (testStore != null) { - await testStore.DisposeAsync(); + await testStore.DisposeAsync().ConfigureAwait(false); } } } @@ -158,20 +159,20 @@ protected override async Task InitializeAsync(Func createContext, Fun if (_dataFilePath == null) { - await base.InitializeAsync(createContext ?? (() => _storeContext), seed, clean); + await base.InitializeAsync(createContext ?? (() => _storeContext), seed, clean).ConfigureAwait(false); } else { using var context = createContext(); - await CreateFromFile(context); + await CreateFromFile(context).ConfigureAwait(false); } } private async Task CreateFromFile(DbContext context) { - if (await EnsureCreatedAsync(context)) + if (await EnsureCreatedAsync(context).ConfigureAwait(false)) { - await context.Database.EnsureCreatedAsync(); + await context.Database.EnsureCreatedAsync().ConfigureAwait(false); var cosmosClient = context.GetService(); var serializer = CosmosClientWrapper.Serializer; using var fs = new FileStream(_dataFilePath!, FileMode.Open, FileAccess.Read); @@ -208,7 +209,7 @@ private async Task CreateFromFile(DbContext context) document["Discriminator"] = entityName; await cosmosClient.CreateItemAsync( - "NorthwindContext", document, new FakeUpdateEntry()); + "NorthwindContext", document, new FakeUpdateEntry()).ConfigureAwait(false); } else if (reader.TokenType == JsonToken.EndObject) { @@ -234,104 +235,80 @@ public async Task EnsureCreatedAsync(DbContext context, CancellationToken if (!TestEnvironment.UseTokenCredential) { var cosmosClientWrapper = context.GetService(); - return await cosmosClientWrapper.CreateDatabaseIfNotExistsAsync(null, cancellationToken); + return await cosmosClientWrapper.CreateDatabaseIfNotExistsAsync(null, cancellationToken).ConfigureAwait(false); } - var databaseAccount = await GetDBAccount(cancellationToken); + var databaseAccount = await GetDBAccount(cancellationToken).ConfigureAwait(false); var collection = databaseAccount.Value.GetCosmosDBSqlDatabases(); var sqlDatabaseCreateUpdateOptions = new CosmosDBSqlDatabaseCreateOrUpdateContent(TestEnvironment.AzureLocation, new CosmosDBSqlDatabaseResourceInfo(Name)); + if (await collection.ExistsAsync(Name, cancellationToken)) + { + return false; + } + var databaseResponse = (await collection.CreateOrUpdateAsync( - WaitUntil.Completed, Name, sqlDatabaseCreateUpdateOptions, cancellationToken)).GetRawResponse(); - return databaseResponse.Status == (int)HttpStatusCode.Created; + WaitUntil.Completed, Name, sqlDatabaseCreateUpdateOptions, cancellationToken).ConfigureAwait(false)).GetRawResponse(); + return databaseResponse.Status == (int)HttpStatusCode.OK; } private async Task EnsureDeletedAsync(DbContext context, CancellationToken cancellationToken = default) { if (!TestEnvironment.UseTokenCredential) { - return await context.Database.EnsureDeletedAsync(cancellationToken); + return await context.Database.EnsureDeletedAsync(cancellationToken).ConfigureAwait(false); + } + + var databaseAccount = await GetDBAccount(cancellationToken).ConfigureAwait(false); + var collection = databaseAccount.Value.GetCosmosDBSqlDatabases(); + var database = (await collection.GetIfExistsAsync(Name, cancellationToken).ConfigureAwait(false)); + if (database == null + || !database.HasValue) + { + return false; } - var databaseAccount = await GetDBAccount(cancellationToken); - var databaseResponse = (await databaseAccount.Value.GetCosmosDBSqlDatabase(Name, cancellationToken).Value.DeleteAsync( - WaitUntil.Completed, cancellationToken)).GetRawResponse(); - return databaseResponse.Status == (int)HttpStatusCode.Created; + var databaseResponse = (await database.Value!.DeleteAsync(WaitUntil.Completed, cancellationToken).ConfigureAwait(false)).GetRawResponse(); + return databaseResponse.Status == (int)HttpStatusCode.OK; } - private async Task> GetDBAccount(CancellationToken cancellationToken) + private Task> GetDBAccount(CancellationToken cancellationToken = default) { var accountName = new Uri(ConnectionUri).Host.Split('.').First(); var databaseAccountIdentifier = CosmosDBAccountResource.CreateResourceIdentifier( TestEnvironment.SubscriptionId, TestEnvironment.ResourceGroup, accountName); - return await _armClient.GetCosmosDBAccountResource(databaseAccountIdentifier).GetAsync(cancellationToken); + return _armClient.GetCosmosDBAccountResource(databaseAccountIdentifier).GetAsync(cancellationToken); } public override async Task CleanAsync(DbContext context) { - var created = await EnsureCreatedAsync(context); + var created = await EnsureCreatedAsync(context).ConfigureAwait(false); try { if (!created) { - var cosmosClient = context.Database.GetCosmosClient(); - var database = cosmosClient.GetDatabase(Name); - var containerIterator = database.GetContainerQueryIterator(); - while (containerIterator.HasMoreResults) - { - foreach (var containerProperties in await containerIterator.ReadNextAsync()) - { - var container = database.GetContainer(containerProperties.Id); - var partitionKeys = containerProperties.PartitionKeyPaths.Select(p => p[1..]).ToList(); - var itemIterator = container.GetItemQueryIterator( - new QueryDefinition("SELECT * FROM c")); - - var items = new List<(string Id, PartitionKey PartitionKeyValue)>(); - while (itemIterator.HasMoreResults) - { - foreach (var item in await itemIterator.ReadNextAsync()) - { - var partitionKeyValue = PartitionKey.None; - if (partitionKeys.Count >= 1 - && item[partitionKeys[0]] is not null) - { - var builder = new PartitionKeyBuilder(); - foreach (var partitionKey in partitionKeys) - { - builder.Add((string?)item[partitionKey]); - } - - partitionKeyValue = builder.Build(); - } - - items.Add((item["id"]!.ToString(), partitionKeyValue)); - } - } - - foreach (var item in items) - { - await container.DeleteItemAsync(item.Id, item.PartitionKeyValue); - } - } - } + await DeleteContainers(context).ConfigureAwait(false); + } - created = await context.Database.EnsureCreatedAsync(); + if (!TestEnvironment.UseTokenCredential) + { + created = await context.Database.EnsureCreatedAsync().ConfigureAwait(false); if (!created) { - var creator = (CosmosDatabaseCreator)context.GetService(); - await creator.SeedAsync(); + await SeedAsync(context).ConfigureAwait(false); } } else { - await context.Database.EnsureCreatedAsync(); + await CreateContainersAsync(context).ConfigureAwait(false); + await SeedAsync(context).ConfigureAwait(false); } } catch (Exception) { try { - await EnsureDeletedAsync(context); + await EnsureDeletedAsync(context).ConfigureAwait(false); } catch (Exception) { @@ -341,6 +318,165 @@ public override async Task CleanAsync(DbContext context) } } + private async Task CreateContainersAsync(DbContext context) + { + var databaseAccount = await GetDBAccount().ConfigureAwait(false); + var collection = databaseAccount.Value.GetCosmosDBSqlDatabases(); + var database = await collection.GetAsync(Name).ConfigureAwait(false); + var model = context.GetService().Model; + + foreach (var container in GetContainersToCreate(model)) + { + var resource = new CosmosDBSqlContainerResourceInfo(container.Id) + { + AnalyticalStorageTtl = container.AnalyticalStoreTimeToLiveInSeconds, + DefaultTtl = container.DefaultTimeToLive, + PartitionKey = new CosmosDBContainerPartitionKey + { + Version = 2 + } + }; + + if (container.PartitionKeyStoreNames.Count > 1) + { + resource.PartitionKey.Kind = "MultiHash"; + } + foreach (var partitionKey in container.PartitionKeyStoreNames) + { + resource.PartitionKey.Paths.Add("/" + partitionKey); + } + + var content = new CosmosDBSqlContainerCreateOrUpdateContent(TestEnvironment.AzureLocation, resource); + if (container.Throughput != null) + { + content.Options.AutoscaleMaxThroughput = container.Throughput.AutoscaleMaxThroughput; + content.Options.Throughput = container.Throughput.Throughput; + } + + await database.Value.GetCosmosDBSqlContainers().CreateOrUpdateAsync( + WaitUntil.Completed, container.Id, content).ConfigureAwait(false); + } + } + + private static IEnumerable GetContainersToCreate(IModel model) + { + var containers = new Dictionary>(); + foreach (var entityType in model.GetEntityTypes().Where(et => et.FindPrimaryKey() != null)) + { + var container = entityType.GetContainer(); + if (container == null) + { + continue; + } + + if (!containers.TryGetValue(container, out var mappedTypes)) + { + mappedTypes = []; + containers[container] = mappedTypes; + } + + mappedTypes.Add(entityType); + } + + foreach (var (containerName, mappedTypes) in containers) + { + IReadOnlyList partitionKeyStoreNames = Array.Empty(); + int? analyticalTtl = null; + int? defaultTtl = null; + ThroughputProperties? throughput = null; + + foreach (var entityType in mappedTypes) + { + if (!partitionKeyStoreNames.Any()) + { + partitionKeyStoreNames = GetPartitionKeyStoreNames(entityType); + } + analyticalTtl ??= entityType.GetAnalyticalStoreTimeToLive(); + defaultTtl ??= entityType.GetDefaultTimeToLive(); + throughput ??= entityType.GetThroughput(); + } + + yield return new( + containerName, + partitionKeyStoreNames, + analyticalTtl, + defaultTtl, + throughput); + } + } + + private static IReadOnlyList GetPartitionKeyStoreNames(IEntityType entityType) + { + var properties = entityType.GetPartitionKeyProperties(); + return properties.Any() + ? properties.Select(p => p.GetJsonPropertyName()).ToList() + : [CosmosClientWrapper.DefaultPartitionKey]; + } + + private async Task DeleteContainers(DbContext context) + { + if (!TestEnvironment.UseTokenCredential) + { + var cosmosClient = context.Database.GetCosmosClient(); + var database = cosmosClient.GetDatabase(Name); + var containerIterator = database.GetContainerQueryIterator(); + while (containerIterator.HasMoreResults) + { + foreach (var containerProperties in await containerIterator.ReadNextAsync().ConfigureAwait(false)) + { + var container = database.GetContainer(containerProperties.Id); + var partitionKeys = containerProperties.PartitionKeyPaths.Select(p => p[1..]).ToList(); + var itemIterator = container.GetItemQueryIterator( + new QueryDefinition("SELECT * FROM c")); + + var items = new List<(string Id, PartitionKey PartitionKeyValue)>(); + while (itemIterator.HasMoreResults) + { + foreach (var item in await itemIterator.ReadNextAsync().ConfigureAwait(false)) + { + var partitionKeyValue = PartitionKey.None; + if (partitionKeys.Count >= 1 + && item[partitionKeys[0]] is not null) + { + var builder = new PartitionKeyBuilder(); + foreach (var partitionKey in partitionKeys) + { + builder.Add((string?)item[partitionKey]); + } + + partitionKeyValue = builder.Build(); + } + + items.Add((item["id"]!.ToString(), partitionKeyValue)); + } + } + + foreach (var item in items) + { + await container.DeleteItemAsync(item.Id, item.PartitionKeyValue).ConfigureAwait(false); + } + } + } + } + else + { + var databaseAccount = await GetDBAccount().ConfigureAwait(false); + var collection = databaseAccount.Value.GetCosmosDBSqlDatabases(); + var database = await collection.GetAsync(Name).ConfigureAwait(false); + var containers = await database.Value.GetCosmosDBSqlContainers().GetAllAsync().ToListAsync().ConfigureAwait(false); + foreach (var container in containers) + { + await container.DeleteAsync(WaitUntil.Completed).ConfigureAwait(false); + } + } + } + + private static async Task SeedAsync(DbContext context) + { + var creator = (CosmosDatabaseCreator)context.GetService(); + await creator.SeedAsync().ConfigureAwait(false); + } + public override void Dispose() => throw new InvalidOperationException("Calling Dispose can cause deadlocks. Use DisposeAsync instead."); @@ -359,7 +495,7 @@ public override async Task DisposeAsync() GetTestStoreIndex(ServiceProvider).RemoveShared(GetType().Name + Name); } - await EnsureDeletedAsync(_storeContext); + await EnsureDeletedAsync(_storeContext).ConfigureAwait(false); } _storeContext.Dispose();