Add Throughput properties to Create and Replace Container #4

alexandru-slobodcicov · Feb 15, 2021 · dc78103 · dc78103
1 parent f4765c1
commit dc78103
Show file tree

Hide file tree

Showing 17 changed files with 524 additions and 164 deletions.
diff --git a/README.md b/README.md
@@ -50,19 +50,22 @@ All subsequent changes are applied to the created DB.
 
 * createContainer [REST](https://docs.microsoft.com/en-us/rest/api/cosmos-db/create-a-collection) [SDK](https://docs.microsoft.com/en-us/java/api/com.azure.cosmos.cosmosdatabase.createcontainer?view=azure-java-stable)
 <p> 
-Creates a Cosmos container while passing additional request options. 
+Creates a Cosmos container while passing additional request properties.
+There is a possibility to pass throughput properties either manual as a number or auto as a json.
 There is a flag to skip if exists and do not fail.
-If no options are specified then a ``/null`` partition key path is the default one.
+If no properties are specified then a ``/null`` partition key path is the default one.
 </p>
 
 * replaceContainer [REST](https://docs.microsoft.com/en-us/rest/api/cosmos-db/replace-a-collection) [SDK](https://docs.microsoft.com/en-us/java/api/com.azure.cosmos.cosmoscontainer.replace?view=azure-java-stable)
 <p>
-Replaces the container properties by container name.
+Replaces the container properties by container id.
+There is a possibility to pass throughput properties either manual as a number or auto as a json.
+If only properties specified no throughput properties will be amended and viceversa.
 </p>
 
 * deleteContainer [REST](https://docs.microsoft.com/en-us/rest/api/cosmos-db/delete-a-collection) [SDK](https://docs.microsoft.com/en-us/java/api/com.azure.cosmos.cosmoscontainer.delete?view=azure-java-stable)
 <p>
-Deletes the Cosmos container by name.
+Deletes the Cosmos container by id.
 There is a flag to skip if missing and do not fail.
 </p>
 

diff --git a/src/main/java/liquibase/ext/cosmosdb/change/CreateContainerChange.java b/src/main/java/liquibase/ext/cosmosdb/change/CreateContainerChange.java
@@ -23,16 +23,12 @@
 import liquibase.change.ChangeMetaData;
 import liquibase.change.DatabaseChange;
 import liquibase.database.Database;
-import liquibase.ext.cosmosdb.statement.CreateContainerIfNotExistsStatement;
 import liquibase.ext.cosmosdb.statement.CreateContainerStatement;
 import liquibase.statement.SqlStatement;
 import lombok.Getter;
 import lombok.NoArgsConstructor;
 import lombok.Setter;
 
-import static java.lang.Boolean.FALSE;
-import static java.util.Optional.ofNullable;
-
 @DatabaseChange(name = "createContainer",
         description = "Create container " +
                 "https://docs.microsoft.com/en-us/java/api/com.azure.cosmos.cosmosdatabase.createcontainer?view=azure-java-stable\n" +
@@ -45,6 +41,7 @@ public class CreateContainerChange extends AbstractCosmosChange {
 
     private String containerName;
     private String options;
+    private String throughput;
     private Boolean skipExisting;
 
     @Override
@@ -56,7 +53,7 @@ public String getConfirmationMessage() {
     public SqlStatement[] generateStatements(final Database database) {
 
         final CreateContainerStatement createContainerStatement =
-                ofNullable(skipExisting).orElse(FALSE) ? new CreateContainerIfNotExistsStatement(containerName, options) : new CreateContainerStatement(containerName, options);
+                        new CreateContainerStatement(containerName, options, throughput, skipExisting);
 
         return new SqlStatement[]{
                 createContainerStatement

diff --git a/src/main/java/liquibase/ext/cosmosdb/change/ReplaceContainerChange.java b/src/main/java/liquibase/ext/cosmosdb/change/ReplaceContainerChange.java
@@ -24,6 +24,7 @@
 import liquibase.change.DatabaseChange;
 import liquibase.database.Database;
 import liquibase.ext.cosmosdb.statement.CreateContainerStatement;
+import liquibase.ext.cosmosdb.statement.ReplaceContainerStatement;
 import liquibase.statement.SqlStatement;
 import lombok.Getter;
 import lombok.NoArgsConstructor;
@@ -42,6 +43,7 @@ public class ReplaceContainerChange extends AbstractCosmosChange {
 
     private String containerName;
     private String options;
+    private String throughput;
 
     @Override
     public String getConfirmationMessage() {
@@ -52,7 +54,7 @@ public String getConfirmationMessage() {
     public SqlStatement[] generateStatements(final Database database) {
 
         final CreateContainerStatement createContainerStatement
-                = new CreateContainerStatement(containerName, options);
+                = new ReplaceContainerStatement(containerName, options, throughput);
 
         return new SqlStatement[]{
                 createContainerStatement

diff --git a/src/main/java/liquibase/ext/cosmosdb/statement/AbstractNoSqlStatement.java b/src/main/java/liquibase/ext/cosmosdb/statement/AbstractNoSqlStatement.java
@@ -38,4 +38,9 @@ public boolean skipOnUnsupported() {
 
     public abstract String toJs();
 
+    @Override
+    public String toString() {
+        return toJs();
+    }
+
 }
diff --git a/src/main/java/liquibase/ext/cosmosdb/statement/CreateContainerIfNotExistsStatement.java b/src/main/java/liquibase/ext/cosmosdb/statement/CreateContainerIfNotExistsStatement.java
diff --git a/src/main/java/liquibase/ext/cosmosdb/statement/CreateContainerStatement.java b/src/main/java/liquibase/ext/cosmosdb/statement/CreateContainerStatement.java
@@ -22,11 +22,17 @@
 
 import com.azure.cosmos.CosmosDatabase;
 import com.azure.cosmos.models.CosmosContainerProperties;
+import com.azure.cosmos.models.ThroughputProperties;
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.NoArgsConstructor;
 
+import static java.lang.Boolean.FALSE;
+import static java.util.Optional.ofNullable;
+import static liquibase.ext.cosmosdb.statement.JsonUtils.toContainerProperties;
+import static liquibase.ext.cosmosdb.statement.JsonUtils.toThroughputProperties;
+
 @AllArgsConstructor
 @NoArgsConstructor
 @Getter
@@ -35,8 +41,18 @@ public class CreateContainerStatement extends AbstractNoSqlStatement implements
 
     public static final String COMMAND_NAME = "createContainer";
 
-    protected String containerName;
-    protected String options;
+    private String containerName;
+    private String options;
+    private String throughput;
+    private Boolean skipExisting;
+
+    public CreateContainerStatement(final String containerName, final String options, final String throughput) {
+        this(containerName, options, throughput, FALSE);
+    }
+
+    public CreateContainerStatement(final String containerName, final String options) {
+        this(containerName, options, null);
+    }
 
     public CreateContainerStatement(final String containerName) {
         this(containerName, null);
@@ -53,22 +69,25 @@ public String toJs() {
                 "db."
                         + getCommandName()
                         + "("
-                        + containerName
+                        + getContainerName()
+                        + ", "
+                        + getOptions()
                         + ", "
-                        + options
+                        + getThroughput()
+                        + ", "
+                        + getSkipExisting()
                         + ");";
     }
 
     @Override
     public void execute(final CosmosDatabase cosmosDatabase) {
-
-        final CosmosContainerProperties cosmosContainerProperties = JsonUtils.toContainerProperties(containerName, options);
-        cosmosDatabase.createContainer(cosmosContainerProperties);
+        final CosmosContainerProperties cosmosContainerProperties = toContainerProperties(getContainerName(), getOptions());
+        final ThroughputProperties throughputProperties = toThroughputProperties(getThroughput());
+        if (ofNullable(skipExisting).orElse(FALSE)) {
+            cosmosDatabase.createContainerIfNotExists(cosmosContainerProperties, throughputProperties);
+        } else {
+            cosmosDatabase.createContainer(cosmosContainerProperties, throughputProperties);
+        }
     }
 
-
-    @Override
-    public String toString() {
-        return toJs();
-    }
 }
diff --git a/src/main/java/liquibase/ext/cosmosdb/statement/JsonUtils.java b/src/main/java/liquibase/ext/cosmosdb/statement/JsonUtils.java
@@ -24,15 +24,20 @@
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.models.*;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.TreeNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import liquibase.util.StringUtil;
+import com.fasterxml.jackson.databind.node.ContainerNode;
+import com.fasterxml.jackson.databind.node.ValueNode;
 import lombok.NoArgsConstructor;
 import com.azure.cosmos.implementation.Document;
 
 import java.util.Map;
 
+import static com.azure.cosmos.implementation.Constants.Properties.AUTOPILOT_MAX_THROUGHPUT;
 import static java.util.Objects.nonNull;
 import static java.util.Optional.ofNullable;
+import static liquibase.util.StringUtil.isNotEmpty;
 import static liquibase.util.StringUtil.trimToNull;
 import static lombok.AccessLevel.PRIVATE;
 
@@ -97,7 +102,7 @@ public static Document mergeDocuments(final Document destination, final Document
     public static CosmosContainerProperties toContainerProperties(final String containerName, final String optionsJson) {
 
         final CosmosContainerProperties cosmosContainerProperties = new CosmosContainerProperties(containerName, DEFAULT_PARTITION_KEY_PATH);
-        if (StringUtil.isNotEmpty(StringUtil.trimToNull(optionsJson))) {
+        if (isNotEmpty(trimToNull(optionsJson))) {
             final DocumentCollection documentCollection = new DocumentCollection(optionsJson);
             if(nonNull(documentCollection.getPartitionKey())) {
                 cosmosContainerProperties.setPartitionKeyDefinition(documentCollection.getPartitionKey());
@@ -120,4 +125,23 @@ public static CosmosContainerProperties toContainerProperties(final String conta
         }
         return cosmosContainerProperties;
     }
+
+    public static ThroughputProperties toThroughputProperties(final String throughput) {
+
+        if (nonNull(trimToNull(throughput))) {
+            final TreeNode node;
+            try {
+                node = OBJECT_MAPPER.readTree(throughput);
+            } catch (final JsonProcessingException e) {
+                throw new IllegalArgumentException(String.format("Unable to parse JSON %s", throughput), e);
+            }
+            if(node.isValueNode()) {
+                return ThroughputProperties.createManualThroughput(((ValueNode)node).asInt());
+            }
+            if(node.isContainerNode() && ((ContainerNode<?>)node).has(AUTOPILOT_MAX_THROUGHPUT)) {
+                return ThroughputProperties.createAutoscaledThroughput(((ValueNode)node.get(AUTOPILOT_MAX_THROUGHPUT)).asInt());
+            }
+        }
+        return null;
+    }
 }
diff --git a/src/main/java/liquibase/ext/cosmosdb/statement/ReplaceContainerStatement.java b/src/main/java/liquibase/ext/cosmosdb/statement/ReplaceContainerStatement.java
@@ -22,19 +22,25 @@
 
 import com.azure.cosmos.CosmosDatabase;
 import com.azure.cosmos.models.CosmosContainerProperties;
+import com.azure.cosmos.models.ThroughputProperties;
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 
+import static java.util.Objects.nonNull;
+import static liquibase.ext.cosmosdb.statement.JsonUtils.toContainerProperties;
+import static liquibase.ext.cosmosdb.statement.JsonUtils.toThroughputProperties;
+import static liquibase.util.StringUtil.trimToNull;
+
 @AllArgsConstructor
 @Getter
 @EqualsAndHashCode(callSuper = true)
 public class ReplaceContainerStatement extends CreateContainerStatement {
 
     public static final String COMMAND_NAME = "replaceContainer";
 
-    public ReplaceContainerStatement(final String containerName, final String options) {
-        super(containerName, options);
+    public ReplaceContainerStatement(final String containerName, final String options, String throughput) {
+        super(containerName, options, throughput);
     }
 
     @Override
@@ -44,8 +50,14 @@ public String getCommandName() {
 
     @Override
     public void execute(final CosmosDatabase cosmosDatabase) {
-        final CosmosContainerProperties cosmosContainerProperties = JsonUtils.toContainerProperties(containerName, options);
-        cosmosDatabase.getContainer(containerName).replace(cosmosContainerProperties);
+        if(nonNull(trimToNull(getOptions()))) {
+            final CosmosContainerProperties cosmosContainerProperties = toContainerProperties(getContainerName(), getOptions());
+            cosmosDatabase.getContainer(getContainerName()).replace(cosmosContainerProperties);
+        }
+        if(nonNull(trimToNull(getThroughput()))) {
+            final ThroughputProperties throughputProperties = toThroughputProperties(getThroughput());
+            cosmosDatabase.getContainer(getContainerName()).replaceThroughput(throughputProperties);
+        }
     }
 
     @Override

diff --git a/src/main/resources/liquibase.parser.core.xml/dbchangelog-ext.xsd b/src/main/resources/liquibase.parser.core.xml/dbchangelog-ext.xsd
@@ -9,9 +9,10 @@
 
         <xsd:complexType>
 
-            <xsd:all>
+            <xsd:sequence>
                 <xsd:element name="options" type="xsd:string" minOccurs="0" maxOccurs="1"/>
-            </xsd:all>
+                <xsd:element name="throughput" type="xsd:string" minOccurs="0" maxOccurs="1"/>
+            </xsd:sequence>
 
             <xsd:attribute name="containerName" type="xsd:string" use="required"/>
             <xsd:attribute name="skipExisting" type="xsd:boolean" use="optional" default="false"/>
@@ -24,9 +25,10 @@
 
         <xsd:complexType>
 
-            <xsd:all>
-                <xsd:element name="options" type="xsd:string" minOccurs="0"/>
-            </xsd:all>
+            <xsd:sequence>
+                <xsd:element name="options" type="xsd:string" minOccurs="0" maxOccurs="1"/>
+                <xsd:element name="throughput" type="xsd:string" minOccurs="0" maxOccurs="1"/>
+            </xsd:sequence>
 
             <xsd:attribute name="containerName" type="xsd:string" use="required"/>
 

diff --git a/src/test/java/liquibase/ext/cosmosdb/CosmosLiquibaseIT.java b/src/test/java/liquibase/ext/cosmosdb/CosmosLiquibaseIT.java
@@ -23,6 +23,8 @@
 import com.azure.cosmos.CosmosContainer;
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.CosmosStoredProcedureProperties;
+import com.azure.cosmos.models.IndexingMode;
+import com.azure.cosmos.models.ThroughputProperties;
 import liquibase.Liquibase;
 import liquibase.ext.cosmosdb.changelog.CosmosRanChangeSet;
 import liquibase.resource.ClassLoaderResourceAccessor;
@@ -77,7 +79,33 @@ void testUpdateCreateContainer() {
         final CosmosContainerProperties maximal = containerProperties.stream().filter(c -> c.getId().equals("maximal")).findFirst().orElse(null);
         assertThat(maximal).isNotNull();
 
-        assertThat(containerProperties).hasSize(5);
+        assertThat(containerProperties).hasSize(7);
+    }
+
+    @SneakyThrows
+    @Test
+    void testUpdateReplaceContainer() {
+        final Liquibase liquibase = new Liquibase("liquibase/ext/changelog.replace-container.test.xml", new ClassLoaderResourceAccessor(), cosmosLiquibaseDatabase);
+        liquibase.update("");
+        assertThat(cosmosDatabase.getContainer(cosmosLiquibaseDatabase.getDatabaseChangeLogLockTableName()).read()).isNotNull();
+        assertThat(cosmosDatabase.getContainer(cosmosLiquibaseDatabase.getDatabaseChangeLogTableName()).read()).isNotNull();
+
+        assertThat(cosmosDatabase.getContainer("minimal").read()).isNotNull();
+        assertThat(cosmosDatabase.getContainer("minimal").read().getProperties()).isNotNull();
+        assertThat(cosmosDatabase.getContainer("minimal").readThroughput()).isNotNull();
+
+        final CosmosContainerProperties maximalProperties = cosmosDatabase.getContainer("maximal").read().getProperties();
+        assertThat(maximalProperties).isNotNull();
+        assertThat(maximalProperties.getId()).isEqualTo("maximal");
+        //TODO: Review after fixed replace
+        assertThat(maximalProperties.getIndexingPolicy().getIndexingMode()).isEqualTo(IndexingMode.CONSISTENT);
+        assertThat(maximalProperties.getIndexingPolicy().isAutomatic()).isTrue();
+        assertThat(maximalProperties.getIndexingPolicy().getIncludedPaths()).hasSize(1);
+        assertThat(maximalProperties.getIndexingPolicy().getExcludedPaths()).hasSize(1);
+
+        final ThroughputProperties maximalThroughput = cosmosDatabase.getContainer("maximal").readThroughput().getProperties();
+        assertThat(maximalThroughput).isNotNull();
+        assertThat(maximalThroughput.getAutoscaleMaxThroughput()).isEqualTo(8000);
     }
 
     @SneakyThrows