openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-11-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+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-20T14:26:20+00:00
+#   timestamp: 2024-11-29T19:36:27+00:00
 
 from __future__ import annotations
 

metadata/ingestion/api/models.py

@@ -14,20 +14,24 @@ Generic models
 from typing import Generic, Optional, TypeVar
 
 from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
 
-# Entities are instances of BaseModel
 Entity = BaseModel
+
 T = TypeVar("T")
 
 
 class Either(BaseModel, Generic[T]):
     """Any execution should return us Either an Entity of an error for us to handle"""
 
-    left: Optional[StackTraceError] = Field(
-        None, description="Error encountered during execution"
-    )
-    right: Optional[T] = Field(None, description="Correct instance of an Entity")
+    left: Annotated[
+        Optional[StackTraceError],
+        Field(description="Error encountered during execution", default=None),
+    ]
+    right: Annotated[
+        Optional[T], Field(description="Correct instance of an Entity", default=None)
+    ]

metadata/ingestion/models/patch_request.py

@@ -94,6 +94,7 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "tableConstraints": True,
     "tablePartition": True,
     "location": True,
+    "locationPath": True,
     "schemaDefinition": True,
     "sampleData": True,
     "fileFormat": True,

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/ometa/models.py

@@ -30,3 +30,4 @@ class EntityList(BaseModel, Generic[T]):
     entities: List[T]
     total: int
     after: Optional[str] = None
+    before: Optional[str] = None

metadata/ingestion/ometa/ometa_api.py

@@ -378,12 +378,13 @@ class OpenMetadata(
         logger.debug("Cannot find the Entity %s", fqn)
         return None
 
-    # pylint: disable=too-many-locals
+    # pylint: disable=too-many-locals, too-many-arguments
     def list_entities(
         self,
         entity: Type[T],
         fields: Optional[List[str]] = None,
         after: Optional[str] = None,
+        before: Optional[str] = None,
         limit: int = 100,
         params: Optional[Dict[str, str]] = None,
         skip_on_failure: bool = False,
@@ -395,9 +396,10 @@ class OpenMetadata(
         suffix = self.get_suffix(entity)
         url_limit = f"?limit={limit}"
         url_after = f"&after={after}" if after else ""
+        url_before = f"&before={before}" if before else ""
         url_fields = f"&fields={','.join(fields)}" if fields else ""
         resp = self.client.get(
-            path=f"{suffix}{url_limit}{url_after}{url_fields}", data=params
+            path=f"{suffix}{url_limit}{url_after}{url_before}{url_fields}", data=params
         )
 
         if self._use_raw_data:
@@ -421,7 +423,8 @@ class OpenMetadata(
 
         total = resp["paging"]["total"]
         after = resp["paging"]["after"] if "after" in resp["paging"] else None
-        return EntityList(entities=entities, total=total, after=after)
+        before = resp["paging"]["before"] if "before" in resp["paging"] else None
+        return EntityList(entities=entities, total=total, after=after, before=before)
 
     def list_all_entities(
         self,

metadata/ingestion/sink/metadata_rest.py

@@ -84,6 +84,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardUsage
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.profiler.api.models import ProfilerResponse
+from metadata.sampler.models import SamplerResponse
 from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import get_log_name, ingestion_logger
 
@@ -549,22 +550,8 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         )
 
     @_run_dispatch.register
-    def write_profiler_response(self, record: ProfilerResponse) -> Either[Table]:
-        """Cleanup "`" character in columns and ingest"""
-        column_profile = record.profile.columnProfile
-        for column in column_profile:
-            column.name = column.name.replace("`", "")
-
-        record.profile.columnProfile = column_profile
-
-        table = self.metadata.ingest_profile_data(
-            table=record.table,
-            profile_request=record.profile,
-        )
-        logger.debug(
-            f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.root}"
-        )
-
+    def write_sampler_response(self, record: SamplerResponse) -> Either[Table]:
+        """Ingest the sample data - if needed - and the PII tags"""
         if record.sample_data and record.sample_data.store:
             table_data = self.metadata.ingest_table_sample_data(
                 table=record.table, sample_data=record.sample_data.data
@@ -572,7 +559,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             if not table_data:
                 self.status.failed(
                     StackTraceError(
-                        name=table.fullyQualifiedName.root,
+                        name=record.table.fullyQualifiedName.root,
                         error="Error trying to ingest sample data for table",
                     )
                 )
@@ -587,7 +574,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             )
             if not patched:
                 self.status.warning(
-                    key=table.fullyQualifiedName.root,
+                    key=record.table.fullyQualifiedName.root,
                     reason="Error patching tags for table",
                 )
             else:
@@ -595,6 +582,24 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
                     f"Successfully patched tag {record.column_tags} for {record.table.fullyQualifiedName.root}"
                 )
 
+        return Either(right=record.table)
+
+    @_run_dispatch.register
+    def write_profiler_response(self, record: ProfilerResponse) -> Either[Table]:
+        """Cleanup "`" character in columns and ingest"""
+        column_profile = record.profile.columnProfile
+        for column in column_profile:
+            column.name = column.name.replace("`", "")
+
+        record.profile.columnProfile = column_profile
+
+        table = self.metadata.ingest_profile_data(
+            table=record.table,
+            profile_request=record.profile,
+        )
+        logger.debug(
+            f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.root}"
+        )
         return Either(right=table)
 
     @_run_dispatch.register

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

@@ -0,0 +1,252 @@
+#  Copyright 2023 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+REST Auth & Client for MicroStrategy
+"""
+import traceback
+from typing import List, Optional
+
+import requests
+
+from metadata.generated.schema.entity.services.connections.dashboard.microStrategyConnection import (
+    MicroStrategyConnection,
+)
+from metadata.ingestion.connections.test_connections import SourceConnectionException
+from metadata.ingestion.ometa.client import REST, ClientConfig
+from metadata.ingestion.source.dashboard.microstrategy.models import (
+    AuthHeaderCookie,
+    MstrDashboard,
+    MstrDashboardDetails,
+    MstrDashboardList,
+    MstrProject,
+    MstrProjectList,
+    MstrSearchResult,
+    MstrSearchResultList,
+)
+from metadata.utils.helpers import clean_uri
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+API_VERSION = "MicroStrategyLibrary/api"
+APPLICATION_TYPE = 35
+
+
+class MicroStrategyClient:
+    """
+    Client Handling API communication with Metabase
+    """
+
+    def _get_base_url(self, path=None):
+        if not path:
+            return f"{clean_uri(str(self.config.hostPort))}/{API_VERSION}"
+        return f"{clean_uri(str(self.config.hostPort))}/{API_VERSION}/{path}"
+
+    def __init__(
+        self,
+        config: MicroStrategyConnection,
+    ):
+        self.config = config
+
+        self.auth_params: AuthHeaderCookie = self._get_auth_header_and_cookies()
+
+        client_config = ClientConfig(
+            base_url=clean_uri(str(self.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": self.config.loginMode,
+                "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)
+
+    def get_projects_list(self) -> List[MstrProject]:
+        """
+        Get List of all projects
+        """
+        try:
+            resp_projects = self.client.get(
+                path="/projects",
+            )
+
+            project_list = MstrProjectList(projects=resp_projects)
+            return project_list.projects
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to fetch the project list due to [{exc}]")
+
+        return []
+
+    def get_project_by_name(self) -> Optional[MstrProject]:
+        """
+        Get Project By Name
+        """
+        try:
+            resp_projects = self.client.get(
+                path=f"/projects/{self.config.projectName}",
+            )
+
+            project = MstrProject.model_validate(resp_projects)
+            return project
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning("Failed to fetch the project list")
+
+        return None
+
+    def get_search_results_list(
+        self, project_id, object_type
+    ) -> 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:
+            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,
+            )
+
+            results_list = MstrSearchResultList.model_validate(resp_results).result
+            return results_list
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning("Failed to fetch the Search Result list")
+
+        return []
+
+    def get_dashboards_list(self, project_id, project_name) -> List[MstrDashboard]:
+        """
+        Get Dashboard
+        """
+        try:
+            results = self.get_search_results_list(
+                project_id=project_id, object_type=55
+            )
+
+            dashboards = []
+            for result in results:
+                dashboards.append(
+                    MstrDashboard(projectName=project_name, **result.model_dump())
+                )
+
+            dashboards_list = MstrDashboardList(dashboards=dashboards)
+            return dashboards_list.dashboards
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning("Failed to fetch the dashboard list")
+
+        return []
+
+    def get_dashboard_details(
+        self, project_id, project_name, dashboard_id
+    ) -> Optional[MstrDashboardDetails]:
+        """
+        Get Dashboard Details
+        """
+        try:
+            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
+            )
+
+            return MstrDashboardDetails(
+                projectId=project_id, projectName=project_name, **resp_dashboard
+            )
+
+        except Exception:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to fetch the dashboard with id: {dashboard_id}")
+
+        return None

metadata/ingestion/source/dashboard/{mstr → microstrategy}/connection.py

@@ -17,31 +17,29 @@ from typing import Optional
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
-from metadata.generated.schema.entity.services.connections.dashboard.mstrConnection import (
-    MstrConnection,
+from metadata.generated.schema.entity.services.connections.dashboard.microStrategyConnection import (
+    MicroStrategyConnection,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.dashboard.mstr.client import MSTRClient
-from metadata.utils.constants import THREE_MIN
+from metadata.ingestion.source.dashboard.microstrategy.client import MicroStrategyClient
 
 
-def get_connection(connection: MstrConnection) -> MSTRClient:
+def get_connection(connection: MicroStrategyConnection) -> MicroStrategyClient:
     """
     Create connection
     """
-    return MSTRClient(connection)
+    return MicroStrategyClient(connection)
 
 
 def test_connection(
     metadata: OpenMetadata,
-    client: MSTRClient,
-    service_connection: MstrConnection,
+    client: MicroStrategyClient,
+    service_connection: MicroStrategyConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-    timeout_seconds: Optional[int] = THREE_MIN,
 ) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
@@ -55,5 +53,4 @@ def test_connection(
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
-        timeout_seconds=timeout_seconds,
     )