castor-extractor 0.18.5__tar.gz → 0.18.7__tar.gz

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 0.18.7 - 2024-08-01
+
+* Salesforce: extract table descriptions
+
+## 0.18.6 - 2024-07-30
+
+* BigQuery: introduce extended regions to extract missing queries
+
 ## 0.18.5 - 2024-07-17
 
 * Salesforce: extract DeveloperName and tooling url

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/PKG-INFO

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

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/bigquery/client.py

@@ -1,13 +1,14 @@
+import itertools
 import logging
-from typing import List, Optional, Set, Tuple
+from typing import List, Optional, Set
 
-from google.api_core.exceptions import Forbidden
-from google.api_core.page_iterator import Iterator as PageIterator
+from google.api_core.exceptions import Forbidden  # type: ignore
 from google.cloud.bigquery import Client as GoogleCloudClient  # type: ignore
 from google.cloud.bigquery.dataset import Dataset  # type: ignore
 from google.oauth2.service_account import Credentials  # type: ignore
 
 from ...utils import SqlalchemyClient, retry
+from .types import SetTwoString
 
 logger = logging.getLogger(__name__)
 
@@ -117,16 +118,50 @@ class BigQueryClient(SqlalchemyClient):
             ]
         return self._projects
 
-    def get_regions(self) -> Set[Tuple[str, str]]:
+    def get_regions(self) -> SetTwoString:
         """
-        Returns distinct (project_id, region) available for the given GCP client
+        Returns (project_id, region) available for the given GCP client
+        - Loops trough projects -> datasets -> region
+        - Returns distinct values
+        Example:
+            project_A
+                -> dataset_1:US
+            project_B
+                -> empty
+            project_C
+                -> dataset_2:EU
+                -> dataset_3:EU
+        Will return:
+            { (p_A, US), (p_C, EU) }
         """
         return {
             (ds.project, ds._properties["location"])
             for ds in self._list_datasets()
         }
 
-    def get_datasets(self) -> Set[Tuple[str, str]]:
+    def get_extended_regions(self) -> SetTwoString:
+        """
+        Returns all combinations of (project_id, region) for the given client
+        - Fetch all projects
+        - Fetch all regions (cross projects)
+        - Returns a combination of the two lists
+        Example:
+            project_A
+                -> dataset_1:US
+            project_B
+                -> empty
+            project_C
+                -> dataset_2:EU
+                -> dataset_3:EU
+        Will return:
+            { (p_A, EU), (p_A, US), (p_B, EU), (p_B, US), (p_C, EU), (p_C, US) }
+        """
+        projects = self.get_projects()
+        regions = {ds._properties["location"] for ds in self._list_datasets()}
+        combinations = itertools.product(projects, regions)
+        return set(combinations)
+
+    def get_datasets(self) -> SetTwoString:
         """
         Returns distinct (project_id, dataset_id) available for the given GCP client
         """

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/bigquery/extract.py

@@ -68,6 +68,7 @@ def extract_all(**kwargs) -> None:
     query_builder = BigQueryQueryBuilder(
         regions=client.get_regions(),
         datasets=client.get_datasets(),
+        extended_regions=client.get_extended_regions(),
     )
 
     storage = LocalStorage(directory=output_directory)

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/bigquery/query.py

@@ -2,18 +2,16 @@ import logging
 from typing import List, Optional
 
 from ..abstract import (
-    QUERIES_DIR,
     AbstractQueryBuilder,
     ExtractionQuery,
     TimeFilter,
     WarehouseAsset,
 )
-
-# Those queries must be formatted with {region}
-from .types import IterTwoString
+from .types import SetTwoString
 
 logger = logging.getLogger(__name__)
 
+# Those queries must be formatted with {region}
 REGION_REQUIRED = (
     WarehouseAsset.COLUMN,
     WarehouseAsset.DATABASE,
@@ -23,6 +21,11 @@ REGION_REQUIRED = (
     WarehouseAsset.USER,
 )
 
+# Some clients use empty projects (no datasets) to run their SQL queries
+# The extended regions is a combination of all regions with all projects
+# It allows to extract those queries which were left apart before
+EXTENDED_REGION_REQUIRED = (WarehouseAsset.QUERY,)
+
 # Those queries must be formatted with {dataset}
 DATASET_REQUIRED = (WarehouseAsset.VIEW_DDL,)
 
@@ -38,7 +41,7 @@ SHARDED_ASSETS = (WarehouseAsset.TABLE, WarehouseAsset.COLUMN)
 SHARDED_FILE_PATH = "cte/sharded.sql"
 
 
-def _database_formatted(datasets: IterTwoString) -> str:
+def _database_formatted(datasets: SetTwoString) -> str:
     databases = {db for _, db in datasets}
     if not databases:
         # when no datasets are provided condition should pass
@@ -55,10 +58,11 @@ class BigQueryQueryBuilder(AbstractQueryBuilder):
 
     def __init__(
         self,
-        regions: IterTwoString,
-        datasets: IterTwoString,
+        regions: SetTwoString,
+        datasets: SetTwoString,
         time_filter: Optional[TimeFilter] = None,
         sync_tags: Optional[bool] = False,
+        extended_regions: Optional[SetTwoString] = None,
     ):
         super().__init__(
             time_filter=time_filter,
@@ -67,6 +71,7 @@ class BigQueryQueryBuilder(AbstractQueryBuilder):
         self._regions = regions
         self._datasets = datasets
         self._sync_tags = sync_tags
+        self._extended_regions = extended_regions or regions
 
     @staticmethod
     def _format(query: ExtractionQuery, values: dict) -> ExtractionQuery:
@@ -97,6 +102,13 @@ class BigQueryQueryBuilder(AbstractQueryBuilder):
         sharded_statement = self._load_from_file(SHARDED_FILE_PATH)
         return statement.format(sharded_statement=sharded_statement)
 
+    def _get_regions(self, asset: WarehouseAsset) -> SetTwoString:
+        return (
+            self._extended_regions
+            if asset in EXTENDED_REGION_REQUIRED
+            else self._regions
+        )
+
     def build(self, asset: WarehouseAsset) -> List[ExtractionQuery]:
         """
         It would be easier to stitch data directly in the query statement (UNION ALL).
@@ -110,12 +122,14 @@ class BigQueryQueryBuilder(AbstractQueryBuilder):
         query = super().build_default(asset)
 
         if asset in REGION_REQUIRED:
+            regions = self._get_regions(asset)
+
             logger.info(
-                f"\tWill run queries with following region params: {self._regions}",
+                f"\tWill run queries with following region params: {regions}",
             )
             return [
                 self._format(query, {"project": project, "region": region})
-                for project, region in self._regions
+                for project, region in regions
             ]
 
         if asset in DATASET_REQUIRED:

castor_extractor-0.18.7/castor_extractor/warehouse/bigquery/types.py

@@ -0,0 +1,4 @@
+from typing import Set, Tuple
+
+SetString = Set[str]
+SetTwoString = Set[Tuple[str, str]]

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/salesforce/client.py

@@ -1,11 +1,15 @@
 import logging
-from typing import Dict, Iterator, List, Tuple
+from typing import Dict, Iterator, List, Optional, Tuple
 
 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
+from .soql import (
+    DESCRIPTION_QUERY_TPL,
+    SOBJECT_FIELDS_QUERY_TPL,
+    SOBJECTS_QUERY_TPL,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -90,13 +94,34 @@ class SalesforceClient(SalesforceBaseClient):
         response = self._call(self.tooling_url, params={"q": query})
         return response["records"]
 
+    def fetch_description(self, table_name: str) -> Optional[str]:
+        """Retrieve description of a table"""
+        query = DESCRIPTION_QUERY_TPL.format(table_name=table_name)
+        response = self._call(self.tooling_url, params={"q": query})
+        if not response["records"]:
+            return None
+        return response["records"][0]["Description"]
+
+    def add_table_descriptions(self, sobjects: List[dict]) -> List[dict]:
+        """
+        Add table descriptions.
+        We use the tooling API which does not handle well the LIMIT in SOQL
+        so we have to retrieve descriptions individually
+        """
+        described_sobjects = []
+        for sobject in sobjects:
+            description = self.fetch_description(sobject["QualifiedApiName"])
+            described_sobjects.append({**sobject, "Description": description})
+        return described_sobjects
+
     def tables(self) -> List[dict]:
         """
         Get Salesforce sobjects as tables
         """
         sobjects = self.fetch_sobjects()
         logger.info(f"Extracted {len(sobjects)} sobjects")
-        return list(self.formatter.tables(sobjects))
+        described_sobjects = self.add_table_descriptions(sobjects)
+        return list(self.formatter.tables(described_sobjects))
 
     def columns(
         self, sobject_names: List[Tuple[str, str]], show_progress: bool = True

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/salesforce/format.py

@@ -44,7 +44,7 @@ def _to_table_payload(sobject: dict, table_name: str) -> dict:
         "label": sobject["Label"],
         "schema_id": SCHEMA_NAME,
         "table_name": table_name,
-        "description": "",
+        "description": sobject.get("Description"),
         "tags": [],
         "type": "TABLE",
     }

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/salesforce/format_test.py

@@ -59,7 +59,7 @@ def test__merge_label_and_api_name():
         "label": "foo",
         "schema_id": SCHEMA_NAME,
         "table_name": expected_name,
-        "description": "",
+        "description": None,
         "tags": [],
         "type": "TABLE",
     }

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/castor_extractor/warehouse/salesforce/soql.py

@@ -1,3 +1,9 @@
+DESCRIPTION_QUERY_TPL = """
+    SELECT Description
+    FROM EntityDefinition
+    WHERE QualifiedApiName = '{table_name}'
+"""
+
 SOBJECTS_QUERY_TPL = """
     SELECT
         DeveloperName,
@@ -13,7 +19,6 @@ SOBJECTS_QUERY_TPL = """
     LIMIT {limit}
 """
 
-
 SOBJECT_FIELDS_QUERY_TPL = """
     SELECT
         DataType,

{castor_extractor-0.18.5 → castor_extractor-0.18.7}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.18.5"
+version = "0.18.7"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"

castor_extractor-0.18.5/castor_extractor/warehouse/bigquery/types.py

@@ -1,5 +0,0 @@
-from typing import Iterable, Set, Tuple
-
-SetString = Set[str]
-SetTwoString = Set[Tuple[str, str]]
-IterTwoString = Iterable[Tuple[str, str]]