castor-extractor 0.22.0__tar.gz → 0.22.1__tar.gz

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/CHANGELOG.md

@@ -1,6 +1,10 @@
 
 # Changelog
 
+## 0.22.1 - 2024-12-05
+
+* Salesforce: deduplicate tables
+
 ## 0.22.0 - 2024-12-04
 
 * Stop supporting python3.8

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.22.0
+Version: 0.22.1
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -19,6 +19,7 @@ Provides-Extra: bigquery
 Provides-Extra: databricks
 Provides-Extra: dbt
 Provides-Extra: looker
+Provides-Extra: lookerstudio
 Provides-Extra: metabase
 Provides-Extra: mysql
 Provides-Extra: postgres
@@ -31,6 +32,7 @@ Provides-Extra: tableau
 Requires-Dist: cryptography (>=43.0.0,<44.0.0) ; extra == "snowflake"
 Requires-Dist: databricks-sql-connector (>=3.2.0,<4.0.0) ; extra == "databricks" or extra == "all"
 Requires-Dist: google-api-core (>=2.1.1,<3.0.0)
+Requires-Dist: google-api-python-client (>=2.121.0,<3.0.0) ; extra == "lookerstudio" or extra == "all"
 Requires-Dist: google-auth (>=2,<3)
 Requires-Dist: google-cloud-core (>=2.1.0,<3.0.0)
 Requires-Dist: google-cloud-storage (>=2,<3)
@@ -205,6 +207,10 @@ For any questions or bug report, contact us at [support@castordoc.com](mailto:su
 
 # Changelog
 
+## 0.22.1 - 2024-12-05
+
+* Salesforce: deduplicate tables
+
 ## 0.22.0 - 2024-12-04
 
 * Stop supporting python3.8

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/castor_extractor/utils/__init__.py

@@ -45,6 +45,7 @@ from .time import (
     current_timestamp,
     date_after,
     format_date,
+    format_rfc_3339_date,
     past_date,
     timestamp_ms,
     yesterday,

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/castor_extractor/utils/time.py

@@ -63,5 +63,9 @@ def format_date(timestamp: Union[datetime, date]) -> str:
     return timestamp.strftime(ISO_FORMAT)
 
 
+def format_rfc_3339_date(timestamp: datetime) -> str:
+    return timestamp.isoformat(timespec="seconds") + "Z"
+
+
 def yesterday() -> date:
     return current_date() - timedelta(days=1)

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/castor_extractor/utils/time_test.py

@@ -1,6 +1,6 @@
 from datetime import date, datetime
 
-from .time import at_midnight, date_after, timestamp_ms
+from .time import at_midnight, date_after, format_rfc_3339_date, timestamp_ms
 
 
 def test_at_midnight():
@@ -17,3 +17,10 @@ def test_timestamp_ms():
     result = timestamp_ms(dt)
     expected = 670636800000
     assert result == expected
+
+
+def test_format_rfc_3339_date():
+    dt = datetime(1995, 4, 3, 2, 1)
+    result = format_rfc_3339_date(dt)
+    expected = "1995-04-03T02:01:00Z"
+    assert result == expected

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/__init__.py

@@ -0,0 +1,6 @@
+from .assets import LookerStudioAsset
+from .client import (
+    LookerStudioAssetType,
+    LookerStudioClient,
+    LookerStudioCredentials,
+)

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/assets.py

@@ -0,0 +1,6 @@
+from ...types import ExternalAsset
+
+
+class LookerStudioAsset(ExternalAsset):
+    ASSETS = "assets"
+    VIEW_ACTIVITY = "view_activity"

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/__init__.py

@@ -0,0 +1,3 @@
+from .client import LookerStudioClient
+from .credentials import LookerStudioCredentials
+from .enums import LookerStudioAssetType

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/admin_sdk_client.py

@@ -0,0 +1,90 @@
+from typing import Iterator, Optional
+
+from google.oauth2.service_account import Credentials
+from googleapiclient import discovery  # type: ignore
+
+from ....utils import (
+    at_midnight,
+    current_date,
+    fetch_all_pages,
+    format_rfc_3339_date,
+    past_date,
+)
+from .credentials import LookerStudioCredentials
+from .pagination import LookerStudioPagination
+from .scopes import SCOPES
+
+USER_EMAIL_FIELD = "primaryEmail"
+
+
+class AdminSDKClient:
+    """
+    Client to call the Report API and Directory API.
+    The service account must impersonate and admin account.
+    """
+
+    def __init__(self, credentials: LookerStudioCredentials):
+        self._credentials = Credentials.from_service_account_info(
+            credentials.model_dump(),
+            scopes=SCOPES,
+            subject=credentials.admin_email,  # impersonates an admin
+        )
+        self.directory_api = discovery.build(
+            "admin", "directory_v1", credentials=self._credentials
+        )
+        self.report_api = discovery.build(
+            "admin", "reports_v1", credentials=self._credentials
+        )
+
+    def list_users(self) -> Iterator[dict]:
+        """
+        Lists all users in the domain; only the primaryEmail field is selected.
+        Note:
+        * `my_customer` is an alias to represent the account's `customerId`
+        * `domain_public` allows non-admins to list users. This is technically
+           not necessary here because an admin account is impersonated, but it
+           avoids tapping into unnecessary data & serves for future reference.
+        See
+            https://googleapis.github.io/google-api-python-client/docs/dyn/admin_directory_v1.users.html#list
+            https://developers.google.com/admin-sdk/directory/reference/rest/v1/users/list
+            https://developers.google.com/admin-sdk/directory/v1/guides/manage-users#retrieve_users_non_admin
+            https://stackoverflow.com/a/71083443/14448410
+        """
+
+        def _users(pagination_params: Optional[dict] = None) -> dict:
+            parameters = {
+                "viewType": "domain_public",
+                "customer": "my_customer",
+                "fields": f"users({USER_EMAIL_FIELD}), nextPageToken",
+                **(pagination_params or {}),
+            }
+
+            return self.directory_api.users().list(**parameters).execute()
+
+        yield from fetch_all_pages(_users, LookerStudioPagination)
+
+    def list_view_events(self) -> Iterator[dict]:
+        """
+        Lists all Data Studio View events of the past day.
+        See
+            https://googleapis.github.io/google-api-python-client/docs/dyn/admin_reports_v1.activities.html
+            https://developers.google.com/admin-sdk/reports/reference/rest/v1/activities/list
+            https://developers.google.com/admin-sdk/reports/v1/appendix/activity/data-studio#VIEW
+        """
+
+        def _activity(pagination_params: Optional[dict] = None) -> dict:
+            yesterday = format_rfc_3339_date(at_midnight(past_date(1)))
+            today = format_rfc_3339_date(at_midnight(current_date()))
+
+            parameters = {
+                "userKey": "all",
+                "applicationName": "data_studio",
+                "eventName": "VIEW",
+                "startTime": yesterday,
+                "endTime": today,
+                **(pagination_params or {}),
+            }
+
+            return self.report_api.activities().list(**parameters).execute()
+
+        yield from fetch_all_pages(_activity, LookerStudioPagination)

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/client.py

@@ -0,0 +1,37 @@
+from typing import Iterator
+
+from .. import LookerStudioAsset
+from .admin_sdk_client import USER_EMAIL_FIELD, AdminSDKClient
+from .credentials import LookerStudioCredentials
+from .looker_studio_api_client import LookerStudioAPIClient
+
+
+class LookerStudioClient:
+    """
+    Acts as a wrapper class to fetch Looker Studio assets, which requires
+    coordinating calls between the Admin SDK API and the Looker Studio API.
+    """
+
+    def __init__(self, credentials: LookerStudioCredentials):
+        self.admin_sdk_client = AdminSDKClient(credentials)
+        self.looker_studio_client = LookerStudioAPIClient(credentials)
+
+    def _get_assets(self) -> Iterator[dict]:
+        """
+        Extracts reports and data sources user by user.
+        """
+        users = self.admin_sdk_client.list_users()
+
+        for user in users:
+            email = user[USER_EMAIL_FIELD]
+            yield from self.looker_studio_client.fetch_user_assets(email)
+
+    def fetch(self, asset: LookerStudioAsset) -> Iterator[dict]:
+        if asset == LookerStudioAsset.VIEW_ACTIVITY:
+            yield from self.admin_sdk_client.list_view_events()
+
+        elif asset == LookerStudioAsset.ASSETS:
+            yield from self._get_assets()
+
+        else:
+            raise ValueError(f"The asset {asset}, is not supported")

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/credentials.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel, SecretStr, field_serializer
+
+
+class LookerStudioCredentials(BaseModel):
+    admin_email: str
+    auth_provider_x509_cert_url: str
+    auth_uri: str
+    client_email: str
+    client_id: str
+    client_x509_cert_url: str
+    private_key: SecretStr
+    private_key_id: str
+    project_id: str
+    token_uri: str
+    type: str
+
+    @field_serializer("private_key")
+    def dump_secret(self, pk):
+        """When using model_dump, show private_key value"""
+        return pk.get_secret_value()

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/endpoints.py

@@ -0,0 +1,18 @@
+class LookerStudioAPIEndpoint:
+    BASE_PATH = "https://datastudio.googleapis.com"
+
+    @classmethod
+    def search(cls) -> str:
+        """
+        Search a user's assets.
+        See https://developers.google.com/looker-studio/integrate/api/reference/assets/search
+        """
+        return f"{cls.BASE_PATH}/v1/assets:search"
+
+    @classmethod
+    def permissions(cls, asset_name: str) -> str:
+        """
+        Get the permissions of an asset. The user must be the owner of the asset.
+        See https://developers.google.com/looker-studio/integrate/api/reference/permissions/get
+        """
+        return f"{cls.BASE_PATH}/v1/assets/{asset_name}/permissions"

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/enums.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class LookerStudioAssetType(Enum):
+    DATA_SOURCE = "DATA_SOURCE"
+    EXPLORER = "EXPLORER"
+    REPORT = "REPORT"
+    WORKSPACE = "WORKSPACE"

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/looker_studio_api_client.py

@@ -0,0 +1,102 @@
+from functools import partial
+from typing import Iterator, Optional
+
+from google.auth.transport.requests import Request
+from google.oauth2.service_account import Credentials
+
+from ....utils import (
+    APIClient,
+    BearerAuth,
+    fetch_all_pages,
+)
+from .credentials import LookerStudioCredentials
+from .endpoints import LookerStudioAPIEndpoint
+from .enums import LookerStudioAssetType
+from .pagination import LookerStudioPagination
+from .scopes import SCOPES
+
+
+class LookerStudioAPIAuth(BearerAuth):
+    def __init__(
+        self,
+        credentials: LookerStudioCredentials,
+        subject: Optional[str] = None,
+    ):
+        """
+        Instantiates the service account credentials.
+        If a `subject` email is passed, the service account will impersonate
+        that user and make requests on that user's behalf.
+        """
+        self._credentials = Credentials.from_service_account_info(
+            credentials.model_dump(), scopes=SCOPES
+        )
+        if subject:
+            self._credentials = self._credentials.with_subject(subject)
+
+    def fetch_token(self):
+        self._credentials.refresh(Request())
+        return self._credentials.token
+
+
+class LookerStudioAPIClient(APIClient):
+    def __init__(self, credentials: LookerStudioCredentials):
+        auth = LookerStudioAPIAuth(credentials=credentials)
+        super().__init__(auth=auth)
+
+        self._credentials = credentials
+
+    def _is_private_asset(self, asset_name: str) -> bool:
+        """
+        Returns True if the asset is not viewable by anyone other than the owner.
+
+        The permissions	dict contains `Role: Member[]` key-value pairs and has
+        at least one key-value pair to define the asset's unique OWNER.
+        If another key is present, it means the asset was shared with
+        another person or group.
+
+        See also https://developers.google.com/looker-studio/integrate/api/reference/types#Permissions
+        """
+        data = self._get(LookerStudioAPIEndpoint.permissions(asset_name))
+        permissions = data["permissions"]
+        return len(permissions.keys()) == 1
+
+    def _user_assets(
+        self, asset_type: LookerStudioAssetType, user_email: str
+    ) -> Iterator[dict]:
+        """
+        Yields all assets of the given type, owned by the given user and visible
+        by other members.
+        """
+        request = partial(
+            self._get,
+            LookerStudioAPIEndpoint.search(),
+            params={"assetTypes": [asset_type.value]},
+        )
+        assets = fetch_all_pages(request, LookerStudioPagination)
+
+        for asset in assets:
+            asset_name = asset["name"]
+            owner = asset["owner"]
+            if owner == user_email and not self._is_private_asset(asset_name):
+                yield asset
+
+    def _impersonate_user(self, user_email: str):
+        self._auth = LookerStudioAPIAuth(
+            credentials=self._credentials, subject=user_email
+        )
+
+    def fetch_user_assets(self, user_email: str) -> Iterator[dict]:
+        """Yields assets (reports and data sources) shared by the given user."""
+        self._impersonate_user(user_email)
+
+        reports = self._user_assets(
+            asset_type=LookerStudioAssetType.REPORT,
+            user_email=user_email,
+        )
+        data_sources = self._user_assets(
+            asset_type=LookerStudioAssetType.DATA_SOURCE,
+            user_email=user_email,
+        )
+
+        yield from reports
+        yield from data_sources

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/pagination.py

@@ -0,0 +1,31 @@
+from typing import Optional
+
+from pydantic import AliasChoices, ConfigDict, Field
+from pydantic.alias_generators import to_camel
+
+from ....utils import PaginationModel
+
+NEXT_PAGE_KEY = "pageToken"
+
+
+class LookerStudioPagination(PaginationModel):
+    items: list = Field(
+        default_factory=list,
+        validation_alias=AliasChoices("items", "users", "assets"),
+    )
+    next_page_token: Optional[str] = None
+
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+    def is_last(self) -> bool:
+        return self.next_page_token is None
+
+    def next_page_payload(self) -> dict:
+        return {NEXT_PAGE_KEY: self.next_page_token}
+
+    def page_results(self) -> list:
+        return self.items

castor_extractor-0.22.1/castor_extractor/visualization/looker_studio/client/scopes.py

@@ -0,0 +1,6 @@
+SCOPES = (
+    "https://www.googleapis.com/auth/datastudio",
+    "https://www.googleapis.com/auth/userinfo.profile",
+    "https://www.googleapis.com/auth/admin.reports.audit.readonly",
+    "https://www.googleapis.com/auth/admin.directory.user.readonly",
+)

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/castor_extractor/warehouse/salesforce/format.py

@@ -4,7 +4,7 @@ from typing import Any
 from ...utils import group_by
 from .constants import SCHEMA_NAME
 
-_HAS_DUPLICATE_KEY = "#has_duplicate"
+_HAS_DUPLICATE_KEY = "#has_duplicate_label"
 
 
 def _clean(raw: str) -> str:
@@ -70,9 +70,15 @@ def _to_table_payload(sobject: dict) -> dict:
     }
 
 
-def _detect_duplicates(sobjects: list[dict]) -> list[dict]:
+def _remove_duplicates(sobjects: list[dict]) -> list[dict]:
+    """only keep one object per QualifiedApiName"""
+    by_name = group_by("QualifiedApiName", sobjects)
+    return [objects[0] for _, objects in by_name.items()]
+
+
+def _detect_duplicate_labels(sobjects: list[dict]) -> list[dict]:
     """
-    enrich the given data with "has_duplicate" flag:
+    enrich the given data with "has_duplicate_label" flag:
     - True when another asset has the same Label in the list
     - False otherwise
     """
@@ -94,7 +100,8 @@ class SalesforceFormatter:
         """
         formats the raw list of sobjects to tables
         """
-        sobjects = _detect_duplicates(sobjects)
+        sobjects = _remove_duplicates(sobjects)
+        sobjects = _detect_duplicate_labels(sobjects)
         for sobject in sobjects:
             yield _to_table_payload(sobject)
 
@@ -102,6 +109,6 @@ class SalesforceFormatter:
     def columns(sobject_fields: dict[str, list[dict]]) -> Iterator[dict]:
         """formats the raw list of sobject fields to columns"""
         for table_name, fields in sobject_fields.items():
-            fields = _detect_duplicates(fields)
+            fields = _detect_duplicate_labels(fields)
             for index, field in enumerate(fields):
                 yield _to_column_payload(field, index, table_name)

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/castor_extractor/warehouse/salesforce/format_test.py

@@ -1,9 +1,10 @@
 from .format import (
     _HAS_DUPLICATE_KEY,
     SalesforceFormatter,
-    _detect_duplicates,
+    _detect_duplicate_labels,
     _field_description,
     _name,
+    _remove_duplicates,
 )
 
 
@@ -11,9 +12,10 @@ def _tables_sobjects() -> tuple[dict[str, str], ...]:
     """Returns 4 sobjects with 2 sharing the same label"""
     a = {"Label": "a", "QualifiedApiName": "a_one"}
     b = {"Label": "b", "QualifiedApiName": "b"}
-    c = {"Label": "c", "QualifiedApiName": "c"}
+    c = {"Label": "c", "QualifiedApiName": "c_unique_so_doesnt_matter"}
     a_prime = {"Label": "a", "QualifiedApiName": "a_two"}
-    return a, b, c, a_prime
+    b_exact_duplicate = {"Label": "b", "QualifiedApiName": "b"}
+    return a, b, c, a_prime, b_exact_duplicate
 
 
 def _columns_sobjects() -> dict[str, list[dict]]:
@@ -79,14 +81,14 @@ def test__name():
     assert _name(empty_label_sobject) == "empty_label"
 
 
-def test__detect_duplicates():
+def test__detect_duplicate_labels():
     objects = [
         {"Label": "Foo"},
         {"Label": "Bar"},
         {"Label": "Foo"},
     ]
 
-    objects = _detect_duplicates(objects)
+    objects = _detect_duplicate_labels(objects)
     assert objects == [
         {"Label": "Foo", _HAS_DUPLICATE_KEY: True},
         {"Label": "Bar", _HAS_DUPLICATE_KEY: False},
@@ -94,11 +96,25 @@ def test__detect_duplicates():
     ]
 
 
+def test__remove_duplicates():
+    objects = [
+        {"QualifiedApiName": "Foo"},
+        {"QualifiedApiName": "Bar"},
+        {"QualifiedApiName": "Foo"},
+    ]
+
+    objects = _remove_duplicates(objects)
+    assert len(objects) == 2
+    names = {sobject["QualifiedApiName"] for sobject in objects}
+    assert names == {"Foo", "Bar"}
+
+
 def test_salesforce_formatter_tables():
     sobjects = [*_tables_sobjects()]
-    tables = SalesforceFormatter.tables(sobjects)
+    tables = [t for t in SalesforceFormatter.tables(sobjects)]
     expected_names = {"a (a_one)", "a (a_two)", "b", "c"}
     payload_names = {t["table_name"] for t in tables}
+    assert len(tables) == 4  # we only keep one "b"
     assert payload_names == expected_names
 
 

{castor_extractor-0.22.0 → castor_extractor-0.22.1}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.22.0"
+version = "0.22.1"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"
@@ -29,6 +29,7 @@ python = ">=3.9, <3.13"
 
 cryptography = { version = "^43.0.0", optional = true }
 google-api-core = ">=2.1.1, <3.0.0"
+google-api-python-client = {version = "^2.121.0", optional = true}
 google-auth = "^2"
 google-cloud-core = ">=2.1.0, <3.0.0"
 google-cloud-storage = "^2"
@@ -71,6 +72,7 @@ bigquery = ["sqlalchemy-bigquery", "numpy"]
 databricks = ["databricks-sql-connector", "pandas", "numpy"]
 dbt = []
 looker = ["looker-sdk"]
+lookerstudio = ["google-api-python-client"]
 metabase = ["psycopg2-binary", "pycryptodome"]
 mysql = ["pymysql"]
 powerbi = ["msal"]
@@ -86,6 +88,7 @@ sqlserver = ["pymssql"]
 tableau = ["tableauserverclient"]
 all = [
   "databricks-sql-connector",
+  "google-api-python-client",
   "looker-sdk",
   "msal",
   "numpy",