feat: add AuthTarget to AuthData (#459)

datalens-tech · May 28, 2024 · 24a595d · 24a595d
1 parent 5543f94
commit 24a595d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/lib/dl_api_commons/dl_api_commons/base_models.py b/lib/dl_api_commons/dl_api_commons/base_models.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import enum
 from typing import (
     Any,
     Optional,
@@ -37,13 +38,17 @@ def get_reporting_extra(self) -> dict[str, str | None]:
         )
 
 
+class AuthTarget(str, enum.Enum):
+    UNITED_STORAGE = "united_storage"
+
+
 class AuthData(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def get_headers(self) -> dict[DLHeaders, str]:
+    def get_headers(self, target: Optional[AuthTarget] = None) -> dict[DLHeaders, str]:
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def get_cookies(self) -> dict[DLCookies, str]:
+    def get_cookies(self, target: Optional[AuthTarget] = None) -> dict[DLCookies, str]:
         raise NotImplementedError()
 
 
@@ -169,8 +174,8 @@ def get_tenant_id(self) -> str:
 
 @attr.s()
 class NoAuthData(AuthData):
-    def get_headers(self) -> dict[DLHeaders, str]:
+    def get_headers(self, target: Optional[AuthTarget] = None) -> dict[DLHeaders, str]:
         return {}
 
-    def get_cookies(self) -> dict[DLCookies, str]:
+    def get_cookies(self, target: Optional[AuthTarget] = None) -> dict[DLCookies, str]:
         return {}
diff --git a/lib/dl_core/dl_core/united_storage_client.py b/lib/dl_core/dl_core/united_storage_client.py
@@ -24,6 +24,7 @@
 
 from dl_api_commons.base_models import (
     AuthData,
+    AuthTarget,
     TenantCommon,
     TenantDef,
 )
@@ -104,12 +105,12 @@ def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders,
 
         return {
             **(self.tenant.get_outbound_tenancy_headers() if include_tenancy else {}),
-            **self.auth_data.get_headers(),
+            **self.auth_data.get_headers(AuthTarget.UNITED_STORAGE),
             **{name: val for name, val in flags.items() if val is not None},
         }
 
     def get_outbound_cookies(self) -> dict[DLCookies, str]:
-        return self.auth_data.get_cookies()
+        return self.auth_data.get_cookies(AuthTarget.UNITED_STORAGE)
 
 
 @attr.s(frozen=True)
@@ -176,11 +177,11 @@ def get_tenant(self) -> TenantDef:
     def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
         return {
             **(self.tenant.get_outbound_tenancy_headers() if include_tenancy else {}),
-            **self.auth_data.get_headers(),
+            **self.auth_data.get_headers(AuthTarget.UNITED_STORAGE),
         }
 
     def get_outbound_cookies(self) -> dict[DLCookies, str]:
-        return self.auth_data.get_cookies()
+        return self.auth_data.get_cookies(AuthTarget.UNITED_STORAGE)
 
 
 # noinspection PyCallByClass