update sdk client to accept headers

Author: Deepthi-Chand

dataspace_sdk/base.py

@@ -25,6 +25,7 @@ def __init__(self, base_url: str, auth_client: Any = None):
         """
         self.base_url = base_url.rstrip("/")
         self.auth_client = auth_client
+        self.default_headers: Dict[str, str] = {}
 
     def _get_headers(self, additional_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
         """
@@ -38,6 +39,9 @@ def _get_headers(self, additional_headers: Optional[Dict[str, str]] = None) -> D
         """
         headers = {"Content-Type": "application/json"}
 
+        if self.default_headers:
+            headers.update(self.default_headers)
+
         if self.auth_client and self.auth_client.is_authenticated():
             headers["Authorization"] = f"Bearer {self.auth_client.access_token}"
 

dataspace_sdk/client.py

@@ -187,6 +187,22 @@ def is_authenticated(self) -> bool:
         """
         return self._auth.is_authenticated()
 
+    def set_organization(self, organization_id: str) -> None:
+        """
+        Set the organization header for all subsequent API requests.
+
+        The DataSpace backend reads the 'organization' header to scope
+        queries to a specific organization.
+
+        Args:
+            organization_id: Organization ID to include in requests
+        """
+        self.datasets.default_headers["organization"] = organization_id
+        self.aimodels.default_headers["organization"] = organization_id
+        self.usecases.default_headers["organization"] = organization_id
+        self.sectors.default_headers["organization"] = organization_id
+        self.auditors.default_headers["organization"] = organization_id
+
     @property
     def user(self) -> Optional[dict]:
         """

dataspace_sdk/resources/auditors.py

@@ -25,10 +25,13 @@ def __init__(self, base_url: str, auth_client: Any):
         """
         self._base_url = base_url.rstrip("/")
         self._auth = auth_client
+        self.default_headers: Dict[str, str] = {}
 
     def _get_headers(self) -> Dict[str, str]:
         """Get request headers with authentication."""
         headers = {"Content-Type": "application/json"}
+        if self.default_headers:
+            headers.update(self.default_headers)
         if self._auth and self._auth.access_token:
             headers["Authorization"] = f"Bearer {self._auth.access_token}"
         return headers