add partition key for greater flexibility

internal/api/arm/types_cosmosdata.go

@@ -30,6 +30,14 @@ import (
 type CosmosMetadata struct {
 	ResourceID *azcorearm.ResourceID `json:"resourceID"`
 
+	// PartitionKey is the Cosmos partition-key value for this document. It is stored
+	// lower-cased; SetPartitionKey normalises on input and GetPartitionKey on output.
+	// When unset, GetPartitionKey falls back to the lower-cased ResourceID.SubscriptionID,
+	// which preserves the historical behaviour for every container that is partitioned
+	// by subscription ID. Containers partitioned by something else (e.g. management
+	// cluster name for the kube-applier container) populate this field explicitly.
+	PartitionKey string `json:"partitionKey,omitempty"`
+
 	// ExistingCosmosUID exists to allow for a migration path from where we are today to a uuid based cosmosID
 	// and this will be deleted afterwards.
 	ExistingCosmosUID string `json:"-"`
@@ -57,10 +65,26 @@ func (o *CosmosMetadata) GetCosmosUID() string {
 	return Must(ResourceIDToCosmosID(o.ResourceID))
 }
 
+// GetPartitionKey returns the lower-cased partition key for this document.
+// If the explicit PartitionKey field has been set it is used; otherwise the
+// historical default (lower-cased ResourceID.SubscriptionID) is returned so that
+// existing subscription-keyed containers continue to work without changes.
 func (o *CosmosMetadata) GetPartitionKey() string {
+	if len(o.PartitionKey) > 0 {
+		return strings.ToLower(o.PartitionKey)
+	}
+	if o.ResourceID == nil {
+		return ""
+	}
 	return strings.ToLower(o.ResourceID.SubscriptionID)
 }
 
+// SetPartitionKey records partitionKey on the metadata, normalising to lower
+// case so reads and writes always agree on a single canonical form.
+func (o *CosmosMetadata) SetPartitionKey(partitionKey string) {
+	o.PartitionKey = strings.ToLower(partitionKey)
+}
+
 func (o *CosmosMetadata) GetResourceID() *azcorearm.ResourceID {
 	return o.ResourceID
 }
@@ -87,6 +111,8 @@ type CosmosMetadataAccessor interface {
 	SetResourceID(*azcorearm.ResourceID)
 	GetEtag() azcore.ETag
 	SetEtag(cosmosETag azcore.ETag)
+	GetPartitionKey() string
+	SetPartitionKey(partitionKey string)
 }
 
 func ResourceIDToCosmosID(resourceID *azcorearm.ResourceID) (string, error) {

internal/api/arm/types_cosmosdata_test.go

@@ -0,0 +1,161 @@
+// Copyright 2026 Microsoft Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package arm
+
+import (
+	"encoding/json"
+	"testing"
+
+	azcorearm "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+)
+
+func mustParseTestID(t *testing.T, s string) *azcorearm.ResourceID {
+	t.Helper()
+	id, err := azcorearm.ParseResourceID(s)
+	if err != nil {
+		t.Fatalf("parse %q: %v", s, err)
+	}
+	return id
+}
+
+func TestCosmosMetadataPartitionKey(t *testing.T) {
+	const idStr = "/subscriptions/MyUpperCaseSub/resourceGroups/rg/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/c"
+
+	tests := []struct {
+		name string
+		// Each test mutates the metadata and asserts on Get/Set behaviour.
+		runWith func(t *testing.T, m *CosmosMetadata)
+	}{
+		{
+			name: "GetPartitionKey falls back to lowercased subscription ID when field is empty",
+			runWith: func(t *testing.T, m *CosmosMetadata) {
+				if got, want := m.GetPartitionKey(), "myuppercasesub"; got != want {
+					t.Errorf("GetPartitionKey() = %q, want %q (lowercased SubscriptionID)", got, want)
+				}
+			},
+		},
+		{
+			name: "SetPartitionKey lowercases on store",
+			runWith: func(t *testing.T, m *CosmosMetadata) {
+				m.SetPartitionKey("MGMT-Cluster-1")
+				if got, want := m.PartitionKey, "mgmt-cluster-1"; got != want {
+					t.Errorf("PartitionKey field = %q, want %q (lowercased on Set)", got, want)
+				}
+			},
+		},
+		{
+			name: "GetPartitionKey returns the stored field when set, lowercased",
+			runWith: func(t *testing.T, m *CosmosMetadata) {
+				m.SetPartitionKey("MGMT-Cluster-1")
+				if got, want := m.GetPartitionKey(), "mgmt-cluster-1"; got != want {
+					t.Errorf("GetPartitionKey() = %q, want %q", got, want)
+				}
+			},
+		},
+		{
+			name: "GetPartitionKey lowercases even if the field was set directly with mixed case",
+			runWith: func(t *testing.T, m *CosmosMetadata) {
+				m.PartitionKey = "MGMT-Cluster-1" // bypass setter
+				if got, want := m.GetPartitionKey(), "mgmt-cluster-1"; got != want {
+					t.Errorf("GetPartitionKey() = %q, want %q (lowercased on Get)", got, want)
+				}
+			},
+		},
+		{
+			name: "SetPartitionKey then GetPartitionKey is idempotent",
+			runWith: func(t *testing.T, m *CosmosMetadata) {
+				m.SetPartitionKey("foo")
+				m.SetPartitionKey(m.GetPartitionKey())
+				if got := m.PartitionKey; got != "foo" {
+					t.Errorf("idempotent set yielded %q", got)
+				}
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			m := &CosmosMetadata{
+				ResourceID: mustParseTestID(t, idStr),
+			}
+			tc.runWith(t, m)
+		})
+	}
+}
+
+func TestCosmosMetadataPartitionKey_NilResourceID(t *testing.T) {
+	m := &CosmosMetadata{}
+	if got := m.GetPartitionKey(); got != "" {
+		t.Errorf("GetPartitionKey() with nil ResourceID and unset PartitionKey = %q, want empty", got)
+	}
+	m.SetPartitionKey("X")
+	if got := m.GetPartitionKey(); got != "x" {
+		t.Errorf("GetPartitionKey() after Set with nil ResourceID = %q, want %q", got, "x")
+	}
+}
+
+func TestCosmosMetadataJSONRoundTrip_PartitionKey(t *testing.T) {
+	m := &CosmosMetadata{
+		ResourceID: mustParseTestID(t,
+			"/subscriptions/sub/resourceGroups/rg/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/c"),
+	}
+	m.SetPartitionKey("MGMT-1")
+
+	data, err := json.Marshal(m)
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	// Should serialize the lowercased value under "partitionKey".
+	if want := `"partitionKey":"mgmt-1"`; !contains(string(data), want) {
+		t.Errorf("marshalled JSON did not contain %q\n  got: %s", want, data)
+	}
+
+	var got CosmosMetadata
+	if err := json.Unmarshal(data, &got); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	if got.GetPartitionKey() != "mgmt-1" {
+		t.Errorf("after round-trip, GetPartitionKey() = %q, want %q", got.GetPartitionKey(), "mgmt-1")
+	}
+}
+
+func TestCosmosMetadataJSONRoundTrip_OmitsEmptyPartitionKey(t *testing.T) {
+	// Older documents on disk don't have partitionKey in cosmosMetadata. Round-tripping
+	// without setting the field must not introduce one (omitempty), and GetPartitionKey
+	// must still fall back to the subscription ID.
+	m := &CosmosMetadata{
+		ResourceID: mustParseTestID(t,
+			"/subscriptions/sub/resourceGroups/rg/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/c"),
+	}
+	data, err := json.Marshal(m)
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	if contains(string(data), "partitionKey") {
+		t.Errorf("expected omitempty to drop partitionKey from JSON, got: %s", data)
+	}
+	if got, want := m.GetPartitionKey(), "sub"; got != want {
+		t.Errorf("GetPartitionKey() with unset field = %q, want %q (subscription ID fallback)", got, want)
+	}
+}
+
+func contains(s, substr string) bool {
+	for i := 0; i+len(substr) <= len(s); i++ {
+		if s[i:i+len(substr)] == substr {
+			return true
+		}
+	}
+	return false
+}

internal/database/convert_cluster.go

@@ -16,7 +16,6 @@ package database
 
 import (
 	"fmt"
-	"strings"
 
 	"k8s.io/utils/ptr"
 
@@ -29,18 +28,20 @@ func InternalToCosmosCluster(internalObj *api.HCPOpenShiftCluster) (*HCPCluster,
 		return nil, nil
 	}
 
+	partitionKey := internalObj.GetCosmosData().GetPartitionKey()
 	cosmosObj := &HCPCluster{
 		TypedDocument: TypedDocument{
 			BaseDocument: BaseDocument{
 				ID: internalObj.GetCosmosData().GetCosmosUID(),
 			},
-			PartitionKey: strings.ToLower(internalObj.ID.SubscriptionID),
+			PartitionKey: partitionKey,
 			ResourceID:   internalObj.ID,
 			ResourceType: internalObj.ID.ResourceType.String(),
 		},
 		HCPClusterProperties: HCPClusterProperties{
 			CosmosMetadata: api.CosmosMetadata{
-				ResourceID: internalObj.ID,
+				ResourceID:   internalObj.ID,
+				PartitionKey: partitionKey,
 			},
 			IntermediateResourceDoc: &ResourceDocument{
 				ResourceID:        internalObj.ID,

internal/database/convert_externalauth.go

@@ -16,7 +16,6 @@ package database
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/Azure/ARO-HCP/internal/api"
 	"github.com/Azure/ARO-HCP/internal/api/arm"
@@ -28,18 +27,20 @@ func InternalToCosmosExternalAuth(internalObj *api.HCPOpenShiftClusterExternalAu
 		return nil, nil
 	}
 
+	partitionKey := internalObj.GetCosmosData().GetPartitionKey()
 	cosmosObj := &ExternalAuth{
 		TypedDocument: TypedDocument{
 			BaseDocument: BaseDocument{
 				ID: internalObj.GetCosmosData().GetCosmosUID(),
 			},
-			PartitionKey: strings.ToLower(internalObj.ID.SubscriptionID),
+			PartitionKey: partitionKey,
 			ResourceID:   internalObj.ID,
 			ResourceType: internalObj.ID.ResourceType.String(),
 		},
 		ExternalAuthProperties: ExternalAuthProperties{
 			CosmosMetadata: api.CosmosMetadata{
-				ResourceID: internalObj.ID,
+				ResourceID:   internalObj.ID,
+				PartitionKey: partitionKey,
 			},
 			IntermediateResourceDoc: &ResourceDocument{
 				ResourceID:        internalObj.ID,

internal/database/convert_generic.go

@@ -16,7 +16,6 @@ package database
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/Azure/ARO-HCP/internal/api"
 	"github.com/Azure/ARO-HCP/internal/api/arm"
@@ -34,17 +33,24 @@ func InternalToCosmosGeneric[InternalAPIType any](internalObj *InternalAPIType)
 		return nil, fmt.Errorf("internalObj must be an arm.CosmosMetadataAccessor: %T", internalObj)
 	}
 
+	partitionKey := metadata.GetPartitionKey()
 	cosmosObj := &GenericDocument[InternalAPIType]{
 		TypedDocument: TypedDocument{
 			BaseDocument: BaseDocument{
 				ID: metadata.GetCosmosUID(),
 			},
-			PartitionKey: strings.ToLower(metadata.GetResourceID().SubscriptionID),
+			PartitionKey: partitionKey,
 			ResourceID:   metadata.GetResourceID(),
 			ResourceType: metadata.GetResourceID().ResourceType.String(),
 		},
 		Content: *internalObj,
 	}
+	// Mirror the envelope's partitionKey into the inner cosmosMetadata copy so the on-disk
+	// representation has both fields in sync. We mutate the value-copy in cosmosObj.Content
+	// rather than the caller-supplied internalObj.
+	if cm, ok := any(&cosmosObj.Content).(arm.CosmosMetadataAccessor); ok {
+		cm.SetPartitionKey(partitionKey)
+	}
 
 	// this isn't pretty, but on balance it's a better choice so that we can share all the rest.
 	switch any(internalObj).(type) {
@@ -68,6 +74,9 @@ func CosmosGenericToInternal[InternalAPIType any](cosmosObj *GenericDocument[Int
 	cosmosData := ret.(arm.CosmosPersistable).GetCosmosData()
 	cosmosData.ExistingCosmosUID = cosmosObj.ID
 	ret.SetEtag(cosmosObj.CosmosETag)
+	// Round-trip the envelope's partitionKey back into the metadata so callers
+	// can read it without re-deriving from ResourceID.SubscriptionID.
+	ret.SetPartitionKey(cosmosObj.PartitionKey)
 
 	// this isn't pretty, but on balance it's a better choice so that we can share all the rest.
 	switch castObj := any(ret).(type) {

internal/database/convert_nodepool.go

@@ -16,7 +16,6 @@ package database
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/Azure/ARO-HCP/internal/api"
 	"github.com/Azure/ARO-HCP/internal/api/arm"
@@ -28,18 +27,20 @@ func InternalToCosmosNodePool(internalObj *api.HCPOpenShiftClusterNodePool) (*No
 		return nil, nil
 	}
 
+	partitionKey := internalObj.GetCosmosData().GetPartitionKey()
 	cosmosObj := &NodePool{
 		TypedDocument: TypedDocument{
 			BaseDocument: BaseDocument{
 				ID: internalObj.GetCosmosData().GetCosmosUID(),
 			},
-			PartitionKey: strings.ToLower(internalObj.ID.SubscriptionID),
+			PartitionKey: partitionKey,
 			ResourceID:   internalObj.ID,
 			ResourceType: internalObj.ID.ResourceType.String(),
 		},
 		NodePoolProperties: NodePoolProperties{
 			CosmosMetadata: api.CosmosMetadata{
-				ResourceID: internalObj.ID,
+				ResourceID:   internalObj.ID,
+				PartitionKey: partitionKey,
 			},
 			IntermediateResourceDoc: &ResourceDocument{
 				ResourceID:        internalObj.ID,

internal/database/crud_nested_resource.go

@@ -146,12 +146,12 @@ func (d *nestedCosmosResourceCRUD[InternalAPIType, CosmosAPIType]) AddReplaceToT
 }
 
 func (d *nestedCosmosResourceCRUD[InternalAPIType, CosmosAPIType]) Create(ctx context.Context, newObj *InternalAPIType, options *azcosmos.ItemOptions) (*InternalAPIType, error) {
-	partitionKey := strings.ToLower(any(newObj).(arm.CosmosPersistable).GetCosmosData().GetResourceID().SubscriptionID)
+	partitionKey := any(newObj).(arm.CosmosPersistable).GetCosmosData().GetPartitionKey()
 	return create[InternalAPIType, CosmosAPIType](ctx, d.containerClient, partitionKey, newObj, options)
 }
 
 func (d *nestedCosmosResourceCRUD[InternalAPIType, CosmosAPIType]) Replace(ctx context.Context, newObj *InternalAPIType, options *azcosmos.ItemOptions) (*InternalAPIType, error) {
-	partitionKey := strings.ToLower(any(newObj).(arm.CosmosPersistable).GetCosmosData().GetResourceID().SubscriptionID)
+	partitionKey := any(newObj).(arm.CosmosPersistable).GetCosmosData().GetPartitionKey()
 	return replace[InternalAPIType, CosmosAPIType](ctx, d.containerClient, partitionKey, newObj, options)
 }
 

test-integration/backend/controllers/do_nothing/artifacts/sync_cluster/99-cosmosCompare-end-state/do-nothing-success.json

@@ -8,6 +8,7 @@
 	"partitionKey": "4fa75980-6637-4157-9726-84d878a62e83",
 	"properties": {
 		"cosmosMetadata": {
+			"partitionKey": "4fa75980-6637-4157-9726-84d878a62e83",
 			"resourceID": "/subscriptions/4fa75980-6637-4157-9726-84d878a62e83/resourceGroups/shrilleffectiveness/providers/microsoft.redhatopenshift/hcpopenshiftclusters/lavishunhappiness/hcpOpenShiftControllers/DoNothingExample"
 		},
 		"externalId": "/subscriptions/4fa75980-6637-4157-9726-84d878a62e83/resourceGroups/shrilleffectiveness/providers/microsoft.redhatopenshift/hcpopenshiftclusters/lavishunhappiness",

test-integration/backend/controllers/mismatches/artifacts/cluster/remove_orphaned_cluster_descendents/00-load-initial-state/controller-nodepool-fake.json

@@ -8,6 +8,7 @@
 	"partitionKey": "a433a095-1277-44f1-8453-8d61a4d848c2",
 	"properties": {
 		"cosmosMetadata": {
+			"partitionKey": "a433a095-1277-44f1-8453-8d61a4d848c2",
 			"resourceID": "/subscriptions/a433a095-1277-44f1-8453-8d61a4d848c2/resourceGroups/unimportantpostponement/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/monstrousprecinct/nodepools/basic/hcpOpenShiftControllers/DoNothingExample"
 		},
 		"externalId": "/subscriptions/a433a095-1277-44f1-8453-8d61a4d848c2/resourceGroups/unimportantpostponement/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/monstrousprecinct/nodepools/basic",

test-integration/backend/controllers/mismatches/artifacts/delete_orphaned_cosmos/all_parents_exist/00-load-initial-state/controller-cluster.json

@@ -3,6 +3,7 @@
 	"partitionKey": "a433a095-1277-44f1-8453-8d61a4d848c2",
 	"properties": {
 		"cosmosMetadata": {
+			"partitionKey": "a433a095-1277-44f1-8453-8d61a4d848c2",
 			"resourceID": "/subscriptions/a433a095-1277-44f1-8453-8d61a4d848c2/resourceGroups/unimportantpostponement/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/monstrousprecinct/hcpOpenShiftControllers/testcontroller"
 		},
 		"externalId": "/subscriptions/a433a095-1277-44f1-8453-8d61a4d848c2/resourceGroups/unimportantpostponement/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/monstrousprecinct",

test-integration/backend/controllers/mismatches/artifacts/delete_orphaned_cosmos/all_parents_exist/99-cosmosCompare-end-state/controller-cluster.json

@@ -3,6 +3,7 @@
 	"partitionKey": "a433a095-1277-44f1-8453-8d61a4d848c2",
 	"properties": {
 		"cosmosMetadata": {
+			"partitionKey": "a433a095-1277-44f1-8453-8d61a4d848c2",
 			"resourceID": "/subscriptions/a433a095-1277-44f1-8453-8d61a4d848c2/resourceGroups/unimportantpostponement/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/monstrousprecinct/hcpOpenShiftControllers/testcontroller"
 		},
 		"externalId": "/subscriptions/a433a095-1277-44f1-8453-8d61a4d848c2/resourceGroups/unimportantpostponement/providers/Microsoft.RedHatOpenShift/hcpOpenShiftClusters/monstrousprecinct",