Merge pull request #712 from hkad98/ds-credentials

Inject data source credentials from config

Author: hkad98

gooddata-sdk/gooddata_sdk/catalog/data_source/declarative_model/data_source.py

@@ -2,7 +2,8 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, List, Optional, Type
+from typing import Any, List, Optional, Type, Union
+from warnings import warn
 
 import attr
 from gooddata_api_client.model.declarative_data_source import DeclarativeDataSource
@@ -13,7 +14,7 @@
 from gooddata_sdk.catalog.entity import TokenCredentialsFromFile
 from gooddata_sdk.catalog.parameter import CatalogParameter
 from gooddata_sdk.catalog.permission.declarative_model.permission import CatalogDeclarativeDataSourcePermission
-from gooddata_sdk.utils import create_directory, read_layout_from_file, write_layout_to_file
+from gooddata_sdk.utils import create_directory, get_ds_credentials, read_layout_from_file, write_layout_to_file
 
 BIGQUERY_TYPE = "BIGQUERY"
 LAYOUT_DATA_SOURCES_DIR = "data_sources"
@@ -23,10 +24,9 @@
 class CatalogDeclarativeDataSources(Base):
     data_sources: List[CatalogDeclarativeDataSource]
 
-    def to_api(self, credentials: Optional[dict[str, Any]] = None) -> DeclarativeDataSources:
+    def _inject_base(self, credentials: dict[str, Any]) -> DeclarativeDataSources:
         data_sources = []
         client_class = self.client_class()
-        credentials = credentials if credentials is not None else dict()
         for data_source in self.data_sources:
             if data_source.id in credentials:
                 if data_source.type == BIGQUERY_TYPE:
@@ -38,6 +38,34 @@ def to_api(self, credentials: Optional[dict[str, Any]] = None) -> DeclarativeDat
                 data_sources.append(data_source.to_api())
         return client_class(data_sources=data_sources)
 
+    def _inject_credentials_legacy(self, credentials: dict[str, Any]) -> DeclarativeDataSources:
+        return self._inject_base(credentials)
+
+    def _inject_credentials_aac(self, config_file: Union[str, Path]) -> DeclarativeDataSources:
+        ds_ids = {ds.id for ds in self.data_sources}
+        credentials = get_ds_credentials(config_file)
+        missing = set(credentials.keys()).difference(ds_ids)
+        if len(missing) > 0:
+            warn(
+                f"The following data sources are missing credentials: {missing}.",
+                UserWarning,
+                stacklevel=2,
+            )
+        return self._inject_base(credentials)
+
+    def to_api(
+        self, credentials: Optional[dict[str, Any]] = None, config_file: Optional[Union[str, Path]] = None
+    ) -> DeclarativeDataSources:
+        client_class = self.client_class()
+        if credentials is not None and config_file is not None:
+            raise ValueError("Only one of credentials or config_file should be provided")
+        if credentials is None and config_file is None:
+            return client_class(data_sources=[data_source.to_api() for data_source in self.data_sources])
+        if credentials is not None:
+            return self._inject_credentials_legacy(credentials)
+        if config_file is not None:
+            return self._inject_credentials_aac(config_file)
+
     @staticmethod
     def client_class() -> Type[DeclarativeDataSources]:
         return DeclarativeDataSources
@@ -101,9 +129,7 @@ def data_source_folder(data_sources_folder: Path, data_source_id: str) -> Path:
         create_directory(data_source_folder)
         return data_source_folder
 
-    def to_api(
-        self, password: Optional[str] = None, token: Optional[str] = None, include_nested_structures: bool = True
-    ) -> DeclarativeDataSource:
+    def to_api(self, password: Optional[str] = None, token: Optional[str] = None) -> DeclarativeDataSource:
         dictionary = self._get_snake_dict()
         if password is not None:
             dictionary["password"] = password
@@ -114,7 +140,7 @@ def to_api(
     def store_to_disk(self, data_sources_folder: Path) -> None:
         data_source_folder = self.data_source_folder(data_sources_folder, self.id)
         file_path = data_source_folder / f"{self.id}.yaml"
-        data_source_dict = self.to_api(include_nested_structures=False).to_dict(camel_case=True)
+        data_source_dict = self.to_api().to_dict(camel_case=True)
 
         write_layout_to_file(file_path, data_source_dict)
 

gooddata-sdk/gooddata_sdk/catalog/data_source/service.py

@@ -3,7 +3,7 @@
 
 import functools
 from pathlib import Path
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
 from gooddata_api_client.exceptions import NotFoundException
 
@@ -29,7 +29,7 @@
 from gooddata_sdk.catalog.entity import TokenCredentialsFromFile
 from gooddata_sdk.catalog.workspace.declarative_model.workspace.logical_model.ldm import CatalogDeclarativeModel
 from gooddata_sdk.client import GoodDataApiClient
-from gooddata_sdk.utils import load_all_entities_dict, read_layout_from_file
+from gooddata_sdk.utils import get_ds_credentials, load_all_entities_dict, read_layout_from_file
 
 _PDM_DEPRECATION_MSG = "This method is going to be deprecated due to PDM removal."
 
@@ -173,6 +173,7 @@ def put_declarative_data_sources(
         self,
         declarative_data_sources: CatalogDeclarativeDataSources,
         credentials_path: Optional[Path] = None,
+        config_file: Optional[Union[str, Path]] = None,
         test_data_sources: bool = False,
     ) -> None:
         """Set all data sources, including their related physical data model.
@@ -182,16 +183,18 @@ def put_declarative_data_sources(
                 Declarative Data Source object. Can be retrieved by get_declarative_data_sources.
             credentials_path (Optional[Path], optional):
                 Path to the Credentials. Optional, defaults to None.
+            config_file (Optional[Union[str, Path]], optional):
+                Path to the config file. Defaults to None.
             test_data_sources (bool, optional):
                 If True, the connection of data sources is tested. Defaults to False.
 
         Returns:
             None
         """
         if test_data_sources:
-            self.test_data_sources_connection(declarative_data_sources, credentials_path)
+            self.test_data_sources_connection(declarative_data_sources, credentials_path, config_file)
         credentials = self._credentials_from_file(credentials_path) if credentials_path is not None else dict()
-        self._layout_api.put_data_sources_layout(declarative_data_sources.to_api(credentials))
+        self._layout_api.put_data_sources_layout(declarative_data_sources.to_api(credentials, config_file))
 
     def store_declarative_data_sources(self, layout_root_path: Path = Path.cwd()) -> None:
         """Store data sources layouts in a directory hierarchy.
@@ -236,6 +239,7 @@ def load_and_put_declarative_data_sources(
         self,
         layout_root_path: Path = Path.cwd(),
         credentials_path: Optional[Path] = None,
+        config_file: Optional[Union[str, Path]] = None,
         test_data_sources: bool = False,
     ) -> None:
         """Loads and sets layouts stored using `store_declarative_data_sources`.
@@ -247,15 +251,17 @@ def load_and_put_declarative_data_sources(
             layout_root_path (Optional[Path], optional):
                 Path to the root of the layout directory. Defaults to Path.cwd().
             credentials_path (Optional[Path], optional):
-                Path to the credentials. Defaults to Path.cwd().
+                Path to the credentials.
+            config_file (Optional[Union[str, Path]], optional):
+                Path to the config file.
             test_data_sources (bool, optional):
                 If True, the connection of data sources is tested. Defaults to False.
 
         Returns:
             None
         """
         data_sources = self.load_declarative_data_sources(layout_root_path)
-        self.put_declarative_data_sources(data_sources, credentials_path, test_data_sources)
+        self.put_declarative_data_sources(data_sources, credentials_path, config_file, test_data_sources)
 
     @staticmethod
     def store_pdm_to_disk(pdm: CatalogDeclarativeTables, path: Path = Path.cwd()) -> None:
@@ -438,7 +444,10 @@ def scan_sql(self, data_source_id: str, sql_request: ScanSqlRequest) -> ScanSqlR
         return ScanSqlResponse.from_api(self._actions_api.scan_sql(data_source_id, sql_request.to_api()))
 
     def test_data_sources_connection(
-        self, declarative_data_sources: CatalogDeclarativeDataSources, credentials_path: Optional[Path] = None
+        self,
+        declarative_data_sources: CatalogDeclarativeDataSources,
+        credentials_path: Optional[Path] = None,
+        config_file: Optional[Union[str, Path]] = None,
     ) -> None:
         """Tests connection to declarative data sources.
 
@@ -453,6 +462,8 @@ def test_data_sources_connection(
                 Declarative Data Sources object
             credentials_path (Optional[Path], optional):
                 Path to the credentials. Defaults to None.
+            config_file (Optional[Union[str, Path]], optional):
+                Path to the config file. Defaults to None.
 
         Raises:
             ValueError:
@@ -461,7 +472,15 @@ def test_data_sources_connection(
         Returns:
             None
         """
-        credentials = self._credentials_from_file(credentials_path) if credentials_path is not None else dict()
+
+        credentials = dict()
+        if credentials_path is not None and config_file is not None:
+            raise ValueError("Only one of credentials or config_file should be provided")
+        if credentials_path is not None:
+            credentials = self._credentials_from_file(credentials_path)
+        if config_file is not None:
+            credentials = get_ds_credentials(config_file)
+
         errors: dict[str, str] = dict()
         for declarative_data_source in declarative_data_sources.data_sources:
             if credentials.get(declarative_data_source.id) is not None:

gooddata-sdk/gooddata_sdk/catalog/entity.py

@@ -4,7 +4,7 @@
 import base64
 import os
 from pathlib import Path
-from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar
+from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Union
 
 import attr
 
@@ -177,7 +177,7 @@ def from_api(cls, entity: dict[str, Any]) -> TokenCredentialsFromFile:
         raise NotImplementedError
 
     @staticmethod
-    def token_from_file(file_path: Path) -> str:
+    def token_from_file(file_path: Union[str, Path]) -> str:
         with open(file_path, "rb") as fp:
             return base64.b64encode(fp.read()).decode("utf-8")
 

gooddata-sdk/gooddata_sdk/config.py

@@ -51,3 +51,7 @@ class AacConfig(ConfigBase):
     profiles: Dict[str, Profile]
     default_profile: str
     access: Dict[str, str]
+
+    def ds_credentials(self) -> Dict[str, str]:
+        load_dotenv()
+        return {k: os.environ.get(v[1:], v) for k, v in self.access.items()}

gooddata-sdk/gooddata_sdk/utils.py

@@ -282,6 +282,16 @@ def _get_profile(profile: str, content: Dict) -> Dict[str, Any]:
         return _create_profile_legacy(content[profile])
 
 
+def _get_config_content(path: Union[str, Path]) -> Any:
+    path = Path(path) if isinstance(path, str) else path
+    if not path.exists():
+        raise ValueError(f"The file does not exist on path {path}.")
+    content = read_layout_from_file(path)
+    if content is None:
+        raise ValueError(f"The file is empty {path}.")
+    return content
+
+
 def profile_content(profile: str = "default", profiles_path: Path = PROFILES_FILE_PATH) -> dict[str, Any]:
     """Get the profile content from a given file.
 
@@ -301,14 +311,16 @@ def profile_content(profile: str = "default", profiles_path: Path = PROFILES_FIL
         dict[str, Any]:
             Profile content as a dictionary.
     """
-    if not profiles_path.exists():
-        raise ValueError(f"There is no profiles file located for path {profiles_path}.")
-    content = read_layout_from_file(profiles_path)
-    if content is None:
-        raise ValueError(f"The config file is empty {profiles_path}.")
+    content = _get_config_content(profiles_path)
     return _get_profile(profile, content)
 
 
+def get_ds_credentials(config_file: Union[str, Path]) -> Dict[str, str]:
+    content = _get_config_content(config_file)
+    config = AacConfig.from_dict(content)
+    return config.ds_credentials()
+
+
 def good_pandas_profile_content(
     profile: str = "default", profiles_path: Path = PROFILES_FILE_PATH
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]: