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

castor_extractor/visualization/tableau/extract.py

@@ -26,46 +26,70 @@ def iterate_all_data(
     yield TableauAsset.USER, deep_serialize(client.fetch(TableauAsset.USER))
 
     logger.info("Extracting WORKBOOK from Tableau API")
-    yield TableauAsset.WORKBOOK, deep_serialize(
-        client.fetch(TableauAsset.WORKBOOK),
+    yield (
+        TableauAsset.WORKBOOK,
+        deep_serialize(
+            client.fetch(TableauAsset.WORKBOOK),
+        ),
     )
 
     logger.info("Extracting DASHBOARD from Tableau API")
-    yield TableauAsset.DASHBOARD, deep_serialize(
-        client.fetch(TableauAsset.DASHBOARD),
+    yield (
+        TableauAsset.DASHBOARD,
+        deep_serialize(
+            client.fetch(TableauAsset.DASHBOARD),
+        ),
     )
 
     logger.info("Extracting PUBLISHED DATASOURCE from Tableau API")
-    yield TableauAsset.PUBLISHED_DATASOURCE, deep_serialize(
-        client.fetch(TableauAsset.PUBLISHED_DATASOURCE),
+    yield (
+        TableauAsset.PUBLISHED_DATASOURCE,
+        deep_serialize(
+            client.fetch(TableauAsset.PUBLISHED_DATASOURCE),
+        ),
     )
 
     logger.info("Extracting PROJECT from Tableau API")
-    yield TableauAsset.PROJECT, deep_serialize(
-        client.fetch(TableauAsset.PROJECT),
+    yield (
+        TableauAsset.PROJECT,
+        deep_serialize(
+            client.fetch(TableauAsset.PROJECT),
+        ),
     )
 
     logger.info("Extracting USAGE from Tableau API")
     yield TableauAsset.USAGE, deep_serialize(client.fetch(TableauAsset.USAGE))
 
     logger.info("Extracting WORKBOOK_TO_DATASOURCE from Tableau API")
-    yield TableauAsset.WORKBOOK_TO_DATASOURCE, deep_serialize(
-        client.fetch(TableauAsset.WORKBOOK_TO_DATASOURCE),
+    yield (
+        TableauAsset.WORKBOOK_TO_DATASOURCE,
+        deep_serialize(
+            client.fetch(TableauAsset.WORKBOOK_TO_DATASOURCE),
+        ),
     )
 
     logger.info("Extracting DATASOURCE from Tableau API")
-    yield TableauAsset.DATASOURCE, deep_serialize(
-        client.fetch(TableauAsset.DATASOURCE),
+    yield (
+        TableauAsset.DATASOURCE,
+        deep_serialize(
+            client.fetch(TableauAsset.DATASOURCE),
+        ),
     )
 
     logger.info("Extracting CUSTOM_SQL_TABLE from Tableau API")
-    yield TableauAsset.CUSTOM_SQL_TABLE, deep_serialize(
-        client.fetch(TableauAsset.CUSTOM_SQL_TABLE),
+    yield (
+        TableauAsset.CUSTOM_SQL_TABLE,
+        deep_serialize(
+            client.fetch(TableauAsset.CUSTOM_SQL_TABLE),
+        ),
     )
 
     logger.info("Extracting CUSTOM_SQL_QUERY from Tableau API")
-    yield TableauAsset.CUSTOM_SQL_QUERY, deep_serialize(
-        client.fetch(TableauAsset.CUSTOM_SQL_QUERY),
+    yield (
+        TableauAsset.CUSTOM_SQL_QUERY,
+        deep_serialize(
+            client.fetch(TableauAsset.CUSTOM_SQL_QUERY),
+        ),
     )
 
     logger.info("Extracting FIELD from Tableau API")

castor_extractor/visualization/tableau_revamp/client/client.py

@@ -210,7 +210,6 @@ class TableauRevampClient:
         self,
         asset: TableauRevampAsset,
     ) -> SerializedAsset:
-
         if asset == TableauRevampAsset.DATASOURCE:
             data = TSC.Pager(self._server.datasources)
 
@@ -280,13 +279,11 @@ class TableauRevampClient:
 
         return _enrich_workbooks_with_tsc(workbooks, workbook_projects)
 
-    def fetch(
-        self,
-        asset: TableauRevampAsset,
-    ) -> SerializedAsset:
+    def fetch(self, asset: TableauRevampAsset) -> SerializedAsset:
         """
         Extract the given Tableau Asset
         """
+
         if asset == TableauRevampAsset.DATASOURCE:
             # both APIs are required to extract datasources
             return self._fetch_datasources()

castor_extractor/visualization/tableau_revamp/extract.py

@@ -23,8 +23,9 @@ def iterate_all_data(
     """Iterate over the extracted Data from Tableau"""
 
     logger.info("Extracting USER from Tableau API")
-    yield TableauRevampAsset.USER, deep_serialize(
-        client.fetch(TableauRevampAsset.USER)
+    yield (
+        TableauRevampAsset.USER,
+        deep_serialize(client.fetch(TableauRevampAsset.USER)),
     )
 
 

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/databricks/client.py

@@ -4,7 +4,7 @@ from concurrent.futures import ThreadPoolExecutor
 from datetime import date
 from enum import Enum
 from functools import partial
-from typing import Any, Dict, List, Optional, Set, Tuple, cast
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, cast
 
 import requests
 from databricks import sql  # type: ignore
@@ -28,9 +28,13 @@ from .utils import build_path, tag_label
 
 logger = logging.getLogger(__name__)
 
-_DATABRICKS_CLIENT_TIMEOUT = 60
+_DATABRICKS_CLIENT_TIMEOUT = 90
+_DEFAULT_HOUR_MIN = 0
+_DEFAULT_HOUR_MAX = 23
 _MAX_NUMBER_OF_LINEAGE_ERRORS = 1000
+_MAX_NUMBER_OF_QUERY_ERRORS = 1000
 _MAX_THREADS = 10
+_NUM_HOURS_IN_A_DAY = 24
 _RETRY_ATTEMPTS = 3
 _RETRY_BASE_MS = 1000
 _RETRY_EXCEPTIONS = [
@@ -40,7 +44,8 @@ _WORKSPACE_ID_HEADER = "X-Databricks-Org-Id"
 
 _INFORMATION_SCHEMA_SQL = "SELECT * FROM system.information_schema"
 
-safe_params = SafeMode((BaseException,), _MAX_NUMBER_OF_LINEAGE_ERRORS)
+safe_lineage_params = SafeMode((BaseException,), _MAX_NUMBER_OF_LINEAGE_ERRORS)
+safe_query_params = SafeMode((BaseException,), _MAX_NUMBER_OF_QUERY_ERRORS)
 
 
 class TagEntity(Enum):
@@ -112,7 +117,7 @@ class DatabricksClient(APIClient):
         Execute a SQL query on Databricks system tables and return the results.
         https://docs.databricks.com/en/dev-tools/python-sql-connector.html
 
-        /!\ credentials.http_path is required in order to run SQL queries
+        //!\\ credentials.http_path is required in order to run SQL queries
         """
         assert self._http_path, "HTTP_PATH is required to run SQL queries"
         with sql.connect(
@@ -261,7 +266,6 @@ class DatabricksClient(APIClient):
         table_tags = self._get_tags_mapping(TagEntity.TABLE)
         column_tags = self._get_tags_mapping(TagEntity.COLUMN)
         for schema in schemas:
-
             t_to_add, c_to_add = self._tables_columns_of_schema(
                 schema=schema,
                 table_tags=table_tags,
@@ -325,7 +329,7 @@ class DatabricksClient(APIClient):
 
         return list(filter(None, links))
 
-    @safe_mode(safe_params, lambda: [])
+    @safe_mode(safe_lineage_params, lambda: [])
     @retry(
         exceptions=_RETRY_EXCEPTIONS,
         max_retries=_RETRY_ATTEMPTS,
@@ -421,7 +425,7 @@ class DatabricksClient(APIClient):
 
         return list(filter(None, links))
 
-    @safe_mode(safe_params, lambda: [])
+    @safe_mode(safe_lineage_params, lambda: [])
     @retry(
         exceptions=_RETRY_EXCEPTIONS,
         max_retries=_RETRY_ATTEMPTS,
@@ -468,8 +472,20 @@ class DatabricksClient(APIClient):
         return self.formatter.format_lineage(deduplicated)
 
     @staticmethod
-    def _time_filter(time_filter: Optional[TimeFilter]) -> dict:
-        """time filter to retrieve Databricks' queries"""
+    def _time_filter_payload(start_time_ms: int, end_time_ms: int) -> dict:
+        return {
+            "filter_by": {
+                "query_start_time_range": {
+                    "end_time_ms": end_time_ms,
+                    "start_time_ms": start_time_ms,
+                }
+            }
+        }
+
+    def _hourly_time_filters(
+        self, time_filter: Optional[TimeFilter]
+    ) -> Iterable[dict]:
+        """time filters to retrieve Databricks' queries: 1h duration each"""
         # define an explicit time window
         if not time_filter:
             time_filter = TimeFilter.default()
@@ -479,22 +495,13 @@ class DatabricksClient(APIClient):
         hour_min = time_filter.hour_min
         hour_max = time_filter.hour_max
         day = time_filter.day
-        if hour_min is not None and hour_max is not None:  # specific window
-            start_time_ms = _day_hour_to_epoch_ms(day, hour_min)
-            # note: in practice, hour_min == hour_max (hourly query ingestion)
-            end_time_ms = _day_hour_to_epoch_ms(day, hour_max + 1)
-        else:  # fallback to an extraction of the entire day
-            start_time_ms = _day_to_epoch_ms(day)
-            end_time_ms = _day_to_epoch_ms(date_after(day, 1))
+        if hour_min is None or hour_max is None:  # fallback to an entire day
+            hour_min, hour_max = _DEFAULT_HOUR_MIN, _DEFAULT_HOUR_MAX
 
-        return {
-            "filter_by": {
-                "query_start_time_range": {
-                    "end_time_ms": end_time_ms,
-                    "start_time_ms": start_time_ms,
-                }
-            }
-        }
+        for index in range(hour_min, min(hour_max + 1, _NUM_HOURS_IN_A_DAY)):
+            start_time_ms = _day_hour_to_epoch_ms(day, index)
+            end_time_ms = _day_hour_to_epoch_ms(day, index + 1)
+            yield self._time_filter_payload(start_time_ms, end_time_ms)
 
     def query_payload(
         self,
@@ -507,10 +514,11 @@ class DatabricksClient(APIClient):
         if page_token:
             payload: Dict[str, Any] = {"page_token": page_token}
         else:
-            if time_range_filter:
-                payload = {**time_range_filter}
-            else:
-                payload = self._time_filter(None)  # default to yesterday
+            if not time_range_filter:
+                # should never happen.
+                # `time_range_filter` optional to leverage functiontools.partial
+                raise ValueError("Time range not specified")
+            payload = {**time_range_filter}
         if max_results:
             payload["max_results"] = max_results
         return payload
@@ -532,18 +540,29 @@ class DatabricksClient(APIClient):
         content = self.get(path=path, payload=payload)
         return content if content else {}
 
-    def queries(self, time_filter: Optional[TimeFilter] = None) -> List[dict]:
-        """get all queries"""
-        # add a time filter (by default: yesterday)
-        time_range_filter = self._time_filter(time_filter)
+    @safe_mode(safe_query_params, lambda: [])
+    @retry(
+        exceptions=_RETRY_EXCEPTIONS,
+        max_retries=_RETRY_ATTEMPTS,
+        base_ms=_RETRY_BASE_MS,
+    )
+    def _queries(self, filter_: dict) -> List[dict]:
+        """helper to retrieve queries using a given time filter"""
         _time_filtered_scroll_queries = partial(
             self._scroll_queries,
-            time_range_filter=time_range_filter,
+            time_range_filter=filter_,
         )
-
         # retrieve all queries using pagination
-        raw_queries = PagerOnToken(_time_filtered_scroll_queries).all()
+        return PagerOnToken(_time_filtered_scroll_queries).all()
+
+    def queries(self, time_filter: Optional[TimeFilter] = None) -> List[dict]:
+        """get all queries, hour per hour"""
+        time_range_filters = self._hourly_time_filters(time_filter)
 
+        raw_queries = []
+        for _filter in time_range_filters:
+            hourly = self._queries(_filter)
+            raw_queries.extend(hourly)
         return self.formatter.format_query(raw_queries)
 
     def users(self) -> List[dict]:

castor_extractor/warehouse/databricks/client_test.py

@@ -4,7 +4,12 @@ from unittest.mock import Mock, patch
 from freezegun import freeze_time
 
 from ..abstract.time_filter import TimeFilter
-from .client import DatabricksClient, LineageLinks, _day_hour_to_epoch_ms
+from .client import (
+    DatabricksClient,
+    DatabricksCredentials,
+    LineageLinks,
+    _day_hour_to_epoch_ms,
+)
 from .test_constants import (
     CLOSER_DATE,
     MOCK_TABLES_FOR_TABLE_LINEAGE,
@@ -18,45 +23,53 @@ def test__day_hour_to_epoch_ms():
 
 
 @freeze_time("2023-7-4")
-def test_DatabricksClient__time_filter():
+def test_DatabricksClient__hourly_time_filters():
+    credentials = DatabricksCredentials(
+        host="carthago",
+        token="delenda",
+        http_host="est",
+    )
+    client = DatabricksClient(credentials)
+
     # default is yesterday
-    default_time_filter = None
-    default_filter = DatabricksClient._time_filter(default_time_filter)
-    expected_default = {
-        "filter_by": {
-            "query_start_time_range": {
-                "end_time_ms": 1688428800000,  # July 4, 2023 12:00:00 AM GMT
-                "start_time_ms": 1688342400000,  # July 3, 2023 12:00:00 AM GMT
-            }
-        }
-    }
-    assert default_filter == expected_default
+    default_filters = [f for f in client._hourly_time_filters(None)]
+
+    assert len(default_filters) == 24  # number of hours in a day
+
+    first = default_filters[0]
+    start = first["filter_by"]["query_start_time_range"]["start_time_ms"]
+    last = default_filters[-1]
+    end = last["filter_by"]["query_start_time_range"]["end_time_ms"]
+    assert start == 1688342400000  # July 3, 2023 12:00:00 AM GMT
+    assert end == 1688428800000  # July 4, 2023 12:00:00 AM GMT
 
     # custom time (from execution_date in DAG for example)
     time_filter = TimeFilter(day=date(2020, 10, 15))
-    custom_filter = DatabricksClient._time_filter(time_filter)
-    expected_custom = {
-        "filter_by": {
-            "query_start_time_range": {
-                "end_time_ms": 1602806400000,  # October 16, 2020 12:00:00 AM
-                "start_time_ms": 1602720000000,  # October 15, 2020 12:00:00 AM
-            }
-        }
-    }
-    assert custom_filter == expected_custom
+    custom_filters = [f for f in client._hourly_time_filters(time_filter)]
+
+    assert len(custom_filters) == 24
+
+    first = custom_filters[0]
+    start = first["filter_by"]["query_start_time_range"]["start_time_ms"]
+    last = custom_filters[-1]
+    end = last["filter_by"]["query_start_time_range"]["end_time_ms"]
+    assert start == 1602720000000  # Oct 15, 2020 12:00:00 AM
+    assert end == 1602806400000  # Oct 16, 2020 12:00:00 AM
 
     # hourly extraction: note that hour_min == hour_max
     hourly = TimeFilter(day=date(2023, 4, 14), hour_min=4, hour_max=4)
-    hourly_filter = DatabricksClient._time_filter(hourly)
-    expected_hourly = {
-        "filter_by": {
-            "query_start_time_range": {
-                "end_time_ms": 1681448400000,  # April 14, 2023 5:00:00 AM
-                "start_time_ms": 1681444800000,  # April 14, 2023 4:00:00 AM
+    hourly_filters = [f for f in client._hourly_time_filters(hourly)]
+    expected_hourly = [
+        {
+            "filter_by": {
+                "query_start_time_range": {
+                    "end_time_ms": 1681448400000,  # April 14, 2023 5:00:00 AM
+                    "start_time_ms": 1681444800000,  # April 14, 2023 4:00:00 AM
+                }
             }
         }
-    }
-    assert hourly_filter == expected_hourly
+    ]
+    assert hourly_filters == expected_hourly
 
 
 class MockDatabricksClient(DatabricksClient):

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

@@ -19,7 +19,6 @@ def _example_sobjects() -> Tuple[Dict[str, str], ...]:
 
 
 def test__field_description():
-
     field = {}
     assert _field_description(field) == ""
 
@@ -59,7 +58,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,