acryl-datahub 1.0.0rc6__py3-none-any.whl → 1.0.0rc7__py3-none-any.whl

datahub/ingestion/source/metabase.py

@@ -69,9 +69,19 @@ class MetabaseConfig(DatasetLineageProviderConfigBase, StatefulIngestionConfigBa
         default=None,
         description="optional URL to use in links (if `connect_uri` is only for ingestion)",
     )
-    username: Optional[str] = Field(default=None, description="Metabase username.")
+    username: Optional[str] = Field(
+        default=None,
+        description="Metabase username, used when an API key is not provided.",
+    )
     password: Optional[pydantic.SecretStr] = Field(
-        default=None, description="Metabase password."
+        default=None,
+        description="Metabase password, used when an API key is not provided.",
+    )
+
+    # https://www.metabase.com/learn/metabase-basics/administration/administration-and-operation/metabase-api#example-get-request
+    api_key: Optional[pydantic.SecretStr] = Field(
+        default=None,
+        description="Metabase API key. If provided, the username and password will be ignored. Recommended method.",
     )
     # TODO: Check and remove this if no longer needed.
     # Config database_alias is removed from sql sources.
@@ -178,30 +188,40 @@ class MetabaseSource(StatefulIngestionSourceBase):
         self.source_config: MetabaseConfig = config
 
     def setup_session(self) -> None:
-        login_response = requests.post(
-            f"{self.config.connect_uri}/api/session",
-            None,
-            {
-                "username": self.config.username,
-                "password": (
-                    self.config.password.get_secret_value()
-                    if self.config.password
-                    else None
-                ),
-            },
-        )
+        self.session = requests.session()
+        if self.config.api_key:
+            self.session.headers.update(
+                {
+                    "x-api-key": self.config.api_key.get_secret_value(),
+                    "Content-Type": "application/json",
+                    "Accept": "*/*",
+                }
+            )
+        else:
+            # If no API key is provided, generate a session token using username and password.
+            login_response = requests.post(
+                f"{self.config.connect_uri}/api/session",
+                None,
+                {
+                    "username": self.config.username,
+                    "password": (
+                        self.config.password.get_secret_value()
+                        if self.config.password
+                        else None
+                    ),
+                },
+            )
 
-        login_response.raise_for_status()
-        self.access_token = login_response.json().get("id", "")
+            login_response.raise_for_status()
+            self.access_token = login_response.json().get("id", "")
 
-        self.session = requests.session()
-        self.session.headers.update(
-            {
-                "X-Metabase-Session": f"{self.access_token}",
-                "Content-Type": "application/json",
-                "Accept": "*/*",
-            }
-        )
+            self.session.headers.update(
+                {
+                    "X-Metabase-Session": f"{self.access_token}",
+                    "Content-Type": "application/json",
+                    "Accept": "*/*",
+                }
+            )
 
         # Test the connection
         try:
@@ -217,15 +237,17 @@ class MetabaseSource(StatefulIngestionSourceBase):
             )
 
     def close(self) -> None:
-        response = requests.delete(
-            f"{self.config.connect_uri}/api/session",
-            headers={"X-Metabase-Session": self.access_token},
-        )
-        if response.status_code not in (200, 204):
-            self.report.report_failure(
-                title="Unable to Log User Out",
-                message=f"Unable to logout for user {self.config.username}",
+        # API key authentication does not require session closure.
+        if not self.config.api_key:
+            response = requests.delete(
+                f"{self.config.connect_uri}/api/session",
+                headers={"X-Metabase-Session": self.access_token},
             )
+            if response.status_code not in (200, 204):
+                self.report.report_failure(
+                    title="Unable to Log User Out",
+                    message=f"Unable to logout for user {self.config.username}",
+                )
         super().close()
 
     def emit_dashboard_mces(self) -> Iterable[MetadataWorkUnit]:

datahub/ingestion/source/mode.py

@@ -1494,7 +1494,7 @@ class ModeSource(StatefulIngestionSourceBase):
                     sleep_time = error_response.headers.get("retry-after")
                     if sleep_time is not None:
                         time.sleep(float(sleep_time))
-                    raise HTTPError429
+                    raise HTTPError429 from None
 
                 raise http_error
 

datahub/ingestion/source/pulsar.py

@@ -230,8 +230,8 @@ class PulsarSource(StatefulIngestionSourceBase):
                 self.report.report_warning("HTTPError", message)
         except requests.exceptions.RequestException as e:
             raise Exception(
-                f"An ambiguous exception occurred while handling the request: {e}"
-            )
+                "An ambiguous exception occurred while handling the request"
+            ) from e
 
     @classmethod
     def create(cls, config_dict, ctx):

datahub/ingestion/source/sigma/sigma.py

@@ -124,7 +124,7 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
         try:
             self.sigma_api = SigmaAPI(self.config, self.reporter)
         except Exception as e:
-            raise ConfigurationError(f"Unable to connect sigma API. Exception: {e}")
+            raise ConfigurationError("Unable to connect sigma API") from e
 
     @staticmethod
     def test_connection(config_dict: dict) -> TestConnectionReport:

datahub/ingestion/source/snowflake/snowflake_connection.py

@@ -312,7 +312,7 @@ class SnowflakeConnectionConfig(ConfigModel):
             raise ValueError(
                 f"access_token not found in response {response}. "
                 "Please check your OAuth configuration."
-            )
+            ) from None
         connect_args = self.get_options()["connect_args"]
         return snowflake.connector.connect(
             user=self.username,

datahub/ingestion/source/sql/druid.py

@@ -50,11 +50,7 @@ class DruidConfig(BasicSQLAlchemyConfig):
     """
 
     def get_identifier(self, schema: str, table: str) -> str:
-        return (
-            f"{self.platform_instance}.{table}"
-            if self.platform_instance
-            else f"{table}"
-        )
+        return f"{table}"
 
 
 @platform_name("Druid")

datahub/ingestion/source/tableau/tableau.py

@@ -1562,8 +1562,9 @@ class TableauSiteSource:
         query: str,
         connection_type: str,
         page_size: int,
-        query_filter: dict = {},
+        query_filter: Optional[dict] = None,
     ) -> Iterable[dict]:
+        query_filter = query_filter or {}
         query_filter = optimize_query_filter(query_filter)
 
         # Calls the get_connection_object_page function to get the objects,

datahub/ingestion/source/tableau/tableau_common.py

@@ -514,7 +514,8 @@ FIELD_TYPE_MAPPING = {
 }
 
 
-def get_tags_from_params(params: List[str] = []) -> GlobalTagsClass:
+def get_tags_from_params(params: Optional[List[str]] = None) -> GlobalTagsClass:
+    params = params or []
     tags = [
         TagAssociationClass(tag=builder.make_tag_urn(tag.upper()))
         for tag in params

datahub/lite/duckdb_lite.py

@@ -284,9 +284,10 @@ class DuckDBLite(DataHubLiteLocal[DuckDBLiteConfig]):
         self,
         query: str,
         flavor: SearchFlavor,
-        aspects: List[str] = [],
+        aspects: Optional[List[str]] = None,
         snippet: bool = True,
     ) -> Iterable[Searchable]:
+        aspects = aspects or []
         if flavor == SearchFlavor.FREE_TEXT:
             base_query = f"SELECT distinct(urn), 'urn', NULL from metadata_aspect_v2 where urn ILIKE '%{query}%' UNION SELECT urn, aspect_name, metadata from metadata_aspect_v2 where metadata->>'$.name' ILIKE '%{query}%'"
             for r in self.duckdb_client.execute(base_query).fetchall():

datahub/lite/lite_local.py

@@ -90,7 +90,7 @@ class DataHubLiteLocal(Generic[LiteConfig], Closeable, metaclass=ABCMeta):
         self,
         query: str,
         flavor: SearchFlavor,
-        aspects: List[str] = [],
+        aspects: Optional[List[str]] = None,
         snippet: bool = True,
     ) -> Iterable[Searchable]:
         pass

datahub/lite/lite_util.py

@@ -70,9 +70,10 @@ class DataHubLiteWrapper(DataHubLiteLocal):
         self,
         query: str,
         flavor: SearchFlavor,
-        aspects: List[str] = [],
+        aspects: Optional[List[str]] = None,
         snippet: bool = True,
     ) -> Iterable[Searchable]:
+        aspects = aspects or []
         yield from self.lite.search(query, flavor, aspects, snippet)
 
     def ls(self, path: str) -> List[Browseable]:
@@ -96,10 +97,10 @@ def get_datahub_lite(config_dict: dict, read_only: bool = False) -> "DataHubLite
     lite_type = lite_local_config.type
     try:
         lite_class = lite_registry.get(lite_type)
-    except KeyError:
+    except KeyError as e:
         raise Exception(
             f"Failed to find a registered lite implementation for {lite_type}. Valid values are {[k for k in lite_registry.mapping.keys()]}"
-        )
+        ) from e
 
     lite_specific_config = lite_class.get_config_class().parse_obj(
         lite_local_config.config

datahub/utilities/memory_footprint.py

@@ -1,10 +1,10 @@
 from collections import deque
 from itertools import chain
 from sys import getsizeof
-from typing import Any, Iterator
+from typing import Any, Iterator, Optional
 
 
-def total_size(o: Any, handlers: Any = {}) -> int:
+def total_size(o: Any, handlers: Optional[Any] = None) -> int:
     """Returns the approximate memory footprint an object and all of its contents.
     Automatically finds the contents of the following builtin containers and
     their subclasses:  tuple, list, deque, dict, set and frozenset.
@@ -14,6 +14,7 @@ def total_size(o: Any, handlers: Any = {}) -> int:
 
     Based on https://github.com/ActiveState/recipe-577504-compute-mem-footprint/blob/master/recipe.py
     """
+    handlers = handlers or {}
 
     def dict_handler(d: dict) -> Iterator[Any]:
         return chain.from_iterable(d.items())