castor-extractor 0.16.3__tar.gz → 0.16.4__tar.gz

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.16.4 - 2024-04-25
+
+* Salesforce: extract sobjects and fields
+
 ## 0.16.3 - 2024-04-24
 
 * Databricks: Extract table owners

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.16.3
+Version: 0.16.4
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/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-0.16.4/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-0.16.3 → castor_extractor-0.16.4}/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-0.16.4/castor_extractor/utils/client/api.py

@@ -0,0 +1,59 @@
+import logging
+from typing import Any, Callable, Dict, Literal, Optional
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_TIMEOUT_MS = 30_000
+
+# https://requests.readthedocs.io/en/latest/api/#requests.request
+HttpMethod = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
+
+
+class APIClient:
+    """
+    API client
+    - authentication via access token
+    """
+
+    def __init__(self, host: str, token: Optional[str] = None):
+        self._host = host
+        self._token = token or ""
+        self._timeout = DEFAULT_TIMEOUT_MS
+
+    @staticmethod
+    def build_url(host: str, path: str):
+        if not host.startswith("https://"):
+            host = "https://" + host
+        return f"{host.strip('/')}/{path}"
+
+    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 processor:
+            return processor(result)
+
+        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-0.16.4/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-0.16.4/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-0.16.4/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-0.16.4/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-0.16.4/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-0.16.3/castor_extractor/visualization/salesforce_reporting/client → castor_extractor-0.16.4/castor_extractor/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-0.16.4/castor_extractor/visualization/salesforce_reporting/__init__.py

@@ -0,0 +1,3 @@
+from .assets import SalesforceReportingAsset
+from .client import SalesforceReportingClient
+from .extract import extract_all

castor_extractor-0.16.4/castor_extractor/visualization/salesforce_reporting/client/__init__.py

@@ -0,0 +1 @@
+from .rest import SalesforceReportingClient

castor_extractor-0.16.4/castor_extractor/visualization/salesforce_reporting/client/rest.py

@@ -0,0 +1,60 @@
+import logging
+from typing import Dict, Iterator, List, Optional
+
+from ....utils.salesforce import SalesforceBaseClient
+from ..assets import SalesforceReportingAsset
+from .soql import queries
+
+logger = logging.getLogger(__name__)
+
+REQUIRING_URL_ASSETS = (
+    SalesforceReportingAsset.REPORTS,
+    SalesforceReportingAsset.DASHBOARDS,
+    SalesforceReportingAsset.FOLDERS,
+)
+
+
+class SalesforceReportingClient(SalesforceBaseClient):
+    """
+    Salesforce Reporting API client
+    """
+
+    def _get_asset_url(
+        self, asset_type: SalesforceReportingAsset, asset: dict
+    ) -> Optional[str]:
+        """
+        Fetch the given Asset + add the corresponding URL.
+        """
+
+        if asset_type == SalesforceReportingAsset.DASHBOARDS:
+            path = f"lightning/r/Dashboard/{asset['Id']}/view"
+            return self.build_url(self._host, path)
+
+        if asset_type == SalesforceReportingAsset.FOLDERS:
+            path = asset["attributes"]["url"].lstrip("/")
+            return self.build_url(self._host, path)
+
+        if asset_type == SalesforceReportingAsset.REPORTS:
+            path = f"lightning/r/Report/{asset['Id']}/view"
+            return self.build_url(self._host, path)
+
+        return None
+
+    def _fetch_and_add_url(
+        self, asset_type: SalesforceReportingAsset
+    ) -> Iterator[dict]:
+        assets = self._query_all(queries[asset_type])
+        for asset in assets:
+            url = self._get_asset_url(asset_type, asset)
+            yield {**asset, "Url": url}
+
+    def fetch(self, asset: SalesforceReportingAsset) -> List[Dict]:
+        """
+        Fetch Salesforce Reporting assets
+        """
+        logger.info(f"Starting extraction of {asset}")
+
+        if asset in REQUIRING_URL_ASSETS:
+            return list(self._fetch_and_add_url(asset))
+
+        return list(self._query_all(queries[asset]))

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/castor_extractor/visualization/salesforce_reporting/extract.py

@@ -10,14 +10,15 @@ from ...utils import (
     write_json,
     write_summary,
 )
+from ...utils.salesforce import SalesforceCredentials
 from .assets import SalesforceReportingAsset
-from .client import SalesforceClient, SalesforceCredentials
+from .client import SalesforceReportingClient
 
 logger = logging.getLogger(__name__)
 
 
 def iterate_all_data(
-    client: SalesforceClient,
+    client: SalesforceReportingClient,
 ) -> Iterable[Tuple[str, Union[list, dict]]]:
     """Iterate over the extracted data from Salesforce"""
 
@@ -30,10 +31,10 @@ def iterate_all_data(
 def extract_all(
     username: str,
     password: str,
-    consumer_key: str,
-    consumer_secret: str,
+    client_id: str,
+    client_secret: str,
     security_token: str,
-    instance_url: str,
+    base_url: str,
     output_directory: Optional[str] = None,
 ) -> None:
     """
@@ -44,11 +45,12 @@ def extract_all(
     creds = SalesforceCredentials(
         username=username,
         password=password,
-        consumer_key=consumer_key,
-        consumer_secret=consumer_secret,
+        client_id=client_id,
+        client_secret=client_secret,
         security_token=security_token,
+        base_url=base_url,
     )
-    client = SalesforceClient(credentials=creds, instance_url=instance_url)
+    client = SalesforceReportingClient(credentials=creds)
     ts = current_timestamp()
 
     for key, data in iterate_all_data(client):

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/castor_extractor/warehouse/databricks/client.py

@@ -31,7 +31,7 @@ class DatabricksClient(APIClient):
         db_allowed: Optional[Set[str]] = None,
         db_blocked: Optional[Set[str]] = None,
     ):
-        super().__init__(credentials)
+        super().__init__(host=credentials.host, token=credentials.token)
         self._db_allowed = db_allowed
         self._db_blocked = db_blocked
         self.formatter = DatabricksFormatter()

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/castor_extractor/warehouse/databricks/credentials.py

@@ -25,7 +25,4 @@ def to_credentials(params: dict) -> DatabricksCredentials:
     """extract Databricks credentials"""
     host = params.get("host") or from_env(_HOST)
     token = params.get("token") or from_env(_TOKEN)
-    return DatabricksCredentials(
-        host=host,
-        token=token,
-    )
+    return DatabricksCredentials(host=host, token=token)

{castor_extractor-0.16.3 → castor_extractor-0.16.4}/castor_extractor/warehouse/databricks/extract.py

@@ -43,7 +43,7 @@ class DatabricksExtractionProcessor:
         self._storage = storage
         self._skip_existing = skip_existing
 
-    def _should_not_reextract(self, asset_group) -> bool:
+    def _should_not_reextract(self, asset_group: WarehouseAssetGroup) -> bool:
         """helper function to determine whether we need to extract"""
         if not self._skip_existing:
             return False

castor_extractor-0.16.4/castor_extractor/warehouse/salesforce/__init__.py

@@ -0,0 +1,6 @@
+from .client import SalesforceClient
+from .extract import (
+    SALESFORCE_ASSETS,
+    SalesforceExtractionProcessor,
+    extract_all,
+)

castor_extractor-0.16.4/castor_extractor/warehouse/salesforce/client.py

@@ -0,0 +1,112 @@
+import logging
+from typing import Dict, Iterator, List
+
+from tqdm import tqdm  # type: ignore
+
+from ...utils.salesforce import SalesforceBaseClient, SalesforceCredentials
+from .format import SalesforceFormatter
+from .soql import SOBJECT_FIELDS_QUERY_TPL, SOBJECTS_QUERY_TPL
+
+logger = logging.getLogger(__name__)
+
+
+class SalesforceClient(SalesforceBaseClient):
+    """
+    Salesforce API client to extract sobjects
+    """
+
+    # Implicit (hard-coded in Salesforce) limitation when using SOQL of 2,000 rows
+    LIMIT_RECORDS_PER_PAGE = 2000
+
+    def __init__(self, credentials: SalesforceCredentials):
+        super().__init__(credentials)
+        self.formatter = SalesforceFormatter()
+
+    @staticmethod
+    def name() -> str:
+        return "Salesforce"
+
+    def _format_query(self, query_template: str, start_durable_id: str) -> str:
+        return query_template.format(
+            start_durable_id=start_durable_id,
+            limit=self.LIMIT_RECORDS_PER_PAGE,
+        )
+
+    def _next_records(
+        self, url: str, query_template: str, start_durable_id: str = "0000"
+    ) -> List[dict]:
+        query = self._format_query(
+            query_template, start_durable_id=start_durable_id
+        )
+        records, _ = self._call(
+            url, params={"q": query}, processor=self._query_processor
+        )
+        return records
+
+    def _is_last_page(self, records: List[dict]) -> bool:
+        return len(records) < self.LIMIT_RECORDS_PER_PAGE
+
+    def _should_query_next_page(
+        self, records: List[dict], page_number: int
+    ) -> bool:
+        return not (
+            self._is_last_page(records)
+            or self._has_reached_pagination_limit(page_number)
+        )
+
+    def _query_all(self, query_template: str) -> Iterator[dict]:
+        """
+        Run a SOQL query over salesforce API
+
+        Note, pagination is performed via a LIMIT in the SOQL query and requires
+        that ids are sorted. The SOQL query must support `limit` and
+        `start_durable_id` as parameters.
+        """
+        url = self.query_url
+        logger.info("querying page 0")
+        records = self._next_records(url, query_template)
+        yield from records
+
+        page_count = 1
+        while self._should_query_next_page(records, page_count):
+            logger.info(f"querying page {page_count}")
+            last_durable_id = records[-1]["DurableId"]
+            records = self._next_records(
+                url, query_template, start_durable_id=last_durable_id
+            )
+            yield from records
+            page_count += 1
+
+    def fetch_sobjects(self) -> List[dict]:
+        """Fetch all sobjects"""
+        logger.info("Extracting sobjects")
+        return list(self._query_all(SOBJECTS_QUERY_TPL))
+
+    def fetch_fields(self, sobject_name: str) -> List[dict]:
+        """Fetches fields of a given sobject"""
+        query = SOBJECT_FIELDS_QUERY_TPL.format(
+            entity_definition_id=sobject_name
+        )
+        response = self._call(self.tooling_url, params={"q": query})
+        return response["records"]
+
+    def tables(self) -> List[dict]:
+        """
+        Get Salesforce sobjects as tables
+        """
+        sobjects = self.fetch_sobjects()
+        logger.info(f"Extracted {len(sobjects)} sobjects")
+        return self.formatter.tables(sobjects)
+
+    def columns(
+        self, sobject_names: List[str], show_progress: bool = True
+    ) -> List[dict]:
+        """
+        Get salesforce sobject fields as columns
+        show_progress: optionally deactivate the tqdm progress bar
+        """
+        sobject_fields: Dict[str, List[dict]] = dict()
+        for sobject_name in tqdm(sobject_names, disable=not show_progress):
+            fields = self.fetch_fields(sobject_name)
+            sobject_fields[sobject_name] = fields
+        return self.formatter.columns(sobject_fields)

castor_extractor-0.16.4/castor_extractor/warehouse/salesforce/constants.py

@@ -0,0 +1,2 @@
+DATABASE_NAME = "salesforce"
+SCHEMA_NAME = "schema"