openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.5.12.0__py3-none-any.whl

metadata/generated/schema/type/customProperties/tableConfig.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/customProperties/tableConfig.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/customProperty.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/customProperty.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/dailyCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/dailyCount.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/databaseConnectionConfig.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/databaseConnectionConfig.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityHierarchy.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityHierarchy.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityHistory.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityHistory.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityLineage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityLineage.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityReference.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReference.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityReferenceList.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReferenceList.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityRelationship.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityRelationship.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/filterPattern.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/filterPattern.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/function.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/function.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/include.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/include.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/lifeCycle.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/lifeCycle.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-11-15T14:19:19+00:00
+#   timestamp: 2024-11-26T06:50:54+00:00
 
 from __future__ import annotations
 

metadata/ingestion/ometa/client.py

@@ -14,7 +14,7 @@ Python API REST wrapper and helpers
 import time
 import traceback
 from datetime import datetime, timezone
-from typing import Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import requests
 from requests.exceptions import HTTPError
@@ -115,6 +115,7 @@ class ClientConfig(ConfigModel):
     allow_redirects: Optional[bool] = False
     auth_token_mode: Optional[str] = "Bearer"
     verify: Optional[Union[bool, str]] = None
+    cookies: Optional[Any] = None
     ttl_cache: int = 60
 
 
@@ -138,10 +139,11 @@ class REST:
         self._auth_token = self.config.auth_token
         self._auth_token_mode = self.config.auth_token_mode
         self._verify = self.config.verify
+        self._cookies = self.config.cookies
 
         self._limits_reached = TTLCache(config.ttl_cache)
 
-    def _request(  # pylint: disable=too-many-arguments
+    def _request(  # pylint: disable=too-many-arguments,too-many-branches
         self,
         method,
         path,
@@ -160,10 +162,12 @@ class REST:
         base_url = base_url or self._base_url
         version = api_version if api_version else self._api_version
         url: URL = URL(base_url + "/" + version + path)
+        cookies = self._cookies
         if (
             self.config.expires_in
             and datetime.now(timezone.utc).timestamp() >= self.config.expires_in
             or not self.config.access_token
+            and self._auth_token
         ):
             self.config.access_token, expiry = self._auth_token()
             if not self.config.access_token == "no_token":
@@ -173,12 +177,12 @@ class REST:
                     self.config.expires_in = (
                         datetime.now(timezone.utc).timestamp() + expiry - 120
                     )
-
-        headers[self.config.auth_header] = (
-            f"{self._auth_token_mode} {self.config.access_token}"
-            if self._auth_token_mode
-            else self.config.access_token
-        )
+        if self.config.auth_header:
+            headers[self.config.auth_header] = (
+                f"{self._auth_token_mode} {self.config.access_token}"
+                if self._auth_token_mode
+                else self.config.access_token
+            )
 
         # Merge extra headers if provided.
         # If a header value is provided in modulo string format and matches an existing header,
@@ -198,6 +202,7 @@ class REST:
             # It's better to fail early if the URL isn't right.
             "allow_redirects": self.config.allow_redirects,
             "verify": self._verify,
+            "cookies": cookies,
         }
 
         method_key = "params" if method.upper() == "GET" else "data"

metadata/ingestion/source/dashboard/mstr/client.py

@@ -15,13 +15,14 @@ import traceback
 from typing import List, Optional
 
 import requests
-from mstr.requests import MSTRRESTSession
 
 from metadata.generated.schema.entity.services.connections.dashboard.mstrConnection import (
     MstrConnection,
 )
 from metadata.ingestion.connections.test_connections import SourceConnectionException
+from metadata.ingestion.ometa.client import REST, ClientConfig
 from metadata.ingestion.source.dashboard.mstr.models import (
+    AuthHeaderCookie,
     MstrDashboard,
     MstrDashboardDetails,
     MstrDashboardList,
@@ -30,11 +31,14 @@ from metadata.ingestion.source.dashboard.mstr.models import (
     MstrSearchResult,
     MstrSearchResultList,
 )
+from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 API_VERSION = "MicroStrategyLibrary/api"
+LOGIN_MODE_GUEST = 8
+APPLICATION_TYPE = 35
 
 
 class MSTRClient:
@@ -44,32 +48,90 @@ class MSTRClient:
 
     def _get_base_url(self, path=None):
         if not path:
-            return f"{self.config.hostPort}/{API_VERSION}/"
-        return f"{self.config.hostPort}/{API_VERSION}/{path}"
-
-    def _get_mstr_session(self) -> MSTRRESTSession:
-        try:
-            session = MSTRRESTSession(base_url=self._get_base_url())
-            session.login(
-                username=self.config.username,
-                password=self.config.password.get_secret_value(),
-            )
-            return session
-
-        except KeyError as exe:
-            msg = "Failed to fetch mstr session, please validate credentials"
-            raise SourceConnectionException(msg) from exe
-
-        except Exception as exc:
-            msg = f"Unknown error in connection: {exc}."
-            raise SourceConnectionException(msg) from exc
+            return f"{clean_uri(self.config.hostPort)}/{API_VERSION}"
+        return f"{clean_uri(self.config.hostPort)}/{API_VERSION}/{path}"
 
     def __init__(
         self,
         config: MstrConnection,
     ):
         self.config = config
-        self.session = self._get_mstr_session()
+
+        self.auth_params: AuthHeaderCookie = self._get_auth_header_and_cookies()
+
+        client_config = ClientConfig(
+            base_url=clean_uri(config.hostPort),
+            api_version=API_VERSION,
+            extra_headers=self.auth_params.auth_header,
+            allow_redirects=True,
+            cookies=self.auth_params.auth_cookies,
+        )
+
+        self.client = REST(client_config)
+        self._set_api_session()
+
+    def _get_auth_header_and_cookies(self) -> Optional[AuthHeaderCookie]:
+        """
+        Send a request to authenticate the user and get headers and
+
+        To know about the data params below please visit
+        https://demo.microstrategy.com/MicroStrategyLibrary/api-docs/index.html#/Authentication/postLogin
+        """
+        try:
+            data = {
+                "username": self.config.username,
+                "password": self.config.password.get_secret_value(),
+                "loginMode": LOGIN_MODE_GUEST,
+                "applicationType": APPLICATION_TYPE,
+            }
+            response = requests.post(
+                url=self._get_base_url("auth/login"), data=data, timeout=60
+            )
+            if not response:
+                raise SourceConnectionException()
+            return AuthHeaderCookie(
+                auth_header=response.headers, auth_cookies=response.cookies
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Failed to fetch the auth header and cookies due to [{exc}], please validate credentials"
+            )
+        return None
+
+    def _set_api_session(self) -> bool:
+        """
+        Set the user api session to active this will keep the connection alive
+        """
+        api_session = requests.put(
+            url=self._get_base_url("sessions"),
+            headers=self.auth_params.auth_header,
+            cookies=self.auth_params.auth_cookies,
+            timeout=60,
+        )
+        if api_session.ok:
+            logger.info(
+                f"Connection Successful User {self.config.username} is Authenticated"
+            )
+            return True
+        raise requests.ConnectionError(
+            "Connection Failed, Failed to set an api session, Please validate the credentials"
+        )
+
+    def close_api_session(self) -> None:
+        """
+        Closes the active api session
+        """
+        try:
+            close_api_session = self.client.post(
+                path="/auth/logout",
+            )
+            if close_api_session.ok:
+                logger.info("API Session Closed Successfully")
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to close the api sesison due to [{exc}]")
 
     def is_project_name(self) -> bool:
         return bool(self.config.projectName)
@@ -79,14 +141,11 @@ class MSTRClient:
         Get List of all projects
         """
         try:
-            resp_projects = self.session.get(
-                url=self._get_base_url("projects"), params={"include_auth": True}
+            resp_projects = self.client.get(
+                path="/projects",
             )
 
-            if not resp_projects.ok:
-                raise requests.ConnectionError()
-
-            project_list = MstrProjectList(projects=resp_projects.json())
+            project_list = MstrProjectList(projects=resp_projects)
             return project_list.projects
 
         except Exception as exc:
@@ -100,15 +159,11 @@ class MSTRClient:
         Get Project By Name
         """
         try:
-            resp_projects = self.session.get(
-                url=self._get_base_url(f"projects/{self.config.projectName}"),
-                params={"include_auth": True},
+            resp_projects = self.client.get(
+                path=f"/projects/{self.config.projectName}",
             )
 
-            if not resp_projects.ok:
-                raise requests.ConnectionError()
-
-            project = MstrProject(**resp_projects.json())
+            project = MstrProject.model_validate(resp_projects)
             return project
 
         except Exception:
@@ -122,32 +177,28 @@ class MSTRClient:
     ) -> List[MstrSearchResult]:
         """
         Get Search Results
+
+        To know about the data params below please visit
+        https://demo.microstrategy.com/MicroStrategyLibrary/api-docs/index.html?#/Browsing/doQuickSearch
         """
         try:
-            resp_results = self.session.get(
-                url=self._get_base_url("searches/results"),
-                params={
-                    "include_auth": True,
-                    "project_id": project_id,
-                    "type": object_type,
-                    "getAncestors": False,
-                    "offset": 0,
-                    "limit": -1,
-                    "certifiedStatus": "ALL",
-                    "isCrossCluster": False,
-                    "result.hidden": False,
-                },
+            data = {
+                "project_id": project_id,
+                "type": object_type,
+                "getAncestors": False,
+                "offset": 0,
+                "limit": -1,
+                "certifiedStatus": "ALL",
+                "isCrossCluster": False,
+                "result.hidden": False,
+            }
+            resp_results = self.client.get(
+                path="/searches/results",
+                data=data,
             )
 
-            if not resp_results.ok:
-                raise requests.ConnectionError()
-
-            results = []
-            for resp_result in resp_results.json()["result"]:
-                results.append(resp_result)
-
-            results_list = MstrSearchResultList(results=results)
-            return results_list.results
+            results_list = MstrSearchResultList.model_validate(resp_results).result
+            return results_list
 
         except Exception:
             logger.debug(traceback.format_exc())
@@ -186,19 +237,13 @@ class MSTRClient:
         Get Dashboard Details
         """
         try:
-            resp_dashboard = self.session.get(
-                url=self._get_base_url(f"v2/dossiers/{dashboard_id}/definition"),
-                params={
-                    "include_auth": True,
-                },
-                headers={"X-MSTR-ProjectID": project_id},
+            headers = {"X-MSTR-ProjectID": project_id} | self.auth_params.auth_header
+            resp_dashboard = self.client._request(  # pylint: disable=protected-access
+                "GET", path=f"/v2/dossiers/{dashboard_id}/definition", headers=headers
             )
 
-            if not resp_dashboard.ok:
-                raise requests.ConnectionError()
-
             return MstrDashboardDetails(
-                projectId=project_id, projectName=project_name, **resp_dashboard.json()
+                projectId=project_id, projectName=project_name, **resp_dashboard
             )
 
         except Exception: