castor-extractor 0.18.5__py3-none-any.whl → 0.18.7__py3-none-any.whl

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/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/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/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/warehouse/bigquery/types.py

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

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/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/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/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.dist-info → castor_extractor-0.18.7.dist-info}/METADATA

@@ -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.dist-info → castor_extractor-0.18.7.dist-info}/RECORD

@@ -1,4 +1,4 @@
-CHANGELOG.md,sha256=JT2Xz-gVL-ASKBhU8Hxo07ks_as2hr8MpMjD9APvET8,11835
+CHANGELOG.md,sha256=uL2xlPDomxLQhjD55aaCwT6ItdW_ziWGTIjjOAF0IXo,11992
 Dockerfile,sha256=HcX5z8OpeSvkScQsN-Y7CNMUig_UB6vTMDl7uqzuLGE,303
 DockerfileUsage.md,sha256=2hkJQF-5JuuzfPZ7IOxgM6QgIQW7l-9oRMFVwyXC4gE,998
 LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
@@ -263,10 +263,10 @@ castor_extractor/warehouse/abstract/query.py,sha256=GAgeISCmAdrkTKzFGO79hQDf6SA6
 castor_extractor/warehouse/abstract/time_filter.py,sha256=bggIONfMmUxffkA6TwM3BsjfS2l9WFxPq8krfsau5pw,935
 castor_extractor/warehouse/abstract/time_filter_test.py,sha256=PIkegB7KOKBdpc6zIvmyl_CeQyADeFDplyQ8HTNU5LA,448
 castor_extractor/warehouse/bigquery/__init__.py,sha256=PCGNYdi7dHv-SyanUWzRuBp-ypuQ01PkDaQjVnaNhbM,170
-castor_extractor/warehouse/bigquery/client.py,sha256=ypLKXvvfR0RtKex4T2mNvoef4T-jRF1T_RZGCZ6qbOM,4495
+castor_extractor/warehouse/bigquery/client.py,sha256=UefTefQp0S9kpRQzIzJhlm6VcH5uoCRAHHCgfRB6I58,5606
 castor_extractor/warehouse/bigquery/client_test.py,sha256=Ym8e4d--0YQwiVcNUnXLx0X-X6ZznwNMBMbMaDS5oEA,1514
 castor_extractor/warehouse/bigquery/credentials.py,sha256=oCZ8H7qpudKzwM7PRMpVAmWXt7bjIRa8Harmp-ysQJ4,425
-castor_extractor/warehouse/bigquery/extract.py,sha256=vZFxJC1LtUMph5UhfhYdJLnsEto18IOERKzrt71jqJg,2883
+castor_extractor/warehouse/bigquery/extract.py,sha256=lwiam_9YsaUumoY4OwgRMMV4OwLhPwAwnQNCkt-JIeo,2939
 castor_extractor/warehouse/bigquery/queries/.sqlfluff,sha256=ce8UDW2k39v6RBVxgKqjOHHYMoGN9S9f7BCZNHHhox8,30
 castor_extractor/warehouse/bigquery/queries/column.sql,sha256=NxdTnHwomHTEGSc-UoXFKUwg59I9XAOwrSau7JUqGQE,1815
 castor_extractor/warehouse/bigquery/queries/cte/sharded.sql,sha256=-G7_4lxV7UPe72mYlp4HDGeM_fJjZWuXJ7Q0vxvj5_U,1454
@@ -277,8 +277,8 @@ castor_extractor/warehouse/bigquery/queries/table.sql,sha256=D15UNem03Bfcy0JYvKT
 castor_extractor/warehouse/bigquery/queries/table_with_tags.sql,sha256=mhWQHaLgyumtdJX5XyEbdrn_Qtt-RCu4cH1WLM6TN9o,2660
 castor_extractor/warehouse/bigquery/queries/user.sql,sha256=l-fkNGWJVdZwVhbFZL23B8tve-UKXAI6HRlnQq0gIwM,192
 castor_extractor/warehouse/bigquery/queries/view_ddl.sql,sha256=obCm-IN9V8_YSZTwcgNSBDD0ZXPgRjlxJjrZDSEH2MU,326
-castor_extractor/warehouse/bigquery/query.py,sha256=hrFfjd5jW2oQnZ6ozlkn-gDe6sCIzu5zSX19T9W6fIk,4162
-castor_extractor/warehouse/bigquery/types.py,sha256=LZVWSmE57lOemNbB5hBRyYmDk9bFAU4nbRaJWALl6N8,140
+castor_extractor/warehouse/bigquery/query.py,sha256=5Qc8PEa-kQKpTzguj4RNCAwKyvzWt20vAESYNB0lueo,4768
+castor_extractor/warehouse/bigquery/types.py,sha256=DHK3wUaaLyLMp7LP-7QkXTDYpYTZiPtvptAOkpxgp4g,88
 castor_extractor/warehouse/databricks/__init__.py,sha256=YG3YSIJgCFRjjI8eExy9T7qGnfnjWhMFh8c15KTs_BA,184
 castor_extractor/warehouse/databricks/client.py,sha256=ohqsCyLdUJCJGTBK-IBHCV2BUujPG2VsOkc8qAapaPk,20636
 castor_extractor/warehouse/databricks/client_test.py,sha256=KNp4Hi_CC6GwiW2QDJQQwqALfUebuT9D_qL6FuP_8tY,5246
@@ -329,12 +329,12 @@ castor_extractor/warehouse/redshift/queries/user.sql,sha256=sEXveJAuNvZacvpI6Wfw
 castor_extractor/warehouse/redshift/queries/view_ddl.sql,sha256=Pkyh_QT6d4rhTeyiVcqw6O8CRl7NEhk2p7eM5YIn5kg,719
 castor_extractor/warehouse/redshift/query.py,sha256=0C81rkt2cpkWrJIxxwALDyqr-49vlqQM04y_N6wwStc,540
 castor_extractor/warehouse/salesforce/__init__.py,sha256=NR4aNea5jeE1xYqeZ_29deeN84CkN0_D_Z7CLQdJvFY,137
-castor_extractor/warehouse/salesforce/client.py,sha256=ETnZ3n-GFFH0XohDB2ft74wI1HMspvTefR3k7ne-pmI,3891
+castor_extractor/warehouse/salesforce/client.py,sha256=F3UdD5-9umEU-u_c7uVtaksg81VZeXRW83BVsFvYDkE,4902
 castor_extractor/warehouse/salesforce/constants.py,sha256=GusduVBCPvwpk_Im6F3bDvXeNQ7hRnCMdIAjIg65RnE,52
 castor_extractor/warehouse/salesforce/extract.py,sha256=GaxkGWhdksDT_rlT24KX8DMpWnhKlhDMAUvBPGalli0,3454
-castor_extractor/warehouse/salesforce/format.py,sha256=GsWrMEPCMSjP8XuMpDmC5SEMrSp0_xmbCCYQ8PfQebg,3265
-castor_extractor/warehouse/salesforce/format_test.py,sha256=HBlAYBoCOHaq_QOFudZlpcZb5TyZWV9v-cxK4tklg50,2253
-castor_extractor/warehouse/salesforce/soql.py,sha256=pAEaJE8ZUcyN3ptBsZGzNcGRhCcU81X6RMlnF1HRMw4,1063
+castor_extractor/warehouse/salesforce/format.py,sha256=eiPM_4i_m3FEg_2jkMYlhaRBg3gTvV-9xQuk8ghJZiM,3289
+castor_extractor/warehouse/salesforce/format_test.py,sha256=aWUUYDAX-hN1XQJHlv6ZtI2noXWjRobV-9zdjiXR5n4,2255
+castor_extractor/warehouse/salesforce/soql.py,sha256=XB8ohKwHFfC4Xger7Y84DXLW17IJDye_bZ3FL6DCcOI,1188
 castor_extractor/warehouse/snowflake/__init__.py,sha256=TEGXTyxWp4Tr9gIHb-UFVTRKj6YWmrRtqHruiKSZGiY,174
 castor_extractor/warehouse/snowflake/client.py,sha256=XT0QLVNff_586SDuMe40iu8FCwPDh2uBV5aKc1Ql914,5555
 castor_extractor/warehouse/snowflake/client_test.py,sha256=ihWtOOAQfh8pu5JTr_EWfqefKOVIaJXznACURzaU1Qs,1432
@@ -375,8 +375,8 @@ castor_extractor/warehouse/synapse/queries/schema.sql,sha256=aX9xNrBD_ydwl-znGSF
 castor_extractor/warehouse/synapse/queries/table.sql,sha256=mCE8bR1Vb7j7SwZW2gafcXidQ2fo1HwxcybA8wP2Kfs,1049
 castor_extractor/warehouse/synapse/queries/user.sql,sha256=sTb_SS7Zj3AXW1SggKPLNMCd0qoTpL7XI_BJRMaEpBg,67
 castor_extractor/warehouse/synapse/queries/view_ddl.sql,sha256=3EVbp5_yTgdByHFIPLHmnoOnqqLE77SrjAwFDvu4e54,249
-castor_extractor-0.18.5.dist-info/LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
-castor_extractor-0.18.5.dist-info/METADATA,sha256=MCiHMiLGVUz69hRClDJHsMECllqppMC4yZkg-22ETNM,7209
-castor_extractor-0.18.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-castor_extractor-0.18.5.dist-info/entry_points.txt,sha256=SbyPk58Gh-FRztfCNnUZQ6w7SatzNJFZ6GIJLNsy7tI,1427
-castor_extractor-0.18.5.dist-info/RECORD,,
+castor_extractor-0.18.7.dist-info/LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
+castor_extractor-0.18.7.dist-info/METADATA,sha256=QCkJjnGnFsOfbZ808-jK6dthnNFKqzC_YLydG1sUSqM,7209
+castor_extractor-0.18.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+castor_extractor-0.18.7.dist-info/entry_points.txt,sha256=SbyPk58Gh-FRztfCNnUZQ6w7SatzNJFZ6GIJLNsy7tI,1427
+castor_extractor-0.18.7.dist-info/RECORD,,