castor-extractor 0.16.1__py3-none-any.whl → 0.16.4__py3-none-any.whl

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## 0.16.4 - 2024-04-25
+
+* Salesforce: extract sobjects and fields
+
+## 0.16.3 - 2024-04-24
+
+* Databricks: Extract table owners
+
+## 0.16.2 - 2024-04-09
+
+* PowerBI: Extract pages from report
+
 ## 0.16.1 - 2024-04-02
 
 * Systematically escape nul bytes on CSV write

castor_extractor/commands/extract_databricks.py

@@ -1,7 +1,10 @@
+import logging
 from argparse import ArgumentParser
 
 from castor_extractor.warehouse import databricks  # type: ignore
 
+logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+
 
 def main():
     parser = ArgumentParser()

castor_extractor/commands/extract_salesforce.py

@@ -0,0 +1,43 @@
+import logging
+from argparse import ArgumentParser
+
+from castor_extractor.warehouse import salesforce  # type: ignore
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+
+
+def main():
+    parser = ArgumentParser()
+
+    parser.add_argument("-u", "--username", help="Salesforce username")
+    parser.add_argument("-p", "--password", help="Salesforce password")
+    parser.add_argument("-c", "--client-id", help="Salesforce client id")
+    parser.add_argument(
+        "-s", "--client-secret", help="Salesforce client secret"
+    )
+    parser.add_argument(
+        "-t", "--security-token", help="Salesforce security token"
+    )
+    parser.add_argument("-b", "--base-url", help="Salesforce instance URL")
+    parser.add_argument("-o", "--output", help="Directory to write to")
+
+    parser.add_argument(
+        "--skip-existing",
+        dest="skip_existing",
+        action="store_true",
+        help="Skips files already extracted instead of replacing them",
+    )
+    parser.set_defaults(skip_existing=False)
+
+    args = parser.parse_args()
+
+    salesforce.extract_all(
+        username=args.username,
+        password=args.password,
+        client_id=args.client_id,
+        client_secret=args.client_secret,
+        security_token=args.security_token,
+        base_url=args.base_url,
+        output_directory=args.output,
+        skip_existing=args.skip_existing,
+    )

castor_extractor/commands/extract_salesforce_reporting.py

@@ -11,23 +11,23 @@ def main():
 
     parser.add_argument("-u", "--username", help="Salesforce username")
     parser.add_argument("-p", "--password", help="Salesforce password")
-    parser.add_argument("-k", "--consumer-key", help="Salesforce consumer key")
+    parser.add_argument("-c", "--client-id", help="Salesforce client id")
     parser.add_argument(
-        "-s", "--consumer-secret", help="Salesforce consumer secret"
+        "-s", "--client-secret", help="Salesforce client secret"
     )
     parser.add_argument(
         "-t", "--security-token", help="Salesforce security token"
     )
-    parser.add_argument("-l", "--url", help="Salesforce instance URL")
+    parser.add_argument("-b", "--base-url", help="Salesforce instance URL")
     parser.add_argument("-o", "--output", help="Directory to write to")
 
     args = parser.parse_args()
     salesforce_reporting.extract_all(
         username=args.username,
         password=args.password,
-        consumer_key=args.consumer_key,
-        consumer_secret=args.consumer_secret,
+        client_id=args.client_id,
+        client_secret=args.client_secret,
         security_token=args.security_token,
-        instance_url=args.url,
+        base_url=args.base_url,
         output_directory=args.output,
     )

castor_extractor/utils/client/api.py

@@ -1,23 +1,25 @@
-import json
-from typing import Optional
+import logging
+from typing import Any, Callable, Dict, Literal, Optional
 
 import requests
 
-from ...warehouse.databricks.credentials import DatabricksCredentials
+logger = logging.getLogger(__name__)
 
 DEFAULT_TIMEOUT_MS = 30_000
-APICredentials = DatabricksCredentials
+
+# https://requests.readthedocs.io/en/latest/api/#requests.request
+HttpMethod = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
 
 
 class APIClient:
     """
     API client
-    - used for Databricks Unity Catalog for now
-    - authentication via access token for now
+    - authentication via access token
     """
 
-    def __init__(self, credentials: APICredentials):
-        self.credentials = credentials
+    def __init__(self, host: str, token: Optional[str] = None):
+        self._host = host
+        self._token = token or ""
         self._timeout = DEFAULT_TIMEOUT_MS
 
     @staticmethod
@@ -26,25 +28,32 @@ class APIClient:
             host = "https://" + host
         return f"{host.strip('/')}/{path}"
 
-    def _headers(self):
-        return {
-            "Content-type": "application/json",
-            "Authorization": f"Bearer {self.credentials.token}",
-        }
-
-    def get(self, path: str, payload: Optional[dict] = None) -> dict:
-        """
-        path: REST API operation path, such as /api/2.0/clusters/get
-        """
-        url = self.build_url(self.credentials.host, path)
-        response = requests.get(
-            url,
-            data=json.dumps(payload or dict()),
-            headers=self._headers(),
-            timeout=self._timeout,
+    def _headers(self) -> Dict[str, str]:
+        if self._token:
+            return {"Authorization": f"Bearer {self._token}"}
+        return dict()
+
+    def _call(
+        self,
+        url: str,
+        method: HttpMethod = "GET",
+        *,
+        params: Optional[dict] = None,
+        data: Optional[dict] = None,
+        processor: Optional[Callable] = None,
+    ) -> Any:
+        logger.debug(f"Calling {method} on {url}")
+        result = requests.request(
+            method, url, headers=self._headers(), params=params, json=data
         )
+        result.raise_for_status()
 
-        if response.content:
-            return json.loads(response.content)
+        if processor:
+            return processor(result)
 
-        return {}
+        return result.json()
+
+    def get(self, path: str, payload: Optional[dict] = None) -> dict:
+        """path: REST API operation path, such as /api/2.0/clusters/get"""
+        url = self.build_url(self._host, path)
+        return self._call(url=url, data=payload)

castor_extractor/utils/salesforce/__init__.py

@@ -0,0 +1,3 @@
+from .client import SalesforceBaseClient
+from .constants import Keys
+from .credentials import SalesforceCredentials, to_credentials

castor_extractor/utils/salesforce/client.py

@@ -0,0 +1,84 @@
+import logging
+from typing import Iterator, Optional, Tuple
+
+from requests import Response
+
+from ...utils.client.api import APIClient
+from .constants import DEFAULT_API_VERSION, DEFAULT_PAGINATION_LIMIT
+from .credentials import SalesforceCredentials
+
+logger = logging.getLogger(__name__)
+
+
+class SalesforceBaseClient(APIClient):
+    """
+    Salesforce API client.
+    https://developer.salesforce.com/docs/atlas.en-us.api_rest.meta/api_rest/intro_rest.htm
+    """
+
+    api_version = DEFAULT_API_VERSION
+    pagination_limit = DEFAULT_PAGINATION_LIMIT
+
+    PATH_TPL = "services/data/v{version}/{suffix}"
+
+    def __init__(self, credentials: SalesforceCredentials):
+        super().__init__(host=credentials.base_url)
+        self._token = self._access_token(credentials)
+
+    def _access_token(self, credentials: SalesforceCredentials) -> str:
+        url = self.build_url(self._host, "services/oauth2/token")
+        response = self._call(
+            url, "POST", params=credentials.token_request_payload()
+        )
+        return response["access_token"]
+
+    def _full_url(self, suffix: str) -> str:
+        path = self.PATH_TPL.format(version=self.api_version, suffix=suffix)
+        return self.build_url(self._host, path)
+
+    @property
+    def query_url(self) -> str:
+        """Returns the query API url"""
+        return self._full_url("query")
+
+    @property
+    def tooling_url(self) -> str:
+        """Returns the tooling API url"""
+        return self._full_url("tooling/query")
+
+    @staticmethod
+    def _query_processor(response: Response) -> Tuple[dict, Optional[str]]:
+        results = response.json()
+        return results["records"], results.get("nextRecordsUrl")
+
+    def _has_reached_pagination_limit(self, page_number: int) -> bool:
+        return page_number > self.pagination_limit
+
+    def _query_first_page(self, query: str) -> Tuple[Iterator[dict], str]:
+        url = self.query_url
+        logger.info("querying page 0")
+        records, next_page_url = self._call(
+            url, params={"q": query}, processor=self._query_processor
+        )
+        return records, next_page_url
+
+    def _query_all(self, query: str) -> Iterator[dict]:
+        """
+        Run a SOQL query over salesforce API.
+
+        more: https://developer.salesforce.com/docs/atlas.en-us.api_rest.meta/api_rest/dome_query.htm
+        """
+        records, next_page_path = self._query_first_page(query)
+        yield from records
+
+        page_count = 1
+        while next_page_path and not self._has_reached_pagination_limit(
+            page_count
+        ):
+            logger.info(f"querying page {page_count}")
+            url = self.build_url(self._host, next_page_path)
+            records, next_page = self._call(
+                url, processor=self._query_processor
+            )
+            yield from records
+            page_count += 1

castor_extractor/utils/salesforce/client_test.py

@@ -0,0 +1,21 @@
+from unittest.mock import patch
+
+from .client import SalesforceBaseClient
+from .credentials import SalesforceCredentials
+
+
+@patch.object(SalesforceBaseClient, "_call")
+def test_SalesforceBaseClient__urls(mock_call):
+    mock_call.return_value = {"access_token": "the_token"}
+    credentials = SalesforceCredentials(
+        username="usr",
+        password="pw",
+        client_id="key",
+        client_secret="secret",
+        security_token="token",
+        base_url="url",
+    )
+    client = SalesforceBaseClient(credentials)
+
+    assert client.query_url == "https://url/services/data/v59.0/query"
+    assert client.tooling_url == "https://url/services/data/v59.0/tooling/query"

castor_extractor/utils/salesforce/constants.py

@@ -0,0 +1,13 @@
+DEFAULT_API_VERSION = 59.0
+DEFAULT_PAGINATION_LIMIT = 100
+
+
+class Keys:
+    """Salesforce's credentials keys"""
+
+    USERNAME = "username"
+    PASSWORD = "password"  # noqa: S105
+    CLIENT_ID = "client_id"
+    CLIENT_SECRET = "client_secret"  # noqa: S105
+    SECURITY_TOKEN = "security_token"  # noqa: S105
+    BASE_URL = "base_url"

castor_extractor/utils/salesforce/credentials.py

@@ -0,0 +1,65 @@
+from typing import Dict
+
+from ...utils import from_env
+from .constants import Keys
+
+_USERNAME = "CASTOR_SALESFORCE_USERNAME"
+_PASSWORD = "CASTOR_SALESFORCE_PASSWORD"  # noqa: S105
+_SECURITY_TOKEN = "CASTOR_SALESFORCE_SECURITY_TOKEN"  # noqa: S105
+_CLIENT_ID = "CASTOR_SALESFORCE_CLIENT_ID"
+_CLIENT_SECRET = "CASTOR_SALESFORCE_CLIENT_SECRET"  # noqa: S105
+_BASE_URL = "CASTOR_SALESFORCE_BASE_URL"
+
+
+class SalesforceCredentials:
+    """
+    Class to handle Salesforce rest API permissions
+    """
+
+    def __init__(
+        self,
+        *,
+        username: str,
+        password: str,
+        security_token: str,
+        client_id: str,
+        client_secret: str,
+        base_url: str,
+    ):
+        self.username = username
+        self.password = password + security_token
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.base_url = base_url
+
+    def token_request_payload(self) -> Dict[str, str]:
+        """
+        Params to post to the API in order to retrieve the authentication token
+        """
+        return {
+            "grant_type": "password",
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "username": self.username,
+            "password": self.password,
+        }
+
+
+def to_credentials(params: dict) -> SalesforceCredentials:
+    """extract Salesforce credentials"""
+    username = params.get(Keys.USERNAME) or from_env(_USERNAME)
+    password = params.get(Keys.PASSWORD) or from_env(_PASSWORD)
+    security_token = params.get(Keys.SECURITY_TOKEN) or from_env(
+        _SECURITY_TOKEN
+    )
+    client_id = params.get(Keys.CLIENT_ID) or from_env(_CLIENT_ID)
+    client_secret = params.get(Keys.CLIENT_SECRET) or from_env(_CLIENT_SECRET)
+    base_url = params.get(Keys.BASE_URL) or from_env(_BASE_URL)
+    return SalesforceCredentials(
+        username=username,
+        password=password,
+        client_id=client_id,
+        client_secret=client_secret,
+        security_token=security_token,
+        base_url=base_url,
+    )

castor_extractor/{visualization/salesforce_reporting/client → utils/salesforce}/credentials_test.py

@@ -5,9 +5,10 @@ def test_Credentials_token_request_payload():
     creds = SalesforceCredentials(
         username="giphy",
         password="1312",
-        consumer_key="degenie",
-        consumer_secret="fautpasledire",
+        client_id="degenie",
+        client_secret="fautpasledire",
         security_token="yo",
+        base_url="man",
     )
 
     payload = creds.token_request_payload()

castor_extractor/visualization/domo/client/client.py

@@ -14,7 +14,7 @@ RawData = Iterator[dict]
 
 DOMO_PUBLIC_URL = "https://api.domo.com"
 FORMAT = "%Y-%m-%d %I:%M:%S %p"
-DEFAULT_TIMEOUT = 300
+DEFAULT_TIMEOUT = 500
 TOKEN_EXPIRATION_SECONDS = timedelta(seconds=3000)  # auth token lasts 1 hour
 
 IGNORED_ERROR_CODES = (

castor_extractor/visualization/powerbi/client/constants.py

@@ -14,18 +14,18 @@ POST = "POST"
 class Urls:
     """PowerBi's urls"""
 
-    REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
     CLIENT_APP_BASE = "https://login.microsoftonline.com/"
     DEFAULT_SCOPE = "https://analysis.windows.net/powerbi/api/.default"
+    REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
 
     # PBI rest API Routes
     ACTIVITY_EVENTS = f"{REST_API_BASE_PATH}/admin/activityevents"
-    DATASETS = f"{REST_API_BASE_PATH}/admin/datasets"
     DASHBOARD = f"{REST_API_BASE_PATH}/admin/dashboards"
+    DATASETS = f"{REST_API_BASE_PATH}/admin/datasets"
     GROUPS = f"{REST_API_BASE_PATH}/admin/groups"
+    METADATA_GET = f"{REST_API_BASE_PATH}/admin/workspaces/scanResult"
     METADATA_POST = f"{REST_API_BASE_PATH}/admin/workspaces/getInfo"
     METADATA_WAIT = f"{REST_API_BASE_PATH}/admin/workspaces/scanStatus"
-    METADATA_GET = f"{REST_API_BASE_PATH}/admin/workspaces/scanResult"
     REPORTS = f"{REST_API_BASE_PATH}/admin/reports"
     WORKSPACE_IDS = (
         "https://api.powerbi.com/v1.0/myorg/admin/workspaces/modified"

castor_extractor/visualization/powerbi/client/rest.py

@@ -64,19 +64,15 @@ class Client:
 
     def __init__(self, credentials: Credentials):
         self.creds = credentials
-
-    def _access_token(self) -> dict:
         client_app = f"{Urls.CLIENT_APP_BASE}{self.creds.tenant_id}"
-        app = msal.ConfidentialClientApplication(
+        self.app = msal.ConfidentialClientApplication(
             client_id=self.creds.client_id,
             authority=client_app,
             client_credential=self.creds.secret,
         )
 
-        token = app.acquire_token_silent(self.creds.scopes, account=None)
-
-        if not token:
-            token = app.acquire_token_for_client(scopes=self.creds.scopes)
+    def _access_token(self) -> dict:
+        token = self.app.acquire_token_for_client(scopes=self.creds.scopes)
 
         if Keys.ACCESS_TOKEN not in token:
             raise ValueError(f"No access token in token response: {token}")
@@ -248,7 +244,17 @@ class Client:
         Returns a list of reports for the organization.
         https://learn.microsoft.com/en-us/rest/api/power-bi/admin/reports-get-reports-as-admin
         """
-        return self._get(Urls.REPORTS)[Keys.VALUE]
+        reports = self._get(Urls.REPORTS)[Keys.VALUE]
+        for report in reports:
+            report_id = report.get("id")
+            try:
+                url = Urls.REPORTS + f"/{report_id}/pages"
+                pages = self._get(url)[Keys.VALUE]
+                report["pages"] = pages
+            except (requests.HTTPError, requests.exceptions.Timeout) as e:
+                logger.debug(e)
+                continue
+        return reports
 
     def _dashboards(self) -> List[Dict]:
         """

castor_extractor/visualization/powerbi/client/rest_test.py

@@ -31,7 +31,6 @@ def test__access_token(mock_app):
     # init mocks
     valid_response = {"access_token": "mock_token"}
     returning_valid_token = Mock(return_value=valid_response)
-    mock_app.return_value.acquire_token_silent = Mock(return_value=None)
     mock_app.return_value.acquire_token_for_client = returning_valid_token
 
     # init client
@@ -40,30 +39,29 @@ def test__access_token(mock_app):
     # generated token
     assert client._access_token() == valid_response
 
-    # via silent endpoint token
-    mock_app.return_value.acquire_token_silent = returning_valid_token
-    mock_app.return_value.acquire_token_for_client = None
-    assert client._access_token() == valid_response
-
     # token missing in response
     invalid_response = {"not_access_token": "666"}
     returning_invalid_token = Mock(return_value=invalid_response)
-    mock_app.return_value.acquire_token_silent = returning_invalid_token
+    mock_app.return_value.acquire_token_for_client = returning_invalid_token
 
     with pytest.raises(ValueError):
         client._access_token()
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch.object(Client, "_access_token")
-def test__headers(mock_access_token):
+def test__headers(mock_access_token, mock_app):
+    mock_app.return_value = None
     client = _client()
     mock_access_token.return_value = {Keys.ACCESS_TOKEN: "666"}
     assert client._header() == {"Authorization": "Bearer 666"}
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__get(mocked_access_token, mocked_request):
+def test__get(mocked_access_token, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_access_token.return_value = {Keys.ACCESS_TOKEN: "666"}
     fact = {"fact": "Approximately 24 cat skins can make a coat.", "length": 43}
@@ -81,9 +79,11 @@ def test__get(mocked_access_token, mocked_request):
         result = client._get("https/whatev.er")
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__workspace_ids(_, mocked_request):
+def test__workspace_ids(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_request.return_value = Mock(
         json=lambda: [{"id": 1000}, {"id": 1001}, {"id": 1003}],
@@ -112,9 +112,11 @@ def test__workspace_ids(_, mocked_request):
     )
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__post_default(_, mocked_request):
+def test__post_default(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     url = "https://estcequecestbientotleweekend.fr/"
     params = QueryParams.METADATA_SCAN
@@ -129,9 +131,11 @@ def test__post_default(_, mocked_request):
     )
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__post_with_processor(_, mocked_request):
+def test__post_with_processor(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     url = "https://estcequecestbientotleweekend.fr/"
     params = QueryParams.METADATA_SCAN
@@ -146,9 +150,11 @@ def test__post_with_processor(_, mocked_request):
     assert result == 1000
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__datasets(_, mocked_request):
+def test__datasets(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_request.return_value = Mock(
         json=lambda: {"value": [{"id": 1, "type": "dataset"}]},
@@ -164,27 +170,50 @@ def test__datasets(_, mocked_request):
     assert datasets == [{"id": 1, "type": "dataset"}]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__reports(_, mocked_request):
+def test__reports(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
-    mocked_request.return_value = Mock(
-        json=lambda: {"value": [{"id": 1, "type": "report"}]},
-    )
+    page_url = f"{Urls.REPORTS}/1/pages"
+    calls = [
+        call(GET, Urls.REPORTS, data=None, headers=ANY, params=None),
+        call(
+            GET,
+            page_url,
+            data=None,
+            headers=ANY,
+            params=None,
+        ),
+    ]
+    mocked_request.side_effect = [
+        Mock(json=lambda: {"value": [{"id": 1, "type": "report"}]}),
+        Mock(
+            json=lambda: {
+                "value": [
+                    {"name": "page_name", "displayName": "page", "order": 0}
+                ]
+            }
+        ),
+    ]
     reports = client._reports()
-    mocked_request.assert_called_with(
-        GET,
-        Urls.REPORTS,
-        data=None,
-        headers=ANY,
-        params=None,
-    )
-    assert reports == [{"id": 1, "type": "report"}]
+    mocked_request.assert_has_calls(calls)
+
+    assert reports == [
+        {
+            "id": 1,
+            "type": "report",
+            "pages": [{"name": "page_name", "displayName": "page", "order": 0}],
+        }
+    ]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__dashboards(_, mocked_request):
+def test__dashboards(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_request.return_value = Mock(
         json=lambda: {"value": [{"id": 1, "type": "dashboard"}]},
@@ -200,6 +229,7 @@ def test__dashboards(_, mocked_request):
     assert dashboards == [{"id": 1, "type": "dashboard"}]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch.object(Client, "_workspace_ids")
 @patch.object(Client, "_create_scan")
 @patch.object(Client, "_wait_for_scan_result")
@@ -209,7 +239,9 @@ def test__metadata(
     mocked_wait,
     mocked_create_scan,
     mocked_workspace_ids,
+    mock_app,
 ):
+    mock_app.return_value = None
     mocked_workspace_ids.return_value = list(range(200))
     mocked_create_scan.return_value = 314
     mocked_wait.return_value = True
@@ -240,8 +272,10 @@ _CALLS = [
 ]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch.object(Client, "_call")
-def test__activity_events(mocked):
+def test__activity_events(mocked, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked.side_effect = _CALLS
 

castor_extractor/visualization/salesforce_reporting/__init__.py

@@ -1,4 +1,3 @@
 from .assets import SalesforceReportingAsset
-from .client import SalesforceClient
-from .client.credentials import SalesforceCredentials
+from .client import SalesforceReportingClient
 from .extract import extract_all