diff --git a/sdk/data/azcosmos/CHANGELOG.md b/sdk/data/azcosmos/CHANGELOG.md index 5403efa530c4..6e57cc226045 100644 --- a/sdk/data/azcosmos/CHANGELOG.md +++ b/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) diff --git a/sdk/data/azcosmos/emulator_cosmos_item_test.go b/sdk/data/azcosmos/emulator_cosmos_item_test.go index 9b0c59949d0d..37a62ae2cb88 100644 --- a/sdk/data/azcosmos/emulator_cosmos_item_test.go +++ b/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, diff --git a/sdk/data/azcosmos/partition_key.go b/sdk/data/azcosmos/partition_key.go index ca61aec87ae2..f31a67cecb18 100644 --- a/sdk/data/azcosmos/partition_key.go +++ b/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) diff --git a/sdk/data/azcosmos/partition_key_definition.go b/sdk/data/azcosmos/partition_key_definition.go index f77d3bdff8c0..c15f0cc9c071 100644 --- a/sdk/data/azcosmos/partition_key_definition.go +++ b/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, + }) +} diff --git a/sdk/data/azcosmos/partition_key_definition_test.go b/sdk/data/azcosmos/partition_key_definition_test.go new file mode 100644 index 000000000000..de3e8edbd261 --- /dev/null +++ b/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)) + } +} diff --git a/sdk/data/azcosmos/partition_key_test.go b/sdk/data/azcosmos/partition_key_test.go index bd7e651aba29..934343783439 100644 --- a/sdk/data/azcosmos/partition_key_test.go +++ b/sdk/data/azcosmos/partition_key_test.go @@ -34,6 +34,34 @@ func TestSerialization(t *testing.T) { } } +func TestPartitionKeyAppends(t *testing.T) { + validTypes := map[string]PartitionKey{ + "[\"key0\"]": NewPartitionKey().AppendString("key0"), + "[true]": NewPartitionKey().AppendBool(true), + "[false]": NewPartitionKey().AppendBool(false), + "[10.5]": NewPartitionKey().AppendNumber(10.5), + "[10]": NewPartitionKey().AppendNumber(10), + "[null]": NewPartitionKey().AppendNull(), + "[\"key0\",true,10.5]": NewPartitionKey().AppendString("key0").AppendBool(true).AppendNumber(10.5), + "[null,null,null]": NewPartitionKey().AppendNull().AppendNull().AppendNull(), + } + + for expectedSerialization, pk := range validTypes { + if len(pk.values) < 1 { + t.Errorf("Expected partition key to have at least 1 component, but it has %v", len(pk.values)) + } + + serialization, err := pk.toJsonString() + if err != nil { + t.Errorf("Failed to serialize PK for %v, got %v", pk, err) + } + + if serialization != expectedSerialization { + t.Errorf("Expected serialization %v, but got %v", expectedSerialization, serialization) + } + } +} + func TestPartitionKeyEquality(t *testing.T) { pk := NewPartitionKeyNumber(float64(10.5)) pk2 := NewPartitionKeyNumber(float64(10.5))