castor-extractor 0.18.7__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/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/format_test.py

@@ -19,7 +19,6 @@ def _example_sobjects() -> Tuple[Dict[str, str], ...]:
 
 
 def test__field_description():
-
     field = {}
     assert _field_description(field) == ""
 

{castor_extractor-0.18.7.dist-info → castor_extractor-0.19.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.18.7
+Version: 0.19.0
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -29,7 +29,7 @@ Provides-Extra: redshift
 Provides-Extra: snowflake
 Provides-Extra: sqlserver
 Provides-Extra: tableau
-Requires-Dist: cryptography (>=41.0.5) ; extra == "snowflake"
+Requires-Dist: cryptography (>=42.0.0,<43.0.0) ; extra == "snowflake"
 Requires-Dist: databricks-sql-connector (>=3.2.0,<4.0.0) ; extra == "databricks" or extra == "all"
 Requires-Dist: google-api-core (>=2.1.1,<3.0.0)
 Requires-Dist: google-auth (>=2,<3)
@@ -37,7 +37,7 @@ Requires-Dist: google-cloud-core (>=2.1.0,<3.0.0)
 Requires-Dist: google-cloud-storage (>=2,<3)
 Requires-Dist: google-resumable-media (>=2.0.3,<3.0.0)
 Requires-Dist: googleapis-common-protos (>=1.53.0,<2.0.0)
-Requires-Dist: looker-sdk (>=23.0.0,<24.0.0) ; extra == "looker" or extra == "all"
+Requires-Dist: looker-sdk (>=24.0.0,<25.0.0) ; extra == "looker" or extra == "all"
 Requires-Dist: msal (>=1.20.0,<2.0.0) ; extra == "powerbi" or extra == "all"
 Requires-Dist: numpy (<1.25) ; (python_version >= "3.8" and python_version < "3.9") and (extra == "bigquery" or extra == "databricks" or extra == "all")
 Requires-Dist: numpy (<2) ; extra == "bigquery" or extra == "databricks" or extra == "all"
@@ -52,7 +52,7 @@ Requires-Dist: pymssql (>=2.2.11,<3.0.0) ; extra == "sqlserver" or extra == "all
 Requires-Dist: pymysql[rsa] (>=1.1.0,<2.0.0) ; extra == "mysql" or extra == "all"
 Requires-Dist: python-dateutil (>=2.0.0,<=3.0.0)
 Requires-Dist: requests (>=2.0.0,<3.0.0)
-Requires-Dist: setuptools (>=70,<71)
+Requires-Dist: setuptools (>=72,<73)
 Requires-Dist: snowflake-connector-python (>=3.4.0,<4.0.0) ; extra == "snowflake" or extra == "all"
 Requires-Dist: snowflake-sqlalchemy (!=1.2.5,<2.0.0) ; extra == "snowflake" or extra == "all"
 Requires-Dist: sqlalchemy (>=1.4,<1.5)