diff --git a/state/azure/cosmosdb/cosmosdb.go b/state/azure/cosmosdb/cosmosdb.go
index 588249609e..c1a412c5d7 100644
--- a/state/azure/cosmosdb/cosmosdb.go
+++ b/state/azure/cosmosdb/cosmosdb.go
@@ -114,11 +114,9 @@ func (c *StateStore) Init(meta state.Metadata) error {
 	// Create the client; first, try authenticating with a master key, if present
 	var config *documentdb.Config
 	if m.MasterKey != "" {
-		config = &documentdb.Config{
-			MasterKey: &documentdb.Key{
-				Key: m.MasterKey,
-			},
-		}
+		config = documentdb.NewConfig(&documentdb.Key{
+			Key: m.MasterKey,
+		})
 	} else {
 		// Fallback to using Azure AD
 		env, errB := azure.NewEnvironmentSettings("cosmosdb", meta.Properties)
@@ -273,7 +271,7 @@ func (c *StateStore) Set(req *state.SetRequest) error {
 	if err != nil {
 		return err
 	}
-	_, err = c.client.UpsertDocument(c.collection.Self, doc, options...)
+	_, err = c.client.UpsertDocument(c.collection.Self, &doc, options...)
 
 	if err != nil {
 		if req.ETag != nil {