From 2bb7deb3228a4986b91fa80b11bedadb3eee6dd7 Mon Sep 17 00:00:00 2001
From: Matthew Barnes <mbarnes@fedorapeople.org>
Date: Wed, 16 Oct 2024 13:35:40 -0400
Subject: [PATCH] database: Store all Operations items in a single partition

The azcosmos SDK currently only supports single-partition queries,
so there's no way to list all items in a container unless you know
all the partition keys. The backend needs to list all items in the
Operations container so to work around this limitation we keep all
items in a single partition with a well-known name: "workaround".

Once [1] is fixed we could transition the Operations container to
using subscription IDs as the partition key like other containers.
The items are transient thanks to the container's default TTL, so
GetOperationDoc would just need temporary fallback logic to check
the "workaround" partition.

[1] https://github.com/Azure/azure-sdk-for-go/issues/18578
---
 dev-infrastructure/modules/rp-cosmos.bicep |  1 -
 internal/database/database.go              | 21 ++++++++++++++++++---
 internal/database/document.go              |  2 ++
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/dev-infrastructure/modules/rp-cosmos.bicep b/dev-infrastructure/modules/rp-cosmos.bicep
index 0d2d010c6..2cb6bb1a0 100644
--- a/dev-infrastructure/modules/rp-cosmos.bicep
+++ b/dev-infrastructure/modules/rp-cosmos.bicep
@@ -18,7 +18,6 @@ var containers = [
   {
     name: 'Operations'
     defaultTtl: 604800 // 7 days
-    partitionKeyPaths: ['/id']
   }
   {
     name: 'Resources'
diff --git a/internal/database/database.go b/internal/database/database.go
index eb4c9e378..db04a26a7 100644
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -24,6 +24,21 @@ const (
 	billingContainer    = "Billing"
 	operationsContainer = "Operations"
 	locksContainer      = "Locks"
+
+	// XXX The azcosmos SDK currently only supports single-partition queries,
+	//     so there's no way to list all items in a container unless you know
+	//     all the partition keys. The backend needs to list all items in the
+	//     Operations container so to work around this limitation we keep all
+	//     items in a single partition with a well-known name: "workaround".
+	//
+	//     Once [1] is fixed we could transition the Operations container to
+	//     using subscription IDs as the partition key like other containers.
+	//     The items are transient thanks to the container's default TTL, so
+	//     GetOperationDoc would just need temporary fallback logic to check
+	//     the "workaround" partition.
+	//
+	//     [1] https://github.com/Azure/azure-sdk-for-go/issues/18578
+	operationsPartitionKey = "workaround"
 )
 
 var ErrNotFound = errors.New("DocumentNotFound")
@@ -314,7 +329,7 @@ func (d *CosmosDBClient) GetOperationDoc(ctx context.Context, operationID string
 		return nil, err
 	}
 
-	pk := azcosmos.NewPartitionKeyString(operationID)
+	pk := azcosmos.NewPartitionKeyString(operationsPartitionKey)
 
 	response, err := container.ReadItem(ctx, pk, operationID, nil)
 	if isResponseError(err, http.StatusNotFound) {
@@ -340,7 +355,7 @@ func (d *CosmosDBClient) CreateOperationDoc(ctx context.Context, doc *OperationD
 		return err
 	}
 
-	pk := azcosmos.NewPartitionKeyString(doc.ID)
+	pk := azcosmos.NewPartitionKeyString(operationsPartitionKey)
 
 	data, err := json.Marshal(doc)
 	if err != nil {
@@ -366,7 +381,7 @@ func (d *CosmosDBClient) DeleteOperationDoc(ctx context.Context, operationID str
 		return err
 	}
 
-	pk := azcosmos.NewPartitionKeyString(operationID)
+	pk := azcosmos.NewPartitionKeyString(operationsPartitionKey)
 
 	_, err = container.DeleteItem(ctx, pk, operationID, nil)
 	if isResponseError(err, http.StatusNotFound) {
diff --git a/internal/database/document.go b/internal/database/document.go
index da5efe04d..39d5f8afc 100644
--- a/internal/database/document.go
+++ b/internal/database/document.go
@@ -64,6 +64,7 @@ const (
 type OperationDocument struct {
 	BaseDocument
 
+	PartitionKey string `json:"partitionKey,omitempty"`
 	// TenantID is the tenant ID of the client that requested the operation
 	TenantID string `json:"tenantId,omitempty"`
 	// ClientID is the object ID of the client that requested the operation
@@ -96,6 +97,7 @@ func NewOperationDocument(request OperationRequest) *OperationDocument {
 
 	return &OperationDocument{
 		BaseDocument:       newBaseDocument(),
+		PartitionKey:       operationsPartitionKey,
 		Request:            request,
 		StartTime:          now,
 		LastTransitionTime: now,