castor-extractor 0.19.8__py3-none-any.whl → 0.20.4__py3-none-any.whl

castor_extractor/visualization/powerbi/client/constants.py

@@ -1,62 +1,3 @@
-"""
-File regrouping all constants used in PowerBi client
-"""
-
-DEFAULT_TIMEOUT_IN_SECS = 30
-SCAN_READY = "Succeeded"
-# ModifiedSince params should not be older than 30 days
-RECENT_DAYS = 30
-
-GET = "GET"
-POST = "POST"
-
-
-class Urls:
-    """PowerBi's urls"""
-
-    CLIENT_APP_BASE = "https://login.microsoftonline.com/"
-    DEFAULT_SCOPE = "https://analysis.windows.net/powerbi/api/.default"
-    REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
-
-    # PBI rest API Routes
-    ACTIVITY_EVENTS = f"{REST_API_BASE_PATH}/admin/activityevents"
-    DASHBOARD = f"{REST_API_BASE_PATH}/admin/dashboards"
-    DATASETS = f"{REST_API_BASE_PATH}/admin/datasets"
-    GROUPS = f"{REST_API_BASE_PATH}/admin/groups"
-    METADATA_GET = f"{REST_API_BASE_PATH}/admin/workspaces/scanResult"
-    METADATA_POST = f"{REST_API_BASE_PATH}/admin/workspaces/getInfo"
-    METADATA_WAIT = f"{REST_API_BASE_PATH}/admin/workspaces/scanStatus"
-    REPORTS = f"{REST_API_BASE_PATH}/admin/reports"
-    WORKSPACE_IDS = (
-        "https://api.powerbi.com/v1.0/myorg/admin/workspaces/modified"
-    )
-
-
-class Batches:
-    """Batches used within PowerBI api calls"""
-
-    DEFAULT = 100
-    # The route we use to fetch workspaces info can retrieve a maximum of
-    # 100 workspaces per call
-    # More: https://learn.microsoft.com/en-us/rest/api/power-bi/admin/workspace-info-post-workspace-info#request-body
-    METADATA = 100
-
-
-class QueryParams:
-    """
-    Frequently used PowerBi query params
-    """
-
-    METADATA_SCAN = {
-        "datasetExpressions": True,
-        "datasetSchema": True,
-        "datasourceDetails": True,
-        "getArtifactUsers": True,
-        "lineage": True,
-    }
-    ACTIVE_WORKSPACE_FILTER = "state eq 'Active' and type eq 'Workspace'"
-
-
 class Keys:
     ACCESS_TOKEN = "access_token"  # noqa: S105
     ACTIVITY_EVENT_ENTITIES = "activityEventEntities"
@@ -64,15 +5,7 @@ class Keys:
     ID = "id"
     INACTIVE_WORKSPACES = "excludeInActiveWorkspaces"
     LAST_RESULT_SET = "lastResultSet"
-    MODIFIED_SINCE = "modifiedSince"
     PERSONAL_WORKSPACES = "excludePersonalWorkspaces"
     STATUS = "status"
     VALUE = "value"
     WORKSPACES = "workspaces"
-
-
-class Assertions:
-    """Assertion's messages"""
-
-    BATCH_TOO_BIG = f"Can not retrieve more than {Batches.METADATA} at the time"
-    DATETIME_TOO_OLD = "Date must be within 30 days range"

castor_extractor/visualization/powerbi/client/credentials.py

@@ -3,8 +3,7 @@ from typing import List, Optional
 from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from .constants import Urls
-
+DEFAULT_SCOPE = "https://analysis.windows.net/powerbi/api/.default"
 POWERBI_ENV_PREFIX = "CASTOR_POWERBI_"
 
 
@@ -20,9 +19,9 @@ class PowerbiCredentials(BaseSettings):
     client_id: str
     tenant_id: str
     secret: str = Field(repr=False)
-    scopes: List[str] = [Urls.DEFAULT_SCOPE]
+    scopes: List[str] = [DEFAULT_SCOPE]
 
     @field_validator("scopes", mode="before")
     @classmethod
     def _check_scopes(cls, scopes: Optional[List[str]]) -> List[str]:
-        return scopes if scopes is not None else [Urls.DEFAULT_SCOPE]
+        return scopes if scopes is not None else [DEFAULT_SCOPE]

castor_extractor/visualization/powerbi/client/credentials_test.py

@@ -1,5 +1,4 @@
-from .constants import Urls
-from .credentials import PowerbiCredentials
+from .credentials import DEFAULT_SCOPE, PowerbiCredentials
 
 
 def test_credentials():
@@ -13,7 +12,7 @@ def test_credentials():
         client_id=client_id,
         secret=secret,
     )
-    assert credentials.scopes == [Urls.DEFAULT_SCOPE]
+    assert credentials.scopes == [DEFAULT_SCOPE]
 
     credentials = PowerbiCredentials(
         tenant_id=tenant_id,
@@ -21,7 +20,7 @@ def test_credentials():
         secret=secret,
         scopes=None,
     )
-    assert credentials.scopes == [Urls.DEFAULT_SCOPE]
+    assert credentials.scopes == [DEFAULT_SCOPE]
 
     # empty scopes
     credentials = PowerbiCredentials(

castor_extractor/visualization/powerbi/client/endpoints.py

@@ -0,0 +1,65 @@
+from datetime import date, datetime
+from typing import Optional, Tuple
+
+from ....utils import at_midnight, format_date, yesterday
+
+_CLIENT_APP_BASE = "https://login.microsoftonline.com"
+_REST_API_BASE_PATH = "https://api.powerbi.com/v1.0/myorg"
+
+
+def _time_filter(day: Optional[date]) -> Tuple[datetime, datetime]:
+    target_day = day or yesterday()
+    start = at_midnight(target_day)
+    end = datetime.combine(target_day, datetime.max.time())
+    return start, end
+
+
+class PowerBiEndpointFactory:
+    @classmethod
+    def activity_events(cls, day: Optional[date]) -> str:
+        start, end = _time_filter(day)
+        url = f"{_REST_API_BASE_PATH}/admin/activityevents"
+        url += "?$filter=Activity eq 'viewreport'"
+        url += f"&startDateTime='{format_date(start)}'"
+        url += f"&endDateTime='{format_date(end)}'"
+        return url
+
+    @classmethod
+    def authority(cls, tenant_id: str) -> str:
+        return f"{_CLIENT_APP_BASE}/{tenant_id}"
+
+    @classmethod
+    def dashboards(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/dashboards"
+
+    @classmethod
+    def datasets(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/datasets"
+
+    @classmethod
+    def groups(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/groups"
+
+    @classmethod
+    def metadata_create_scan(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/getInfo"
+
+    @classmethod
+    def metadata_scan_result(cls, scan_id: int) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/scanResult/{scan_id}"
+
+    @classmethod
+    def metadata_scan_status(cls, scan_id: int) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/scanStatus/{scan_id}"
+
+    @classmethod
+    def pages(cls, report_id: str) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/reports/{report_id}/pages"
+
+    @classmethod
+    def reports(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/reports"
+
+    @classmethod
+    def workspace_ids(cls) -> str:
+        return f"{_REST_API_BASE_PATH}/admin/workspaces/modified"

castor_extractor/visualization/powerbi/client/pagination.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from pydantic import ConfigDict
+from pydantic.alias_generators import to_camel
+
+from ....utils import (
+    FetchNextPageBy,
+    PaginationModel,
+)
+
+
+class PowerBiPagination(PaginationModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+    fetch_by: FetchNextPageBy = FetchNextPageBy.URL
+
+    activity_event_entities: list
+    continuation_uri: Optional[str] = None
+    last_result_set: bool = False
+
+    def is_last(self) -> bool:
+        return self.last_result_set
+
+    def next_page_payload(self) -> Optional[str]:
+        return self.continuation_uri
+
+    def page_results(self) -> list:
+        return self.activity_event_entities

castor_extractor/visualization/powerbi/extract.py

@@ -1,4 +1,5 @@
-from typing import Iterable, List, Tuple, Union
+import logging
+from typing import Iterable, Tuple, Union
 
 from ...utils import (
     OUTPUT_DIR,
@@ -9,19 +10,23 @@ from ...utils import (
     write_json,
     write_summary,
 )
-from .assets import METADATA_ASSETS, PowerBiAsset
-from .client import Client, PowerbiCredentials
+from .assets import METADATA_ASSETS, REPORTS_ASSETS, PowerBiAsset
+from .client import PowerbiClient, PowerbiCredentials
+
+logger = logging.getLogger(__name__)
 
 
 def iterate_all_data(
-    client: Client,
-) -> Iterable[Tuple[PowerBiAsset, Union[List, dict]]]:
+    client: PowerbiClient,
+) -> Iterable[Tuple[PowerBiAsset, Union[list, dict]]]:
     for asset in PowerBiAsset:
-        if asset in METADATA_ASSETS:
+        if asset in METADATA_ASSETS + REPORTS_ASSETS:
             continue
 
-        data = client.fetch(asset)
-        yield asset, deep_serialize(data)
+        logger.info(f"Extracting {asset.name} from API")
+        data = list(deep_serialize(client.fetch(asset)))
+        yield asset, data
+        logger.info(f"Extracted {len(data)} {asset.name} from API")
 
 
 def extract_all(**kwargs) -> None:
@@ -31,7 +36,7 @@ def extract_all(**kwargs) -> None:
     """
     _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
     creds = PowerbiCredentials(**kwargs)
-    client = Client(creds)
+    client = PowerbiClient(creds)
     ts = current_timestamp()
 
     for key, data in iterate_all_data(client):

castor_extractor/visualization/tableau_revamp/extract.py

@@ -11,7 +11,7 @@ from ...utils import (
     write_summary,
 )
 from .assets import TableauRevampAsset
-from .client import TableauRevampClient
+from .client import TableauRevampClient, TableauRevampCredentials
 
 logger = logging.getLogger(__name__)
 
@@ -26,16 +26,19 @@ def iterate_all_data(
         yield asset, deep_serialize(data)
 
 
-def extract_all(client: TableauRevampClient, **kwargs: str) -> None:
+def extract_all(**kwargs) -> None:
     """
-    Extract Data from tableau
-    Store data locally in files under the output_directory
-    If errors from Tableau's API are catch store them locally in file under the output_directory
+    Extract Data From tableau and store it locally in files under the
+    output_directory
     """
-    output_directory = kwargs.get("output_directory") or from_env(OUTPUT_DIR)
-
+    output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+    with_pulse = kwargs.get("with_pulse") or False
     timestamp = current_timestamp()
 
+    credentials = TableauRevampCredentials(**kwargs)
+    client = TableauRevampClient(credentials, with_pulse=with_pulse)
+    client.login()
+
     for key, data in iterate_all_data(client):
         filename = get_output_filename(key.value, output_directory, timestamp)
         write_json(filename, data)

castor_extractor/visualization/thoughtspot/__init__.py

@@ -0,0 +1,3 @@
+from .assets import ThoughtspotAsset
+from .client import ThoughtspotClient, ThoughtspotCredentials
+from .extract import extract_all, iterate_all_data

castor_extractor/visualization/thoughtspot/assets.py

@@ -0,0 +1,9 @@
+from ...types import ExternalAsset
+
+
+class ThoughtspotAsset(ExternalAsset):
+    """Thoughtspot assets"""
+
+    LIVEBOARDS = "liveboards"
+    LOGICAL_TABLES = "logical_tables"
+    USAGES = "usages"

castor_extractor/visualization/thoughtspot/client/__init__.py

@@ -0,0 +1,2 @@
+from .client import ThoughtspotClient
+from .credentials import ThoughtspotCredentials

castor_extractor/visualization/thoughtspot/client/client.py

@@ -0,0 +1,120 @@
+from typing import Dict, Iterator, Optional
+
+import requests
+
+from ....utils import (
+    APIClient,
+    BearerAuth,
+    RequestSafeMode,
+    build_url,
+    handle_response,
+)
+from ..assets import (
+    ThoughtspotAsset,
+)
+from .credentials import (
+    ThoughtspotCredentials,
+)
+from .endpoints import (
+    ThoughtspotEndpointFactory,
+)
+from .utils import (
+    usage_liveboard_reader,
+)
+
+_AUTH_TIMEOUT_S = 60
+_THOUGHTSPOT_HEADERS = {
+    "X-Requested-By": "ThoughtSpot",
+    "Accept": "application/json",
+    "Content-Type": "application/json",
+}
+_METADATA_BATCH_SIZE = 100
+_USAGE_LIVEBOARD_ID = "bea79810-145f-4ad0-a02c-4177a6e7d861"
+# By default, no errors are ignored for the moment
+THOUGHTSPOT_SAFE_MODE = RequestSafeMode()
+
+
+class ThoughtspotBearerAuth(BearerAuth):
+    def __init__(self, host: str, token_payload: Dict[str, str]):
+        auth_endpoint = ThoughtspotEndpointFactory.authentication()
+        self.authentication_url = build_url(host, auth_endpoint)
+        self.token_payload = token_payload
+
+    def fetch_token(self):
+        token_api_path = self.authentication_url
+        token_response = requests.post(
+            token_api_path, data=self.token_payload, timeout=_AUTH_TIMEOUT_S
+        )
+        return handle_response(token_response)["token"]
+
+
+class ThoughtspotClient(APIClient):
+    def __init__(
+        self,
+        credentials: ThoughtspotCredentials,
+        safe_mode: Optional[RequestSafeMode] = None,
+    ):
+        auth = ThoughtspotBearerAuth(
+            host=credentials.base_url,
+            token_payload=credentials.dict(),
+        )
+        super().__init__(
+            host=credentials.base_url,
+            auth=auth,
+            headers=_THOUGHTSPOT_HEADERS,
+            safe_mode=safe_mode or THOUGHTSPOT_SAFE_MODE,
+        )
+
+    def _metadata_search(
+        self,
+        metadata_type: str,
+    ) -> Iterator[dict]:
+        offset = 0
+        while True:
+            metadata = self._post(
+                ThoughtspotEndpointFactory.metadata_search(),
+                data={
+                    "metadata": [{"type": metadata_type}],
+                    "include_details": True,
+                    "record_size": _METADATA_BATCH_SIZE,
+                    "record_offset": offset,
+                },
+            )
+            yield from metadata
+            if len(metadata) < _METADATA_BATCH_SIZE:
+                break
+            offset = offset + _METADATA_BATCH_SIZE
+
+    def _get_all_liveboards(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="LIVEBOARD")
+
+    def _get_all_columns(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="LOGICAL_COLUMN")
+
+    def _get_all_tables(self) -> Iterator[dict]:
+        yield from self._metadata_search(metadata_type="LOGICAL_TABLE")
+
+    def _get_liveboards_usages(self) -> Iterator[dict]:
+        data = self._post(
+            endpoint=ThoughtspotEndpointFactory.liveboard(),
+            headers={"Accept": "application/octet-stream"},
+            data={
+                "metadata_identifier": _USAGE_LIVEBOARD_ID,
+                "file_format": "CSV",
+                "visualization_identifiers": [
+                    "Popular Liveboards Last 30 Days"
+                ],
+            },
+            handler=lambda x: x.text,
+        )
+        yield from usage_liveboard_reader(data)
+
+    def fetch(self, asset: ThoughtspotAsset):
+        if asset == ThoughtspotAsset.LIVEBOARDS:
+            yield from self._get_all_liveboards()
+
+        if asset == ThoughtspotAsset.USAGES:
+            yield from self._get_liveboards_usages()
+
+        if asset == ThoughtspotAsset.LOGICAL_TABLES:
+            yield from self._get_all_tables()

castor_extractor/visualization/thoughtspot/client/credentials.py

@@ -0,0 +1,18 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+THOUGHTSPOT_ENV_PREFIX = "CASTOR_THOUGHTSPOT_"
+
+
+class ThoughtspotCredentials(BaseSettings):
+    """Class to handle Thoughtspot rest API permissions"""
+
+    model_config = SettingsConfigDict(
+        env_prefix=THOUGHTSPOT_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    username: str
+    password: str = Field(repr=False)
+    base_url: str

castor_extractor/visualization/thoughtspot/client/endpoints.py

@@ -0,0 +1,12 @@
+class ThoughtspotEndpointFactory:
+    @classmethod
+    def authentication(cls) -> str:
+        return "api/rest/2.0/auth/token/full"
+
+    @classmethod
+    def metadata_search(cls) -> str:
+        return "api/rest/2.0/metadata/search"
+
+    @classmethod
+    def liveboard(cls) -> str:
+        return "api/rest/2.0/report/liveboard"

castor_extractor/visualization/thoughtspot/client/utils.py

@@ -0,0 +1,25 @@
+import csv
+from io import StringIO
+from typing import Iterator
+
+
+def usage_liveboard_reader(usage_liveboard_csv: str) -> Iterator[dict]:
+    """
+    Converts a CSV string into an iterator of dictionaries after
+    ignoring the first 6 lines, using the 7th line as the header.
+    First 6 lines looks like the following:
+
+        "Data extract produced by Castor on 09/19/2024 06:54"
+        "Filters applied on data :"
+        "User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
+        "Pinboard NOT IN [mlm - availability pinboard,null]"
+        "Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
+        "Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+
+    """
+    csv_file = StringIO(usage_liveboard_csv)
+
+    for _ in range(7):
+        next(csv_file)
+
+    yield from csv.DictReader(csv_file)

castor_extractor/visualization/thoughtspot/client/utils_test.py

@@ -0,0 +1,57 @@
+from .utils import (
+    usage_liveboard_reader,
+)
+
+VALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
+"Filters applied on data :"
+"User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
+"Pinboard NOT IN [mlm - availability pinboard,null]"
+"Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
+"Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+""
+"Pinboard","Pinboard Views","Unique Number of User"
+"Market Report","559","19"
+"Retailer report","204","14"
+"Second-hand market","72","6"
+"September test","25","2"'''
+
+
+# Invalid CSV input (missing data rows)
+INVALID_CSV = '''"Data extract produced by Castor on 09/19/2024 06:54"
+"Filters applied on data :"
+"User Action IN [pinboard_embed_view,pinboard_tspublic_no_runtime_filter,pinboard_tspublic_runtime_filter,pinboard_view]"
+"Pinboard NOT IN [mlm - availability pinboard,null]"
+"Timestamp >= 20240820 00:00:00 < 20240919 00:00:00"
+"Timestamp >= 20240919 00:00:00 < 20240920 00:00:00"
+""'''
+
+
+def test_usage_liveboard_reader():
+    expected_output = [
+        {
+            "Pinboard": "Market Report",
+            "Pinboard Views": "559",
+            "Unique Number of User": "19",
+        },
+        {
+            "Pinboard": "Retailer report",
+            "Pinboard Views": "204",
+            "Unique Number of User": "14",
+        },
+        {
+            "Pinboard": "Second-hand market",
+            "Pinboard Views": "72",
+            "Unique Number of User": "6",
+        },
+        {
+            "Pinboard": "September test",
+            "Pinboard Views": "25",
+            "Unique Number of User": "2",
+        },
+    ]
+
+    result = list(usage_liveboard_reader(VALID_CSV))
+    assert result == expected_output
+
+    result = list(usage_liveboard_reader(INVALID_CSV))
+    assert result == []  # Expect an empty result since there is no data

castor_extractor/visualization/thoughtspot/extract.py

@@ -0,0 +1,49 @@
+import logging
+from typing import Iterable, Iterator, Tuple, Union
+
+from ...utils import (
+    OUTPUT_DIR,
+    current_timestamp,
+    deep_serialize,
+    from_env,
+    get_output_filename,
+    write_json,
+    write_summary,
+)
+from .assets import ThoughtspotAsset
+from .client import (
+    ThoughtspotClient,
+    ThoughtspotCredentials,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def iterate_all_data(
+    client: ThoughtspotClient,
+) -> Iterable[Tuple[ThoughtspotAsset, Union[list, Iterator, dict]]]:
+    """Iterate over the extracted data from Thoughtspot"""
+
+    for asset in ThoughtspotAsset:
+        logger.info(f"Extracting {asset.value} from API")
+        data = client.fetch(asset)
+        yield asset, deep_serialize(data)
+
+
+def extract_all(**kwargs) -> None:
+    """
+    Extract data from Thoughtspot API
+    Store the output files locally under the given output_directory
+    """
+    _output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+
+    credentials = ThoughtspotCredentials(**kwargs)
+    client = ThoughtspotClient(credentials=credentials)
+
+    ts = current_timestamp()
+
+    for key, data in iterate_all_data(client):
+        filename = get_output_filename(key.name.lower(), _output_directory, ts)
+        write_json(filename, list(data))
+
+    write_summary(_output_directory, ts, host=credentials.base_url)

castor_extractor/warehouse/databricks/api_client.py

@@ -1,5 +1,6 @@
 import logging
 from functools import partial
+from http import HTTPStatus
 from typing import Iterator, List, Optional, Set, Tuple
 
 import requests
@@ -12,6 +13,7 @@ from ...utils import (
     fetch_all_pages,
     handle_response,
     retry,
+    retry_request,
     safe_mode,
 )
 from ..abstract import TimeFilter
@@ -135,6 +137,10 @@ class DatabricksAPIClient(APIClient):
         max_retries=_RETRY_ATTEMPTS,
         base_ms=_RETRY_BASE_MS,
     )
+    @retry_request(
+        status_codes=(HTTPStatus.TOO_MANY_REQUESTS,),
+        max_retries=_RETRY_ATTEMPTS,
+    )
     def get_single_column_lineage(
         self,
         names: Tuple[str, str],

castor_extractor/warehouse/databricks/client.py

@@ -15,7 +15,8 @@ from .types import TablesColumns, TimestampedLink
 
 logger = logging.getLogger(__name__)
 
-_MAX_THREADS = 10
+_THREADS_COLUMN_LINEAGE = 2
+_THREADS_TABLE_LINEAGE = 10
 
 
 class DatabricksClient:
@@ -99,7 +100,7 @@ class DatabricksClient:
         Wrapper function that retrieves all table lineage
         """
         # retrieve table lineage
-        with ThreadPoolExecutor(max_workers=_MAX_THREADS) as executor:
+        with ThreadPoolExecutor(max_workers=_THREADS_TABLE_LINEAGE) as executor:
             table_paths = [
                 ".".join([table["schema_id"], table["table_name"]])
                 for table in tables
@@ -121,10 +122,15 @@ class DatabricksClient:
         candidate_paths = paths_for_column_lineage(
             tables, columns, table_lineage
         )
+        # retrieve column lineage
+        with ThreadPoolExecutor(
+            max_workers=_THREADS_COLUMN_LINEAGE
+        ) as executor:
+            results = executor.map(
+                self.api_client.get_single_column_lineage, candidate_paths
+            )
         lineages: List[TimestampedLink] = [
-            link
-            for paths in candidate_paths
-            for link in self.api_client.get_single_column_lineage(paths)
+            link for links in results for link in links
         ]
         deduplicated = deduplicate_lineage(lineages)
         return self.formatter.format_lineage(deduplicated)