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

castor_extractor/visualization/powerbi/client/credentials.py

@@ -3,8 +3,7 @@ from typing import List, Optional
 from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from .constants import Urls
-
+DEFAULT_SCOPE = "https://analysis.windows.net/powerbi/api/.default"
 POWERBI_ENV_PREFIX = "CASTOR_POWERBI_"
 
 
@@ -20,9 +19,9 @@ class PowerbiCredentials(BaseSettings):
     client_id: str
     tenant_id: str
     secret: str = Field(repr=False)
-    scopes: List[str] = [Urls.DEFAULT_SCOPE]
+    scopes: List[str] = [DEFAULT_SCOPE]
 
     @field_validator("scopes", mode="before")
     @classmethod
     def _check_scopes(cls, scopes: Optional[List[str]]) -> List[str]:
-        return scopes if scopes is not None else [Urls.DEFAULT_SCOPE]
+        return scopes if scopes is not None else [DEFAULT_SCOPE]

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/redshift/extract.py

@@ -34,6 +34,7 @@ REDSHIFT_PASSWORD = "CASTOR_REDSHIFT_PASSWORD"  # noqa: S105
 REDSHIFT_HOST = "CASTOR_REDSHIFT_HOST"
 REDSHIFT_PORT = "CASTOR_REDSHIFT_PORT"
 REDSHIFT_DATABASE = "CASTOR_REDSHIFT_DATABASE"
+REDSHIFT_SERVERLESS = "CASTOR_REDSHIFT_SERVERLESS"
 
 
 def _credentials(params: dict) -> dict:
@@ -48,6 +49,14 @@ def _credentials(params: dict) -> dict:
     }
 
 
+def _query_builder(params: dict) -> RedshiftQueryBuilder:
+    env_parameter = from_env(REDSHIFT_SERVERLESS, allow_missing=True)
+    from_env_ = str(env_parameter).lower() == "true"
+    from_params_ = params.get("serverless", False)
+    is_serverless = from_params_ or from_env_
+    return RedshiftQueryBuilder(is_serverless=is_serverless)
+
+
 def extract_all(**kwargs) -> None:
     """
     Extract all assets from Redshift and store the results in CSV files
@@ -56,7 +65,7 @@ def extract_all(**kwargs) -> None:
 
     client = RedshiftClient(credentials=_credentials(kwargs))
 
-    query_builder = RedshiftQueryBuilder()
+    query_builder = _query_builder(kwargs)
 
     storage = LocalStorage(directory=output_directory)
 

castor_extractor/warehouse/redshift/extract_test.py

@@ -0,0 +1,26 @@
+import pytest
+
+from .extract import (
+    REDSHIFT_SERVERLESS,
+    _query_builder,
+)
+
+
+@pytest.mark.parametrize(
+    "serverless_param,env_param,expected",
+    [
+        (True, "False", True),
+        (False, "True", True),
+        (None, "TRUE", True),
+        (None, "TrUe", True),
+        (None, "FAlSE", False),
+        (None, "False", False),
+        (None, None, False),
+        (True, None, True),
+    ],
+)
+def test__query_builder(serverless_param, env_param, expected, monkeypatch):
+    params = {"serverless": serverless_param}
+    monkeypatch.setenv(REDSHIFT_SERVERLESS, env_param)
+
+    assert _query_builder(params).is_serverless == expected

castor_extractor/warehouse/redshift/queries/query_serverless.sql

@@ -0,0 +1,69 @@
+WITH parameters AS (
+    SELECT
+        :day AS day_start,
+        :hour_min AS hour_min,
+        :hour_max AS hour_max
+),
+
+queries_deduplicated AS (
+    SELECT DISTINCT q.query_id
+    FROM SYS_QUERY_HISTORY AS q
+        CROSS JOIN parameters AS p
+    WHERE TRUE
+        AND DATE(q.start_time) = p.day_start
+        AND EXTRACT('hour' FROM q.start_time) BETWEEN p.hour_min AND p.hour_max
+),
+
+query AS (
+    SELECT
+        q.query_id,
+        qt.text,
+        qt.sequence,
+        COUNT(*) OVER(PARTITION BY q.query_id) AS sequence_count
+    FROM queries_deduplicated AS q
+        INNER JOIN SYS_QUERY_TEXT AS qt ON q.query_id = qt.query_id
+),
+
+raw_query_text AS
+(
+    SELECT
+        q.query_id,
+        LISTAGG(q.text, '') WITHIN GROUP (ORDER BY q.sequence) AS agg_text
+    FROM query AS q
+    WHERE TRUE
+    -- LISTAGG raises an error when total length >= 65535
+    -- each query text contains 4000 char max
+        AND q.sequence_count < (65535 / 4000)
+    GROUP BY q.query_id
+),
+
+query_text AS (
+    SELECT
+        query_id,
+        CASE
+            WHEN agg_text ILIKE 'INSERT INTO%%'
+                THEN REGEXP_REPLACE(agg_text, 'VALUES (.*)', 'DEFAULT VALUES')
+            ELSE agg_text
+        END AS agg_text
+    FROM raw_query_text
+)
+SELECT
+    q.query_id::VARCHAR(256) AS query_id,
+    qt.agg_text::VARCHAR(60000) AS query_text,
+    q.database_name AS database_id,
+    q.database_name AS database_name,
+    q.session_id AS process_id,
+    0 as aborted,
+    q.start_time AS start_time,
+    q.end_time AS end_time,
+    q.user_id AS user_id,
+    q.query_label,
+    u.usename AS user_name
+FROM SYS_QUERY_HISTORY AS q
+    JOIN query_text AS qt ON q.query_id = qt.query_id
+    JOIN pg_catalog.pg_user AS u ON u.usesysid = q.user_id
+    CROSS JOIN parameters AS p
+WHERE TRUE
+    AND DATE(q.start_time) = p.day_start
+    AND EXTRACT('hour' FROM q.start_time) BETWEEN p.hour_min AND p.hour_max
+    AND q.status = 'success'

castor_extractor/warehouse/redshift/query.py

@@ -15,10 +15,21 @@ class RedshiftQueryBuilder(AbstractQueryBuilder):
 
     def __init__(
         self,
+        is_serverless: bool = False,
         time_filter: Optional[TimeFilter] = None,
     ):
         super().__init__(time_filter=time_filter)
+        self.is_serverless = is_serverless
+
+    def build_query_serverless(self) -> ExtractionQuery:
+        """To get the query history in Redshift Serverless, we cannot use STL tables."""
+        statement = self._load_from_file("query_serverless.sql")
+        params = self._time_filter.to_dict()
+        return ExtractionQuery(statement, params)
 
     def build(self, asset: WarehouseAsset) -> List[ExtractionQuery]:
-        query = self.build_default(asset)
+        if asset == WarehouseAsset.QUERY and self.is_serverless:
+            query = self.build_query_serverless()
+        else:
+            query = self.build_default(asset)
         return [query]

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))