Add API for creating hierarchical PartitionKeys (Azure#23577)

Co-authored-by: Joel Hendrix <[email protected]>
Co-authored-by: Ashley Stanton-Nurse <[email protected]>

sdk/data/azcosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 1.1.1 (Unreleased)
 
 ### Features Added
+* Added API for creating Hierarchical PartitionKeys. See [PR 23577](https://github.com/Azure/azure-sdk-for-go/pull/23577)
 * Set all Telemetry spans to have the Kind of SpanKindClient. See [PR 23618](https://github.com/Azure/azure-sdk-for-go/pull/23618)
 * Set request_charge and status_code on all trace spans. See [PR 23652](https://github.com/Azure/azure-sdk-for-go/pull/23652)
 

sdk/data/azcosmos/emulator_cosmos_item_test.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"net/http"
+	"reflect"
 	"strings"
 	"sync"
 	"testing"
@@ -503,6 +504,196 @@ func TestItemIdEncodingComputeGW(t *testing.T) {
 	verifyEncodingScenario(t, container, "ComputeGW-IdWithUnicodeCharacters", "WithUnicode鱀", http.StatusCreated, http.StatusOK, http.StatusOK, http.StatusNoContent)
 }
 
+func TestItemCRUDHierarchicalPartitionKey(t *testing.T) {
+	emulatorTests := newEmulatorTests(t)
+	client := emulatorTests.getClient(t, newSpanValidator(t, &spanMatcher{
+		ExpectedSpans: []string{},
+	}))
+
+	database := emulatorTests.createDatabase(t, context.TODO(), client, "itemCRUDHierarchicalPartitionKey")
+	defer emulatorTests.deleteDatabase(t, context.TODO(), database)
+	properties := ContainerProperties{
+		ID: "aContainer",
+		PartitionKeyDefinition: PartitionKeyDefinition{
+			Paths:   []string{"/id", "/type"},
+			Kind:    PartitionKeyKindMultiHash,
+			Version: 2,
+		},
+	}
+
+	_, err := database.CreateContainer(context.TODO(), properties, nil)
+	if err != nil {
+		t.Fatalf("Failed to create container: %v", err)
+	}
+
+	container, err := database.NewContainer("aContainer")
+	if err != nil {
+		t.Fatalf("Failed to get container: %v", err)
+	}
+
+	itemAlpha := map[string]interface{}{
+		"id":    "1",
+		"type":  "alpha",
+		"value": "0",
+	}
+
+	itemBeta := map[string]interface{}{
+		"id":    "1",
+		"type":  "beta",
+		"value": "0",
+	}
+
+	pkAlpha := NewPartitionKey().AppendString("1").AppendString("alpha")
+	pkBeta := NewPartitionKey().AppendString("1").AppendString("beta")
+
+	marshalledAlpha, err := json.Marshal(itemAlpha)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	marshalledBeta, err := json.Marshal(itemBeta)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	item0Res, err := container.CreateItem(context.TODO(), pkAlpha, marshalledAlpha, nil)
+	if err != nil {
+		t.Fatalf("Failed to create item: %v", err)
+	}
+
+	if item0Res.SessionToken == nil {
+		t.Fatalf("Session token is empty")
+	}
+
+	if len(item0Res.Value) != 0 {
+		t.Fatalf("Expected empty response, got %v", item0Res.Value)
+	}
+
+	item1Res, err := container.CreateItem(context.TODO(), pkBeta, marshalledBeta, nil)
+	if err != nil {
+		t.Fatalf("Failed to create item: %v", err)
+	}
+
+	if item1Res.SessionToken == nil {
+		t.Fatalf("Session token is empty")
+	}
+
+	if len(item1Res.Value) != 0 {
+		t.Fatalf("Expected empty response, got %v", item1Res.Value)
+	}
+
+	item0Res, err = container.ReadItem(context.TODO(), pkAlpha, "1", nil)
+	if err != nil {
+		t.Fatalf("Failed to read item: %v", err)
+	}
+
+	if len(item0Res.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", item0Res.Value)
+	}
+
+	item1Res, err = container.ReadItem(context.TODO(), pkBeta, "1", nil)
+	if err != nil {
+		t.Fatalf("Failed to read item: %v", err)
+	}
+
+	if len(item1Res.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", item1Res.Value)
+	}
+
+	var item0ResBody map[string]interface{}
+	err = json.Unmarshal(item0Res.Value, &item0ResBody)
+
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+
+	if item0ResBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", item0ResBody["id"])
+	}
+
+	if item0ResBody["type"] != "alpha" {
+		t.Fatalf("Expected type to be alpha, got %v", item0ResBody["type"])
+	}
+
+	if item0ResBody["value"] != "0" {
+		t.Fatalf("Expected value to be 0, got %v", item0ResBody["value"])
+	}
+
+	var item1ResBody map[string]interface{}
+	err = json.Unmarshal(item1Res.Value, &item1ResBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+
+	if item1ResBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", item1ResBody["id"])
+	}
+
+	if item1ResBody["type"] != "beta" {
+		t.Fatalf("Expected type to be beta, got %v", item1ResBody["type"])
+	}
+
+	if item1ResBody["value"] != "0" {
+		t.Fatalf("Expected value to be 0, got %v", item1ResBody["value"])
+	}
+
+	pager := container.NewQueryItemsPager("SELECT * FROM c", pkAlpha, nil)
+
+	var alphaItems []map[string]interface{}
+	for pager.More() {
+		page, err := pager.NextPage(context.TODO())
+		if err != nil {
+			t.Fatalf("Failed to get next page: %v", err)
+		}
+
+		for _, item := range page.Items {
+			var itemBody map[string]interface{}
+			err = json.Unmarshal(item, &itemBody)
+			if err != nil {
+				t.Fatalf("Failed to unmarshal item response: %v", err)
+			}
+
+			alphaItems = append(alphaItems, itemBody)
+		}
+	}
+
+	if len(alphaItems) != 1 {
+		t.Fatalf("Expected 1 item, got %v", len(alphaItems))
+	}
+
+	if !reflect.DeepEqual(alphaItems[0], item0ResBody) {
+		t.Fatalf("Expected %v, got %v", item0ResBody, alphaItems[0])
+	}
+
+	pager = container.NewQueryItemsPager("SELECT * FROM c", pkBeta, nil)
+
+	var betaItems []map[string]interface{}
+	for pager.More() {
+		page, err := pager.NextPage(context.TODO())
+		if err != nil {
+			t.Fatalf("Failed to get next page: %v", err)
+		}
+
+		for _, item := range page.Items {
+			var itemBody map[string]interface{}
+			err = json.Unmarshal(item, &itemBody)
+			if err != nil {
+				t.Fatalf("Failed to unmarshal item response: %v", err)
+			}
+
+			betaItems = append(betaItems, itemBody)
+		}
+	}
+
+	if len(betaItems) != 1 {
+		t.Fatalf("Expected 1 item, got %v", len(betaItems))
+	}
+
+	if !reflect.DeepEqual(betaItems[0], item1ResBody) {
+		t.Fatalf("Expected %v, got %v", item1ResBody, betaItems[0])
+	}
+}
+
 func verifyEncodingScenario(t *testing.T, container *ContainerClient, name string, id string, expectedCreate int, expectedRead int, expectedReplace int, expectedDelete int) {
 	item := map[string]interface{}{
 		"id": id,

sdk/data/azcosmos/partition_key.go

@@ -19,6 +19,13 @@ var NullPartitionKey PartitionKey = PartitionKey{
 	values: []interface{}{nil},
 }
 
+// NewPartitionKey creates a new partition key.
+func NewPartitionKey() PartitionKey {
+	return PartitionKey{
+		values: []interface{}{},
+	}
+}
+
 // NewPartitionKeyString creates a partition key with a string value.
 func NewPartitionKeyString(value string) PartitionKey {
 	components := []interface{}{value}
@@ -43,6 +50,30 @@ func NewPartitionKeyNumber(value float64) PartitionKey {
 	}
 }
 
+// AppendString appends a string value to the partition key.
+func (pk PartitionKey) AppendString(value string) PartitionKey {
+	pk.values = append(pk.values, value)
+	return pk
+}
+
+// AppendBool appends a boolean value to the partition key.
+func (pk PartitionKey) AppendBool(value bool) PartitionKey {
+	pk.values = append(pk.values, value)
+	return pk
+}
+
+// AppendNumber appends a numeric value to the partition key.
+func (pk PartitionKey) AppendNumber(value float64) PartitionKey {
+	pk.values = append(pk.values, value)
+	return pk
+}
+
+// AppendNull appends a null value to the partition key.
+func (pk PartitionKey) AppendNull() PartitionKey {
+	pk.values = append(pk.values, nil)
+	return pk
+}
+
 func (pk *PartitionKey) toJsonString() (string, error) {
 	var completeJson strings.Builder
 	completeJson.Grow(256)

sdk/data/azcosmos/partition_key_definition.go

@@ -3,11 +3,50 @@
 
 package azcosmos
 
+import (
+	"encoding/json"
+)
+
+// PartitionKeyKind represents the type of the partition key that is used in an Azure Cosmos DB container.
+type PartitionKeyKind string
+
+const (
+	PartitionKeyKindHash      PartitionKeyKind = "Hash"
+	PartitionKeyKindMultiHash PartitionKeyKind = "MultiHash"
+)
+
 // PartitionKeyDefinition represents a partition key definition in the Azure Cosmos DB database service.
 // A partition key definition defines the path for the partition key property.
 type PartitionKeyDefinition struct {
+	// Kind returns the kind of partition key definition.
+	Kind PartitionKeyKind `json:"kind"`
 	// Paths returns the list of partition key paths of the container.
 	Paths []string `json:"paths"`
 	// Version returns the version of the hash partitioning of the container.
 	Version int `json:"version,omitempty"`
 }
+
+// MarshalJSON implements the json.Marshaler interface
+// If the Kind is not set, it will be inferred based on the number of paths.
+func (pkd PartitionKeyDefinition) MarshalJSON() ([]byte, error) {
+	var paths_length = len(pkd.Paths)
+
+	var kind PartitionKeyKind
+	if pkd.Kind != "" {
+		kind = pkd.Kind
+	} else if pkd.Kind == "" && paths_length == 1 {
+		kind = PartitionKeyKindHash
+	} else if pkd.Kind == "" && paths_length > 1 {
+		kind = PartitionKeyKindMultiHash
+	}
+
+	return json.Marshal(struct {
+		Kind    PartitionKeyKind `json:"kind"`
+		Paths   []string         `json:"paths"`
+		Version int              `json:"version,omitempty"`
+	}{
+		Kind:    kind,
+		Paths:   pkd.Paths,
+		Version: pkd.Version,
+	})
+}

sdk/data/azcosmos/partition_key_definition_test.go

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"testing"
+)
+
+func TestPartitionKeyDefinitionSerialization(t *testing.T) {
+	pkd_kind_unset_len_one := PartitionKeyDefinition{
+		Paths:   []string{"somePath"},
+		Version: 2,
+	}
+
+	jsonString, err := pkd_kind_unset_len_one.MarshalJSON()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := `{"kind":"Hash","paths":["somePath"],"version":2}`
+	if string(jsonString) != expected {
+		t.Errorf("Expected serialization %v, but got %v", expected, string(jsonString))
+	}
+
+	pkd_kind_unset_len_two := PartitionKeyDefinition{
+		Paths:   []string{"somePath", "someOtherPath"},
+		Version: 2,
+	}
+
+	jsonString, err = pkd_kind_unset_len_two.MarshalJSON()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected = `{"kind":"MultiHash","paths":["somePath","someOtherPath"],"version":2}`
+	if string(jsonString) != expected {
+		t.Errorf("Expected serialization %v, but got %v", expected, string(jsonString))
+	}
+
+	pkd_kind_set := PartitionKeyDefinition{
+		Kind:    PartitionKeyKindMultiHash,
+		Paths:   []string{"somePath"},
+		Version: 2,
+	}
+
+	jsonString, err = pkd_kind_set.MarshalJSON()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected = `{"kind":"MultiHash","paths":["somePath"],"version":2}`
+	if string(jsonString) != expected {
+		t.Errorf("Expected serialization %v, but got %v", expected, string(jsonString))
+	}
+}