castor-extractor 0.20.0__py3-none-any.whl → 0.20.4__py3-none-any.whl

castor_extractor/visualization/powerbi/client/credentials_test.py

@@ -1,5 +1,4 @@
-from .constants import Urls
-from .credentials import PowerbiCredentials
+from .credentials import DEFAULT_SCOPE, PowerbiCredentials
 
 
 def test_credentials():
@@ -13,7 +12,7 @@ def test_credentials():
         client_id=client_id,
         secret=secret,
     )
-    assert credentials.scopes == [Urls.DEFAULT_SCOPE]
+    assert credentials.scopes == [DEFAULT_SCOPE]
 
     credentials = PowerbiCredentials(
         tenant_id=tenant_id,
@@ -21,7 +20,7 @@ def test_credentials():
         secret=secret,
         scopes=None,
     )
-    assert credentials.scopes == [Urls.DEFAULT_SCOPE]
+    assert credentials.scopes == [DEFAULT_SCOPE]
 
     # empty scopes
     credentials = PowerbiCredentials(

castor_extractor/visualization/powerbi/client/endpoints.py

@@ -0,0 +1,65 @@
+from datetime import date, datetime
+from typing import Optional, Tuple
+
+from ....utils import at_midnight, format_date, yesterday
+
+_CLIENT_APP_BASE = "https://login.microsoftonline.com"
+_REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
+
+
+def _time_filter(day: Optional[date]) -> Tuple[datetime, datetime]:
+    target_day = day or yesterday()
+    start = at_midnight(target_day)
+    end = datetime.combine(target_day, datetime.max.time())
+    return start, end
+
+
+class PowerBiEndpointFactory:
+    @classmethod
+    def activity_events(cls, day: Optional[date]) -> str:
+        start, end = _time_filter(day)
+        url = f"{_REST_API_BASE_PATH}/admin/activityevents"
+        url += "?$filter=Activity eq 'viewreport'"
+        url += f"&startDateTime='{format_date(start)}'"
+        url += f"&endDateTime='{format_date(end)}'"
+        return url
+
+    @classmethod
+    def authority(cls, tenant_id: str) -> str:
+        return f"{_CLIENT_APP_BASE}/{tenant_id}"
+
+    @classmethod
+    def dashboards(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/dashboards"
+
+    @classmethod
+    def datasets(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/datasets"
+
+    @classmethod
+    def groups(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/groups"
+
+    @classmethod
+    def metadata_create_scan(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/getInfo"
+
+    @classmethod
+    def metadata_scan_result(cls, scan_id: int) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/scanResult/{scan_id}"
+
+    @classmethod
+    def metadata_scan_status(cls, scan_id: int) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/scanStatus/{scan_id}"
+
+    @classmethod
+    def pages(cls, report_id: str) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/reports/{report_id}/pages"
+
+    @classmethod
+    def reports(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/reports"
+
+    @classmethod
+    def workspace_ids(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/modified"

castor_extractor/visualization/powerbi/client/pagination.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from pydantic import ConfigDict
+from pydantic.alias_generators import to_camel
+
+from ....utils import (
+    FetchNextPageBy,
+    PaginationModel,
+)
+
+
+class PowerBiPagination(PaginationModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+    fetch_by: FetchNextPageBy = FetchNextPageBy.URL
+
+    activity_event_entities: list
+    continuation_uri: Optional[str] = None
+    last_result_set: bool = False
+
+    def is_last(self) -> bool:
+        return self.last_result_set
+
+    def next_page_payload(self) -> Optional[str]:
+        return self.continuation_uri
+
+    def page_results(self) -> list:
+        return self.activity_event_entities

castor_extractor/visualization/powerbi/extract.py

@@ -1,4 +1,5 @@
-from typing import Iterable, List, Tuple, Union
+import logging
+from typing import Iterable, Tuple, Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -9,19 +10,23 @@ from ...utils import (
     write_json,
     write_summary,
 )
-from .assets import METADATA_ASSETS, PowerBiAsset
-from .client import Client, PowerbiCredentials
+from .assets import METADATA_ASSETS, REPORTS_ASSETS, PowerBiAsset
+from .client import PowerbiClient, PowerbiCredentials
+
+logger = logging.getLogger(__name__)
 
 
 def iterate_all_data(
-    client: Client,
-) -> Iterable[Tuple[PowerBiAsset, Union[List, dict]]]:
+    client: PowerbiClient,
+) -> Iterable[Tuple[PowerBiAsset, Union[list, dict]]]:
     for asset in PowerBiAsset:
-        if asset in METADATA_ASSETS:
+        if asset in METADATA_ASSETS + REPORTS_ASSETS:
             continue
 
-        data = client.fetch(asset)
-        yield asset, deep_serialize(data)
+        logger.info(f"Extracting {asset.name} from API")
+        data = list(deep_serialize(client.fetch(asset)))
+        yield asset, data
+        logger.info(f"Extracted {len(data)} {asset.name} from API")
 
 
 def extract_all(**kwargs) -> None:
@@ -31,7 +36,7 @@ def extract_all(**kwargs) -> None:
     """
     _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
     creds = PowerbiCredentials(**kwargs)
-    client = Client(creds)
+    client = PowerbiClient(creds)
     ts = current_timestamp()
 
     for key, data in iterate_all_data(client):

castor_extractor/visualization/thoughtspot/__init__.py

@@ -0,0 +1,3 @@
+from .assets import ThoughtspotAsset
+from .client import ThoughtspotClient, ThoughtspotCredentials
+from .extract import extract_all, iterate_all_data

castor_extractor/visualization/thoughtspot/assets.py

@@ -0,0 +1,9 @@
+from ...types import ExternalAsset
+
+
+class ThoughtspotAsset(ExternalAsset):
+    """Thoughtspot assets"""
+
+    LIVEBOARDS = "liveboards"
+    LOGICAL_TABLES = "logical_tables"
+    USAGES = "usages"

castor_extractor/visualization/thoughtspot/client/__init__.py

@@ -0,0 +1,2 @@
+from .client import ThoughtspotClient
+from .credentials import ThoughtspotCredentials

castor_extractor/visualization/thoughtspot/client/client.py

@@ -0,0 +1,120 @@
+from typing import Dict, Iterator, Optional
+
+import requests
+
+from ....utils import (
+    APIClient,
+    BearerAuth,
+    RequestSafeMode,
+    build_url,
+    handle_response,
+)
+from ..assets import (
+    ThoughtspotAsset,
+)
+from .credentials import (
+    ThoughtspotCredentials,
+)
+from .endpoints import (
+    ThoughtspotEndpointFactory,
+)
+from .utils import (
+    usage_liveboard_reader,
+)
+
+_AUTH_TIMEOUT_S = 60
+_THOUGHTSPOT_HEADERS = {
+    "X-Requested-By": "ThoughtSpot",
+    "Accept": "application/json",
+    "Content-Type": "application/json",
+}
+_METADATA_BATCH_SIZE = 100
+_USAGE_LIVEBOARD_ID = "bea79810-145f-4ad0-a02c-4177a6e7d861"
+# By default, no errors are ignored for the moment
+THOUGHTSPOT_SAFE_MODE = RequestSafeMode()
+
+
+class ThoughtspotBearerAuth(BearerAuth):
+    def __init__(self, host: str, token_payload: Dict[str, str]):
+        auth_endpoint = ThoughtspotEndpointFactory.authentication()
+        self.authentication_url = build_url(host, auth_endpoint)
+        self.token_payload = token_payload
+
+    def fetch_token(self):
+        token_api_path = self.authentication_url
+        token_response = requests.post(
+            token_api_path, data=self.token_payload, timeout=_AUTH_TIMEOUT_S
+        )
+        return handle_response(token_response)["token"]
+
+
+class ThoughtspotClient(APIClient):
+    def __init__(
+        self,
+        credentials: ThoughtspotCredentials,
+        safe_mode: Optional[RequestSafeMode] = None,
+    ):
+        auth = ThoughtspotBearerAuth(
+            host=credentials.base_url,
+            token_payload=credentials.dict(),
+        )
+        super().__init__(
+            host=credentials.base_url,
+            auth=auth,
+            headers=_THOUGHTSPOT_HEADERS,
+            safe_mode=safe_mode or THOUGHTSPOT_SAFE_MODE,
+        )
+
+    def _metadata_search(
+        self,
+        metadata_type: str,
+    ) -> Iterator[dict]:
+        offset = 0
+        while True:
+            metadata = self._post(
+                ThoughtspotEndpointFactory.metadata_search(),
+                data={
+                    "metadata": [{"type": metadata_type}],
+                    "include_details": True,
+                    "record_size": _METADATA_BATCH_SIZE,
+                    "record_offset": offset,
+                },
+            )
+            yield from metadata
+            if len(metadata) < _METADATA_BATCH_SIZE:
+                break
+            offset = offset + _METADATA_BATCH_SIZE
+
+    def _get_all_liveboards(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="LIVEBOARD")
+
+    def _get_all_columns(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="LOGICAL_COLUMN")
+
+    def _get_all_tables(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="LOGICAL_TABLE")
+
+    def _get_liveboards_usages(self) -> Iterator[dict]:
+        data = self._post(
+            endpoint=ThoughtspotEndpointFactory.liveboard(),
+            headers={"Accept": "application/octet-stream"},
+            data={
+                "metadata_identifier": _USAGE_LIVEBOARD_ID,
+                "file_format": "CSV",
+                "visualization_identifiers": [
+                    "Popular Liveboards Last 30 Days"
+                ],
+            },
+            handler=lambda x: x.text,
+        )
+        yield from usage_liveboard_reader(data)
+
+    def fetch(self, asset: ThoughtspotAsset):
+        if asset == ThoughtspotAsset.LIVEBOARDS:
+            yield from self._get_all_liveboards()
+
+        if asset == ThoughtspotAsset.USAGES:
+            yield from self._get_liveboards_usages()
+
+        if asset == ThoughtspotAsset.LOGICAL_TABLES:
+            yield from self._get_all_tables()

castor_extractor/visualization/thoughtspot/client/credentials.py

@@ -0,0 +1,18 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+THOUGHTSPOT_ENV_PREFIX = "CASTOR_THOUGHTSPOT_"
+
+
+class ThoughtspotCredentials(BaseSettings):
+    """Class to handle Thoughtspot rest API permissions"""
+
+    model_config = SettingsConfigDict(
+        env_prefix=THOUGHTSPOT_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    username: str
+    password: str = Field(repr=False)
+    base_url: str

castor_extractor/visualization/thoughtspot/client/endpoints.py

@@ -0,0 +1,12 @@
+class ThoughtspotEndpointFactory:
+    @classmethod
+    def authentication(cls) -> str:
+        return "api/rest/2.0/auth/token/full"
+
+    @classmethod
+    def metadata_search(cls) -> str:
+        return "api/rest/2.0/metadata/search"
+
+    @classmethod
+    def liveboard(cls) -> str:
+        return "api/rest/2.0/report/liveboard"

castor_extractor/visualization/thoughtspot/client/utils.py

@@ -0,0 +1,25 @@
+import csv
+from io import StringIO
+from typing import Iterator
+
+
+def usage_liveboard_reader(usage_liveboard_csv: str) -> Iterator[dict]:
+    """
+    Converts a CSV string into an iterator of dictionaries after
+    ignoring the first 6 lines, using the 7th line as the header.
+    First 6 lines looks like the following:
+
+        "Data extract produced by Castor on 09/19/2024 06:54"
+        "Filters applied on data :"
+        "User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
+        "Pinboard NOT IN [mlm - availability pinboard,null]"
+        "Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
+        "Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+
+    """
+    csv_file = StringIO(usage_liveboard_csv)
+
+    for _ in range(7):
+        next(csv_file)
+
+    yield from csv.DictReader(csv_file)

castor_extractor/visualization/thoughtspot/client/utils_test.py

@@ -0,0 +1,57 @@
+from .utils import (
+    usage_liveboard_reader,
+)
+
+VALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
+"Filters applied on data :"
+"User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
+"Pinboard NOT IN [mlm - availability pinboard,null]"
+"Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
+"Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+""
+"Pinboard","Pinboard Views","Unique Number of User"
+"Market Report","559","19"
+"Retailer report","204","14"
+"Second-hand market","72","6"
+"September test","25","2"'''
+
+
+# Invalid CSV input (missing data rows)
+INVALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
+"Filters applied on data :"
+"User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
+"Pinboard NOT IN [mlm - availability pinboard,null]"
+"Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
+"Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+""'''
+
+
+def test_usage_liveboard_reader():
+    expected_output = [
+        {
+            "Pinboard": "Market Report",
+            "Pinboard Views": "559",
+            "Unique Number of User": "19",
+        },
+        {
+            "Pinboard": "Retailer report",
+            "Pinboard Views": "204",
+            "Unique Number of User": "14",
+        },
+        {
+            "Pinboard": "Second-hand market",
+            "Pinboard Views": "72",
+            "Unique Number of User": "6",
+        },
+        {
+            "Pinboard": "September test",
+            "Pinboard Views": "25",
+            "Unique Number of User": "2",
+        },
+    ]
+
+    result = list(usage_liveboard_reader(VALID_CSV))
+    assert result == expected_output
+
+    result = list(usage_liveboard_reader(INVALID_CSV))
+    assert result == []  # Expect an empty result since there is no data

castor_extractor/visualization/thoughtspot/extract.py

@@ -0,0 +1,49 @@
+import logging
+from typing import Iterable, Iterator, Tuple, Union
+
+from ...utils import (
+    OUTPUT_DIR,
+    current_timestamp,
+    deep_serialize,
+    from_env,
+    get_output_filename,
+    write_json,
+    write_summary,
+)
+from .assets import ThoughtspotAsset
+from .client import (
+    ThoughtspotClient,
+    ThoughtspotCredentials,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def iterate_all_data(
+    client: ThoughtspotClient,
+) -> Iterable[Tuple[ThoughtspotAsset, Union[list, Iterator, dict]]]:
+    """Iterate over the extracted data from Thoughtspot"""
+
+    for asset in ThoughtspotAsset:
+        logger.info(f"Extracting {asset.value} from API")
+        data = client.fetch(asset)
+        yield asset, deep_serialize(data)
+
+
+def extract_all(**kwargs) -> None:
+    """
+    Extract data from Thoughtspot API
+    Store the output files locally under the given output_directory
+    """
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+
+    credentials = ThoughtspotCredentials(**kwargs)
+    client = ThoughtspotClient(credentials=credentials)
+
+    ts = current_timestamp()
+
+    for key, data in iterate_all_data(client):
+        filename = get_output_filename(key.name.lower(), _output_directory, ts)
+        write_json(filename, list(data))
+
+    write_summary(_output_directory, ts, host=credentials.base_url)

castor_extractor/warehouse/salesforce/client.py

@@ -89,4 +89,4 @@ class SalesforceClient(SalesforceBaseClient):
         ):
             fields = self.fetch_fields(api_name)
             sobject_fields[table_name] = fields
-        return self.formatter.columns(sobject_fields)
+        return list(self.formatter.columns(sobject_fields))

castor_extractor/warehouse/salesforce/format.py

@@ -1,12 +1,30 @@
 from typing import Any, Dict, Iterator, List
 
+from ...utils import group_by
 from .constants import SCHEMA_NAME
 
+_HAS_DUPLICATE_KEY = "#has_duplicate"
+
 
 def _clean(raw: str) -> str:
     return raw.strip('"')
 
 
+def _name(sobject: dict) -> str:
+    """
+    compute name for table and columns
+    - when unique: label
+    - when label is empty or has duplicate: label (api_name)
+    """
+    label = sobject["Label"]
+    api_name = sobject["QualifiedApiName"]
+    if not label:
+        return api_name
+    if not sobject[_HAS_DUPLICATE_KEY]:
+        return label
+    return f"{label} ({api_name})"
+
+
 def _field_description(field: Dict[str, Any]) -> str:
     context: Dict[str, str] = {}
 
@@ -24,7 +42,7 @@ def _field_description(field: Dict[str, Any]) -> str:
 
 
 def _to_column_payload(field: dict, position: int, table_name: str) -> dict:
-    field_name = field["QualifiedApiName"]
+    field_name = _name(field)
     return {
         "column_name": field_name,
         "data_type": field.get("DataType"),
@@ -37,33 +55,31 @@ def _to_column_payload(field: dict, position: int, table_name: str) -> dict:
     }
 
 
-def _to_table_payload(sobject: dict, table_name: str) -> dict:
+def _to_table_payload(sobject: dict) -> dict:
+    name = _name(sobject)
     return {
-        "id": table_name,
+        "id": name,
         "api_name": sobject["QualifiedApiName"],
         "label": sobject["Label"],
         "schema_id": SCHEMA_NAME,
-        "table_name": table_name,
+        "table_name": name,
         "description": sobject.get("Description"),
         "tags": [],
         "type": "TABLE",
     }
 
 
-def _merge_label_and_api_name(sobject: dict) -> dict:
-    label = sobject["Label"]
-    api_name = sobject["QualifiedApiName"]
-    table_name = f"{label} ({api_name})"
-    return _to_table_payload(sobject, table_name)
-
-
-def _by_label(sobjects: List[dict]) -> Dict[str, List[dict]]:
-    by_label: Dict[str, List[dict]] = dict()
+def _detect_duplicates(sobjects: List[dict]) -> List[dict]:
+    """
+    enrich the given data with "has_duplicate" flag:
+    - True when another asset has the same Label in the list
+    - False otherwise
+    """
+    by_label = group_by("Label", sobjects)
     for sobject in sobjects:
         label = sobject["Label"]
-        similar_sobjects = by_label.setdefault(label, [])
-        similar_sobjects.append(sobject)
-    return by_label
+        sobject[_HAS_DUPLICATE_KEY] = len(by_label[label]) > 1
+    return sobjects
 
 
 class SalesforceFormatter:
@@ -76,21 +92,15 @@ class SalesforceFormatter:
     def tables(sobjects: List[dict]) -> Iterator[dict]:
         """
         formats the raw list of sobjects to tables
-        if two tables share the same label, then we add the api name as well
         """
-        by_label = _by_label(sobjects)
-        for label, similars in by_label.items():
-            if len(similars) > 1:
-                yield from [_merge_label_and_api_name(s) for s in similars]
-            else:
-                sobject = similars[0]  # unique sobject on label
-                yield _to_table_payload(sobject, label)
+        sobjects = _detect_duplicates(sobjects)
+        for sobject in sobjects:
+            yield _to_table_payload(sobject)
 
     @staticmethod
-    def columns(sobject_fields: Dict[str, List[dict]]) -> List[dict]:
+    def columns(sobject_fields: Dict[str, List[dict]]) -> Iterator[dict]:
         """formats the raw list of sobject fields to columns"""
-        return [
-            _to_column_payload(field, idx, table_name)
-            for table_name, fields in sobject_fields.items()
-            for idx, field in enumerate(fields)
-        ]
+        for table_name, fields in sobject_fields.items():
+            fields = _detect_duplicates(fields)
+            for index, field in enumerate(fields):
+                yield _to_column_payload(field, index, table_name)