From d6474d307b2ed646f70b0234a41c2f773c2549a1 Mon Sep 17 00:00:00 2001
From: Tom Krause <thomas.krause@descarteslabs.com>
Date: Tue, 3 Oct 2023 08:59:39 -0600
Subject: [PATCH] Compute: support admin's being able to access other users
 results (#12205)

GitOrigin-RevId: 2902fd09588310db3f56cba4e3deee047a07791f
---
 README.md                                    |  2 ++
 descarteslabs/core/compute/compute_client.py | 16 +++++++++++++++-
 descarteslabs/core/compute/function.py       | 17 +++++++++++++++++
 descarteslabs/core/compute/job.py            | 18 +++++++++---------
 4 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index a513ed4c..c71f8ea9 100644
--- a/README.md
+++ b/README.md
@@ -80,6 +80,8 @@ Changelog
 
 - The `Job.result_blob()` will return the Catalog Storage Blob holding the result, if any.
 
+- The `Function` object now has attributes `namespace` and `owner`.
+
 - The `Function.wait_for_completion()` and new `Function.as_completed()` methods provide a richer
   set of functionality for waiting on and handling job completion.
 
diff --git a/descarteslabs/core/compute/compute_client.py b/descarteslabs/core/compute/compute_client.py
index 71bc6ce3..79514622 100644
--- a/descarteslabs/core/compute/compute_client.py
+++ b/descarteslabs/core/compute/compute_client.py
@@ -14,7 +14,7 @@
 
 import json
 from datetime import datetime, timezone
-from typing import Iterator
+from typing import Iterator, Optional
 
 from descarteslabs.auth import Auth
 from descarteslabs.config import get_settings
@@ -25,6 +25,8 @@
 
 
 class ComputeClient(ApiService, DefaultClientMixin):
+    _namespace_cache = dict()
+
     def __init__(self, url=None, auth=None, catalog_client=None, retries=None):
         if auth is None:
             auth = Auth.get_default_auth()
@@ -64,3 +66,15 @@ def set_credentials(self):
                 "client_secret": self.auth.client_secret,
             },
         )
+
+    def get_namespace(self, function_id: str) -> Optional[str]:
+        if function_id in self._namespace_cache:
+            return self._namespace_cache[function_id]
+
+        return None
+
+    def set_namespace(self, function_id: str, namespace: str):
+        if not (function_id or namespace):
+            return
+
+        self._namespace_cache[function_id] = namespace
diff --git a/descarteslabs/core/compute/function.py b/descarteslabs/core/compute/function.py
index 8fd5f452..95dc9f52 100644
--- a/descarteslabs/core/compute/function.py
+++ b/descarteslabs/core/compute/function.py
@@ -151,6 +151,18 @@ class Function(Document):
         sortable=True,
         doc="The maximum number of Jobs that execute at the same time for this Function.",
     )
+    namespace: str = Attribute(
+        str,
+        filterable=True,
+        readonly=True,
+        doc="The storage namespace for the Function.",
+    )
+    owner: str = Attribute(
+        str,
+        filterable=True,
+        readonly=True,
+        doc="The owner of the Function.",
+    )
     status: FunctionStatus = Attribute(
         FunctionStatus,
         filterable=True,
@@ -288,6 +300,11 @@ def __init__(
             else:
                 name = self._function.__name__
 
+        # When a Function is hydrated from the server, register the namespace
+        # with the client so that it can be used for subsequent calls.
+        if "id" in extra and "namespace" in extra:
+            self._client.set_namespace(extra["id"], extra["namespace"])
+
         super().__init__(
             name=name,
             image=image,
diff --git a/descarteslabs/core/compute/job.py b/descarteslabs/core/compute/job.py
index 94c80439..819631f7 100644
--- a/descarteslabs/core/compute/job.py
+++ b/descarteslabs/core/compute/job.py
@@ -17,7 +17,6 @@
 import warnings
 from datetime import datetime
 from typing import TYPE_CHECKING, Dict, List, Optional, Type
-
 from strenum import StrEnum
 
 from descarteslabs.exceptions import NotFoundError
@@ -207,12 +206,14 @@ def __eq__(self, other):
 
         return self.id == other.id
 
-    def _get_result_namespace(self):
+    def _get_result_namespace(self) -> str:
         """Returns the namespace for the Job result blob."""
-        auth = self._client.auth
-        namespace = f"{auth.namespace}"
-        if auth.payload["org"]:
-            namespace = f"{auth.payload['org']}:{namespace}"
+        namespace = self._client.get_namespace(self.function_id)
+
+        if not namespace:
+            # Fetching the function from the server will set the namespace
+            # during hydration in Function.__init__
+            namespace = self.function.namespace
 
         return namespace
 
@@ -372,10 +373,9 @@ def result(
             catalog_client = self._client.catalog_client
 
         try:
+            namespace = self._get_result_namespace()
             result = Blob.get_data(
-                name=f"{self.function_id}/{self.id}",
-                namespace=self._get_result_namespace(),
-                storage_type=StorageType.COMPUTE,
+                id=f"{StorageType.COMPUTE}/{namespace}/{self.function_id}/{self.id}",
                 client=catalog_client,
             )
         except NotFoundError: