diff --git a/src/WebJobs.Extensions.CosmosDB/CosmosDBUtility.cs b/src/WebJobs.Extensions.CosmosDB/CosmosDBUtility.cs index fc36986c5..35c1a22eb 100644 --- a/src/WebJobs.Extensions.CosmosDB/CosmosDBUtility.cs +++ b/src/WebJobs.Extensions.CosmosDB/CosmosDBUtility.cs @@ -41,7 +41,7 @@ await CreateDatabaseAndCollectionIfNotExistAsync(context.Service, context.Resolv context.ResolvedAttribute.PartitionKey, context.ResolvedAttribute.ContainerThroughput); } - internal static async Task CreateDatabaseAndCollectionIfNotExistAsync(CosmosClient service, string databaseName, string containerName, string partitionKey, int throughput) + internal static async Task CreateDatabaseAndCollectionIfNotExistAsync(CosmosClient service, string databaseName, string containerName, string partitionKey, int throughput, bool setTTL = false) { await service.CreateDatabaseIfNotExistsAsync(databaseName); @@ -59,7 +59,23 @@ internal static async Task CreateDatabaseAndCollectionIfNotExistAsync(CosmosClie } catch (CosmosException cosmosException) when (cosmosException.StatusCode == System.Net.HttpStatusCode.NotFound) { - await database.CreateContainerAsync(containerName, partitionKey, desiredThroughput); + ContainerProperties containerProperties = new ContainerProperties() + { + Id = containerName + }; + + if (!string.IsNullOrEmpty(partitionKey)) + { + containerProperties.PartitionKeyPath = partitionKey; + } + + if (setTTL) + { + // Enabling TTL on the container without any defined time. TTL is set on the individual items. + containerProperties.DefaultTimeToLive = -1; + } + + await database.CreateContainerAsync(containerProperties, desiredThroughput); } } diff --git a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs index 0c04cc6ae..9493eb97a 100644 --- a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs +++ b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs @@ -140,13 +140,13 @@ private static async Task CreateLeaseCollectionIfNotExistsAsync(CosmosClient cos { try { - await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(cosmosClient, databaseName, collectionName, LeaseCollectionRequiredPartitionKey, throughput); + await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(cosmosClient, databaseName, collectionName, LeaseCollectionRequiredPartitionKey, throughput, setTTL: true); } catch (CosmosException cosmosException) when (cosmosException.StatusCode == System.Net.HttpStatusCode.BadRequest && cosmosException.Message.Contains("invalid for Gremlin API")) { - await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(cosmosClient, databaseName, collectionName, LeaseCollectionRequiredPartitionKeyFromGremlin, throughput); + await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(cosmosClient, databaseName, collectionName, LeaseCollectionRequiredPartitionKeyFromGremlin, throughput, setTTL: true); } } diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBTestUtility.cs b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBTestUtility.cs index 62af71b0a..67dee5f8b 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBTestUtility.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBTestUtility.cs @@ -22,7 +22,7 @@ internal static class CosmosDBTestUtility public const string DatabaseName = "ItemDB"; public const string ContainerName = "ItemCollection"; - public static Mock SetupCollectionMock(Mock mockService, Mock mockDatabase, string partitionKeyPath = null, int throughput = 0) + public static Mock SetupCollectionMock(Mock mockService, Mock mockDatabase, string partitionKeyPath, int throughput = 0, bool setTTL = false) { var mockContainer = new Mock(MockBehavior.Strict); @@ -46,8 +46,7 @@ public static Mock SetupCollectionMock(Mock mockService if (throughput == 0) { mockDatabase - .Setup(m => m.CreateContainerAsync(It.Is(i => i == ContainerName), - It.Is(p => p == partitionKeyPath), + .Setup(m => m.CreateContainerAsync(It.Is(cp => cp.Id == ContainerName && cp.PartitionKeyPath == partitionKeyPath && (!setTTL || cp.DefaultTimeToLive == -1)), It.Is(t => t == null), It.IsAny(), It.IsAny())) @@ -56,8 +55,7 @@ public static Mock SetupCollectionMock(Mock mockService else { mockDatabase - .Setup(m => m.CreateContainerAsync(It.Is(i => i == ContainerName), - It.Is(p => p == partitionKeyPath), + .Setup(m => m.CreateContainerAsync(It.Is(cp => cp.Id == ContainerName && cp.PartitionKeyPath == partitionKeyPath && (!setTTL || cp.DefaultTimeToLive == -1)), It.Is(t => t == throughput), It.IsAny(), It.IsAny())) diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBUtilityTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBUtilityTests.cs index 1760a1ab1..282e49220 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBUtilityTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBUtilityTests.cs @@ -21,7 +21,23 @@ public async Task CreateIfNotExists_DoesNotSetThroughput_IfZero() // Arrange var mockService = new Mock(MockBehavior.Strict); var context = CosmosDBTestUtility.CreateContext(mockService.Object, throughput: 0); - CosmosDBTestUtility.SetupCollectionMock(mockService, CosmosDBTestUtility.SetupDatabaseMock(mockService)); + CosmosDBTestUtility.SetupCollectionMock(mockService, CosmosDBTestUtility.SetupDatabaseMock(mockService), "/id"); + + // Act + await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(context); + + // Assert + mockService.VerifyAll(); + } + + [Fact] + public async Task CreateIfNotExists_SetsThroughput() + { + // Arrange + int throughput = 1000; + var mockService = new Mock(MockBehavior.Strict); + var context = CosmosDBTestUtility.CreateContext(mockService.Object, throughput: throughput); + CosmosDBTestUtility.SetupCollectionMock(mockService, CosmosDBTestUtility.SetupDatabaseMock(mockService), "/id", throughput: throughput); // Act await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(context); @@ -46,13 +62,29 @@ public async Task CreateIfNotExists_SetsPartitionKey_IfSpecified() mockService.VerifyAll(); } + [Fact] + public async Task CreateIfNotExists_SetsTTL_IfSpecified() + { + // Arrange + string partitionKeyPath = "partitionKey"; + var mockService = new Mock(MockBehavior.Strict); + + CosmosDBTestUtility.SetupCollectionMock(mockService, CosmosDBTestUtility.SetupDatabaseMock(mockService), partitionKeyPath, setTTL: true); + + // Act + await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(mockService.Object, CosmosDBTestUtility.DatabaseName, CosmosDBTestUtility.ContainerName, partitionKeyPath, throughput: 0, setTTL: true); + + // Assert + mockService.VerifyAll(); + } + [Fact] public async Task CreateIfNotExist_Succeeds() { // Arrange var mockService = new Mock(MockBehavior.Strict); CosmosDBContext context = CosmosDBTestUtility.CreateContext(mockService.Object); - CosmosDBTestUtility.SetupCollectionMock(mockService, CosmosDBTestUtility.SetupDatabaseMock(mockService)); + CosmosDBTestUtility.SetupCollectionMock(mockService, CosmosDBTestUtility.SetupDatabaseMock(mockService), null); // Act await CosmosDBUtility.CreateDatabaseAndCollectionIfNotExistAsync(context); diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs index e9aa1d063..3ba1e6fab 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs @@ -347,7 +347,7 @@ public async Task ValidCreateIfNotExists(ParameterInfo parameter) .ThrowsAsync(new CosmosException("not found", System.Net.HttpStatusCode.NotFound, 0, Guid.NewGuid().ToString(), 0)); mockDatabase - .Setup(m => m.CreateContainerAsync(It.IsAny(), It.Is(pk => pk == "/id"), It.IsAny(), It.IsAny(), It.IsAny())) + .Setup(m => m.CreateContainerAsync(It.Is(cp => cp.PartitionKeyPath == "/id" && cp.DefaultTimeToLive == -1), It.IsAny(), It.IsAny(), It.IsAny())) .ReturnsAsync(Mock.Of()); var factoryMock = new Mock(MockBehavior.Strict); @@ -362,7 +362,7 @@ public async Task ValidCreateIfNotExists(ParameterInfo parameter) CosmosDBTriggerAttribute cosmosDBTriggerAttribute = parameter.GetCustomAttribute(inherit: false); mockDatabase - .Verify(m => m.CreateContainerAsync(It.IsAny(), It.Is(pk => pk == "/id"), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once); + .Verify(m => m.CreateContainerAsync(It.Is(cp => cp.PartitionKeyPath == "/id" && cp.DefaultTimeToLive == -1), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once); } [Theory] @@ -401,11 +401,11 @@ public async Task ValidCreateIfNotExistsForGremlin(ParameterInfo parameter) .ThrowsAsync(new CosmosException("not found", System.Net.HttpStatusCode.NotFound, 0, Guid.NewGuid().ToString(), 0)); mockDatabase - .Setup(m => m.CreateContainerAsync(It.IsAny(), It.Is(pk => pk == "/id"), It.IsAny(), It.IsAny(), It.IsAny())) + .Setup(m => m.CreateContainerAsync(It.Is(cp => cp.PartitionKeyPath == "/id"), It.IsAny(), It.IsAny(), It.IsAny())) .ThrowsAsync(new CosmosException("invalid for Gremlin API", System.Net.HttpStatusCode.BadRequest, 0, Guid.NewGuid().ToString(), 0)); mockDatabase - .Setup(m => m.CreateContainerAsync(It.IsAny(), It.Is(pk => pk == "/partitionKey"), It.IsAny(), It.IsAny(), It.IsAny())) + .Setup(m => m.CreateContainerAsync(It.Is(cp => cp.PartitionKeyPath == "/partitionKey"), It.IsAny(), It.IsAny(), It.IsAny())) .ReturnsAsync(Mock.Of()); var factoryMock = new Mock(MockBehavior.Strict); @@ -420,10 +420,10 @@ public async Task ValidCreateIfNotExistsForGremlin(ParameterInfo parameter) CosmosDBTriggerAttribute cosmosDBTriggerAttribute = parameter.GetCustomAttribute(inherit: false); mockDatabase - .Verify(m => m.CreateContainerAsync(It.IsAny(), It.Is(pk => pk == "/id"), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once); + .Verify(m => m.CreateContainerAsync(It.Is(cp => cp.PartitionKeyPath == "/id"), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once); mockDatabase - .Verify(m => m.CreateContainerAsync(It.IsAny(), It.Is(pk => pk == "/partitionKey"), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once); + .Verify(m => m.CreateContainerAsync(It.Is(cp => cp.PartitionKeyPath == "/id"), It.IsAny(), It.IsAny(), It.IsAny()), Times.Once); } [Theory]