feat(storage): Simulate object custom contexts (#764)

b/457332685

This simulates the Json and gRPC responses for object contexts with 
the following operations:

1. Insert or Complete Update
    - Input: {contexts: {custom: {key: value}}}
    - Update_mask: contexts
    - Simulation: Add new entries to the object.contexts and furnish with create_time and update_time

2. Patch (Update an existing key)
    - Input: {contexts: {custom: {key: new-value}}}
    - Update_mask: contexts.custom.key
    - Simulation: Update the value of the key to "new-value" and furnish with update_time

3. Patch (Delete an existing key)
    - Input: {contexts: {custom: {key: null}}}
    - Update_mask: contexts.custom.key
    - Simulation: Delete that entry with that key from object.contexts

4. Patch (Delete all the contexts)
    - Input: {contexts: {custom: null}}
    - Update_mask: contexts:custom
    - Simulation: Delete the contexts field from object, making that object looks like its default

Author: xlai20

gcs/object.py

@@ -61,6 +61,7 @@ class Object:
         "event_based_hold",
         "customer_encryption",
         "custom_time",
+        "contexts",
     ]
 
     def __init__(self, metadata, media, bucket, *, upload=None, upload_gen=0):
@@ -123,6 +124,11 @@ def init(
         metadata.checksums.crc32c = actual_crc32c
         metadata.create_time.FromDatetime(timestamp)
         metadata.update_time.FromDatetime(timestamp)
+        if metadata.HasField("contexts"):
+            for _, payload in metadata.contexts.custom.items():
+                payload.create_time.FromDatetime(timestamp)
+                payload.update_time.FromDatetime(timestamp)
+            cls.__validate_object_contexts(metadata.contexts)
         upload_gen = 1
         if upload is None:
             upload_gen = 0
@@ -252,6 +258,29 @@ def __update_metadata(self, source, update_mask):
         self.metadata.etag = Object._metadata_etag(self.metadata)
         self.metadata.update_time.FromDatetime(datetime.datetime.now())
 
+    def __update_contexts_with_timestamps(
+        self, new_metadata, original_metadata, isUpdate
+    ):
+        if not new_metadata.HasField("contexts"):
+            return
+        if not new_metadata.contexts.custom:
+            # Equivalent to {"contexts": {"custom": None}}
+            new_metadata.ClearField("contexts")
+            return
+        timestamp = datetime.datetime.now(datetime.timezone.utc)
+        for key, payload in new_metadata.contexts.custom.items():
+            if isUpdate or key not in original_metadata.contexts.custom:
+                # This is a brand new key, set create and update timestamps
+                payload.create_time.FromDatetime(timestamp)
+                payload.update_time.FromDatetime(timestamp)
+            elif (
+                key in original_metadata.contexts.custom
+                and original_metadata.contexts.custom[key].value != payload.value
+            ):
+                # This is an existing key with new value, set update timestamp
+                payload.update_time.FromDatetime(timestamp)
+        self.__validate_object_contexts(new_metadata.contexts)
+
     def update(self, request, context):
         # Support for `Object: update` over gRPC is not needed (and not implemented).
         assert context is None
@@ -265,6 +294,7 @@ def update(self, request, context):
             testbench.acl.extract_predefined_acl(request, False, context),
             context,
         )
+        self.__update_contexts_with_timestamps(metadata, self.metadata, True)
         self.__update_metadata(metadata, None)
 
     def patch(self, request, context):
@@ -285,8 +315,48 @@ def patch(self, request, context):
             testbench.acl.extract_predefined_acl(request, False, context),
             context,
         )
+        self.__update_contexts_with_timestamps(metadata, self.metadata, False)
         self.__update_metadata(metadata, None)
 
+    @staticmethod
+    def __validate_object_contexts(contexts) -> bool:
+        """Validates an object context map against API layer rules."""
+        assert contexts is not None
+        custom_contexts = contexts.custom
+        if len(custom_contexts) > 50:
+            raise ValueError("The count of object context entries cannot exceed 50.")
+        invalid_chars = {"'", '"', "\\", "/"}
+        total_size_bytes = 0
+        for key, payload in custom_contexts.items():
+            val = payload.value
+            if key.startswith("goog"):
+                raise ValueError(
+                    f"Key '{key}' is invalid. Keys cannot begin with 'goog'."
+                )
+            for item, item_type in ((key, "Key"), (val, "Value")):
+                if not item or not item[0].isalnum():
+                    raise ValueError(
+                        f"{item_type} '{item}' must begin with an alphanumeric character."
+                    )
+                if any(char in invalid_chars for char in item):
+                    raise ValueError(
+                        f"{item_type} '{item}' contains restricted characters (', \", \\, /)."
+                    )
+                encoded_item = item.encode("utf-8")
+                item_length = len(encoded_item)
+                if not (1 <= item_length <= 256):
+                    raise ValueError(
+                        f"{item_type} '{item}' must be between 1 and 256 UTF-8 code units."
+                    )
+                total_size_bytes += item_length
+        max_size_bytes = 25 * 1024
+        if total_size_bytes > max_size_bytes:
+            raise ValueError(
+                f"Aggregate size of keys and values ({total_size_bytes} bytes) exceeds the 25 KiB limit."
+            )
+
+        return True
+
     # === ACL === #
 
     def __search_acl(self, entity, must_exist, context):

testbench/grpc_server.py

@@ -157,6 +157,56 @@ def decorated(self, request, context):
     return decorated
 
 
+def _validate_object_contexts(contexts, grpc_context):
+    """
+    Validates a storage_pb2.Object.contexts message against API layer rules.
+    If validation fails, it aborts the gRPC request using testbench.error.invalid.
+    """
+    if not contexts or not contexts.custom:
+        return
+
+    custom_contexts = contexts.custom
+    if len(custom_contexts) > 50:
+        return testbench.error.invalid(
+            "The count of object context entries cannot exceed 50.", grpc_context
+        )
+
+    invalid_chars = {"'", '"', "\\", "/"}
+    total_size_bytes = 0
+
+    for key, payload in custom_contexts.items():
+        val = payload.value
+        if key.startswith("goog"):
+            return testbench.error.invalid(
+                f"Key '{key}' is invalid. Keys cannot begin with 'goog'.", grpc_context
+            )
+        for item, item_type in ((key, "Key"), (val, "Value")):
+            if not item or not item[0].isalnum():
+                return testbench.error.invalid(
+                    f"{item_type} '{item}' must begin with an alphanumeric character.",
+                    grpc_context,
+                )
+            if any(char in invalid_chars for char in item):
+                return testbench.error.invalid(
+                    f"{item_type} '{item}' contains restricted characters (', \", \\, /).",
+                    grpc_context,
+                )
+            encoded_item = item.encode("utf-8")
+            item_length = len(encoded_item)
+            if not (1 <= item_length <= 256):
+                return testbench.error.invalid(
+                    f"{item_type} '{item}' must be between 1 and 256 UTF-8 code units.",
+                    grpc_context,
+                )
+            total_size_bytes += item_length
+    max_size_bytes = 25 * 1024
+    if total_size_bytes > max_size_bytes:
+        return testbench.error.invalid(
+            f"Aggregate size of keys and values ({total_size_bytes} bytes) exceeds the 25 KiB limit.",
+            grpc_context,
+        )
+
+
 def retry_test(method):
     """
     Decorate a routing function to handle the Retry Test API instructions,
@@ -812,6 +862,10 @@ def UpdateObject(self, request, context):
         updated_metadata = dict()
         removed_metadata_keys = set()
         replace_metadata = False
+        updated_contexts = dict()
+        removed_contexts_keys = set()
+        replace_contexts = False
+        reset_contexts = False
         for path in request.update_mask.paths:
             if path == "metadata":
                 replace_metadata = True
@@ -822,6 +876,17 @@ def UpdateObject(self, request, context):
                     removed_metadata_keys.add(key)
                 else:
                     updated_metadata[key] = value
+            elif path == "contexts":
+                replace_contexts = True
+            elif path == "contexts.custom":
+                reset_contexts = True
+            elif path.startswith("contexts.custom."):
+                key = path[len("contexts.custom.") :]
+                value = request.object.contexts.custom.get(key, None)
+                if value is None:
+                    removed_contexts_keys.add(key)
+                else:
+                    updated_contexts[key] = value
             elif path == "acl":
                 pass
             else:
@@ -873,8 +938,39 @@ def update_impl(blob, live_generation) -> storage_pb2.Object:
                 object.metadata.update(updated_metadata)
                 for k in removed_metadata_keys:
                     object.metadata.pop(k, None)
+            # Manually handle custom contexts due to its complexity.
+            curr_time = datetime.datetime.now()
+            if reset_contexts:
+                object.ClearField("contexts")
+            elif replace_contexts:
+                object.contexts.custom.clear()
+                for k, v in request.object.contexts.custom.items():
+                    dest_payload = object.contexts.custom[k]
+                    dest_payload.CopyFrom(v)
+                    dest_payload.create_time.FromDatetime(curr_time)
+                    dest_payload.update_time.FromDatetime(curr_time)
+            else:
+                for k, v in updated_contexts.items():
+                    if k not in object.contexts.custom:
+                        # This is a brand new key, add create and update time.
+                        updated_contexts[k].create_time.FromDatetime(curr_time)
+                        updated_contexts[k].update_time.FromDatetime(curr_time)
+                    elif v.value != object.contexts.custom[k].value:
+                        # This is an existing key, copy original key's create
+                        # time and add new update time.
+                        updated_contexts[k].create_time.CopyFrom(
+                            object.contexts.custom[k].create_time
+                        )
+                        updated_contexts[k].update_time.FromDatetime(curr_time)
+                    object.contexts.custom[k].CopyFrom(v)
+                for k in removed_contexts_keys:
+                    object.contexts.custom.pop(k, None)
+
+            if object.HasField("contexts"):
+                _validate_object_contexts(object.contexts, context)
+
             object.metageneration += 1
-            object.update_time.FromDatetime(datetime.datetime.now())
+            object.update_time.FromDatetime(curr_time)
             return object
 
         self.db.insert_test_bucket()

tests/test_grpc_server.py

@@ -1142,6 +1142,118 @@ def test_update_object_metadata(self):
         )
         self.assertEqual(response.metadata, {**response.metadata, **expected})
 
+    def test_update_object_contexts(self):
+        media = b"How vexingly quick daft zebras jump!"
+        request = testbench.common.FakeRequest(
+            args={"name": "object-name"},
+            data=media,
+            headers={},
+            environ={},
+        )
+        blob, _ = gcs.object.Object.init_media(request, self.bucket.metadata)
+        self.db.insert_object("bucket-name", blob, None)
+
+        def assert_context_valid(contexts, key, expected_value):
+            payload = contexts.custom.get(key)
+            self.assertIsNotNone(payload, f"Key '{key}' missing from custom contexts")
+            self.assertEqual(payload.value, expected_value)
+            # Verify timestamps are set (non-default)
+            self.assertGreater(
+                payload.create_time.seconds, 0, f"create_time not set for {key}"
+            )
+            self.assertGreater(
+                payload.update_time.seconds, 0, f"update_time not set for {key}"
+            )
+
+        request = storage_pb2.UpdateObjectRequest(
+            object=storage_pb2.Object(
+                bucket="projects/_/buckets/bucket-name",
+                name="object-name",
+                contexts=storage_pb2.ObjectContexts(
+                    custom={
+                        "location": storage_pb2.ObjectCustomContextPayload(
+                            value="Canada"
+                        ),
+                        "year": storage_pb2.ObjectCustomContextPayload(value="2026"),
+                    }
+                ),
+            ),
+            update_mask=field_mask_pb2.FieldMask(
+                paths=["contexts.custom.location", "contexts.custom.year"]
+            ),
+        )
+        context = unittest.mock.Mock()
+        response = self.grpc.UpdateObject(request, context)
+        assert_context_valid(response.contexts, "location", "Canada")
+        assert_context_valid(response.contexts, "year", "2026")
+
+        # Finally verify the changes are persisted
+        context = unittest.mock.Mock()
+        response = self.grpc.GetObject(
+            storage_pb2.GetObjectRequest(
+                bucket="projects/_/buckets/bucket-name", object="object-name"
+            ),
+            context,
+        )
+        assert_context_valid(response.contexts, "location", "Canada")
+        assert_context_valid(response.contexts, "year", "2026")
+
+    def test_update_object_contexts_invalid(self):
+        media = b"How vexingly quick daft zebras jump!"
+        request = testbench.common.FakeRequest(
+            args={"name": "object-name"},
+            data=media,
+            headers={},
+            environ={},
+        )
+        blob, _ = gcs.object.Object.init_media(request, self.bucket.metadata)
+        self.db.insert_object("bucket-name", blob, None)
+
+        def assert_invalid_context(custom_dict):
+            context = unittest.mock.Mock()
+            request = storage_pb2.UpdateObjectRequest(
+                object=storage_pb2.Object(
+                    bucket="projects/_/buckets/bucket-name",
+                    name="object-name",
+                    contexts=storage_pb2.ObjectContexts(
+                        custom={
+                            k: storage_pb2.ObjectCustomContextPayload(value=v)
+                            for k, v in custom_dict.items()
+                        }
+                    ),
+                ),
+                # Using the broader 'contexts' mask to trigger the replacement logic
+                # and evaluate the entire custom dictionary we just built
+                update_mask=field_mask_pb2.FieldMask(paths=["contexts"]),
+            )
+            _ = self.grpc.UpdateObject(request, context)
+            context.abort.assert_called_once_with(
+                grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
+            )
+
+        # --- Keys cannot begin with reserved 'goog' ---
+        assert_invalid_context({"googlekey": "value"})
+
+        # --- Must begin with an alphanumeric character ---
+        assert_invalid_context({"-badkey": "value"})  # Bad Key
+        assert_invalid_context({"validKey": ".badvalue"})  # Bad Value
+
+        # --- Cannot contain ', \", \\, or / ---
+        assert_invalid_context({"bad'key": "value"})
+        assert_invalid_context({"validKey": 'bad"value'})
+        assert_invalid_context({"validKey": "bad/value"})
+        assert_invalid_context({"validKey": "bad\\value"})
+
+        # --- Must be between 1 and 256 UTF-8 code units ---
+        assert_invalid_context({"": "value"})  # 0 length key
+        assert_invalid_context({"key": ""})  # 0 length value
+        assert_invalid_context({"a" * 257: "value"})  # Key too long
+        assert_invalid_context({"key": "v" * 257})  # Value too long
+
+        # --- Limit to 50 entries per object ---
+        too_many_contexts = {f"key{i}": "val" for i in range(51)}
+        assert_invalid_context(too_many_contexts)
+
     def test_update_object_acl(self):
         media = b"How vexingly quick daft zebras jump!"
         request = testbench.common.FakeRequest(