castor-extractor 0.16.11__tar.gz → 0.17.0__tar.gz

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/CHANGELOG.md

@@ -1,5 +1,25 @@
 # Changelog
 
+## 0.17.0 - 2024-06-10
+
+* Uploader: redirect to the proxy, replace credentials with token
+
+## 0.16.15 - 2024-06-07
+
+* Tableau: extract database_name for CustomSQLTables
+
+## 0.16.14 - 2024-06-06
+
+* Snowflake: Extract SQL user defined function
+
+## 0.16.13 - 2024-06-05
+
+* Tableau: extract database_name for tables
+
+## 0.16.12 - 2024-06-04
+
+* Databricks: Extract lineage
+
 ## 0.16.11 - 2024-06-03
 
 * Tableau: add extra fields to optimise storage

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/PKG-INFO

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

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/commands/upload.py

@@ -13,10 +13,10 @@ logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
 def _args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "-c",
-        "--credentials",
+        "-k",
+        "--token",
         required=True,
-        help="""Path to credentials or credentials as string""",
+        help="""API token provided by Castor""",
     )
     parser.add_argument(
         "-s",
@@ -44,7 +44,7 @@ def _args():
     )
     parsed = parser.parse_args()
     return {
-        "credentials": parsed.credentials,
+        "token": parsed.token,
         "source_id": parsed.source_id,
         "file_path": parsed.file_path,
         "directory_path": parsed.directory_path,

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/uploader/constant.py

@@ -1,6 +1,13 @@
 from enum import Enum
 
-EXTRACTION_BUCKET = "extraction-storage"
+from ..utils import RetryStrategy
+
+# url of the gcs proxy
+INGEST_URL = "https://ingest.castordoc.com"
+
+RETRY_BASE_MS = 10_000
+RETRY_JITTER_MS = 1_000
+RETRY_STRATEGY = RetryStrategy.LINEAR
 
 
 class FileType(Enum):

castor_extractor-0.17.0/castor_extractor/uploader/upload.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+import logging
+import ntpath
+from datetime import datetime
+from typing import Dict, Iterable, Optional, Tuple
+from uuid import UUID
+
+import requests
+
+from ..utils.retry import retry
+from .constant import (
+    INGEST_URL,
+    PATH_TEMPLATES,
+    RETRY_BASE_MS,
+    RETRY_JITTER_MS,
+    RETRY_STRATEGY,
+    FileType,
+)
+from .env import get_blob_env
+from .utils import iter_files
+
+logger = logging.getLogger(__name__)
+
+_EXCEPTIONS = (
+    requests.exceptions.Timeout,
+    requests.exceptions.ConnectTimeout,
+)
+
+
+def _path_and_url(
+    source_id: UUID,
+    file_type: FileType,
+    file_path: str,
+) -> Tuple[str, str]:
+
+    now = datetime.utcnow()
+    timestamp = int(now.timestamp())
+    filename = ntpath.basename(file_path)
+    path_template = PATH_TEMPLATES[file_type]
+    path = path_template.format(
+        timestamp=timestamp,
+        source_id=source_id,
+        filename=filename,
+    )
+
+    url = f"{INGEST_URL}/{path}"
+
+    return path, url
+
+
+def _headers(token: str) -> Dict:
+    return {
+        "Authorization": f"Token {token}",
+        "Accept": "text/csv, application/json",
+    }
+
+
+def _upload(
+    token: str,
+    source_id: UUID,
+    file_path: str,
+    file_type: FileType,
+) -> None:
+    """
+    Upload the given file to Google Cloud Storage (GCS)
+    - Don't call GCS API directly
+    - Call the ingestion proxy which handles authorisation and uploading
+    """
+    path, url = _path_and_url(source_id, file_type, file_path)
+    headers = _headers(token)
+    timeout, max_retries = get_blob_env()
+
+    with open(file_path, "rb") as file_content:
+
+        @retry(
+            exceptions=_EXCEPTIONS,
+            max_retries=max_retries,
+            base_ms=RETRY_BASE_MS,
+            jitter_ms=RETRY_JITTER_MS,
+            strategy=RETRY_STRATEGY,
+        )
+        def _request_post():
+            response = requests.post(
+                url=url,
+                headers=headers,
+                data=file_content,
+                timeout=timeout,
+            )
+            response.raise_for_status()
+
+        _request_post()
+
+    logger.info(f"Uploaded {file_path} as {file_type.value} to {path}")
+
+
+def upload_manifest(
+    token: str,
+    source_id: UUID,
+    file_path: str,
+) -> None:
+    """
+    credentials: path to file or dict
+    source_id: id for the source
+    file_path: path to the local manifest to upload
+    """
+    _upload(token, source_id, file_path, FileType.DBT)
+
+
+def upload(
+    token: str,
+    source_id: UUID,
+    file_type: FileType,
+    file_path: Optional[str] = None,
+    directory_path: Optional[str] = None,
+) -> None:
+    """
+    credentials: path to file or dict
+    source_id: id for the source
+    file_type: type of file(s) uploaded - see FileType Enum
+    file_path: path to the local visualization or warehouse file to upload
+    directory_path: path to the local directory containing files to upload
+    """
+    files: Iterable[str]
+    if directory_path:
+        files = iter_files(directory_path)
+    elif file_path:
+        files = [file_path]
+    else:
+        message = "either file_path or directory_path should be defined"
+        raise ValueError(message)
+
+    for file_ in files:
+        _upload(token, source_id, file_, file_type)

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/uploader/upload_test.py

@@ -1,7 +1,7 @@
 from uuid import UUID
 
-from .constant import FileType
-from .upload import _path
+from .constant import INGEST_URL, FileType
+from .upload import _path_and_url
 
 
 def test__path():
@@ -9,5 +9,6 @@ def test__path():
     file_type = FileType.VIZ
     file_path = "filename"
 
-    path = _path(source_id, file_type, file_path)
+    path, url = _path_and_url(source_id, file_type, file_path)
     assert path == f"visualization-{source_id}/{file_path}"
+    assert url == f"{INGEST_URL}/{path}"

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/uploader/utils.py

@@ -1,13 +1,23 @@
+import logging
 import os
 from typing import Iterator
 
+logger = logging.getLogger(__name__)
+
+_ALLOWED_EXTENSION = (".json", ".csv")
+
 
 def iter_files(repository_path: str) -> Iterator[str]:
     """
     Given a repository path yield all files in that repository
+    Removes file whose extension is not allowed
     """
 
     for file in os.listdir(repository_path):
+        _, ext = os.path.splitext(file)
+        if ext not in _ALLOWED_EXTENSION:
+            logger.info(f"Forbidden file extension : skipping {file}")
+            continue
         file_path = os.path.join(repository_path, file)
 
         if os.path.isfile(file_path):

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/utils/client/api.py

@@ -5,7 +5,7 @@ import requests
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_TIMEOUT_MS = 30_000
+DEFAULT_TIMEOUT_S = 30
 
 # https://requests.readthedocs.io/en/latest/api/#requests.request
 HttpMethod = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]
@@ -20,7 +20,7 @@ class APIClient:
     def __init__(self, host: str, token: Optional[str] = None):
         self._host = host
         self._token = token or ""
-        self._timeout = DEFAULT_TIMEOUT_MS
+        self._timeout = DEFAULT_TIMEOUT_S
 
     @staticmethod
     def build_url(host: str, path: str):
@@ -44,7 +44,12 @@ class APIClient:
     ) -> Any:
         logger.debug(f"Calling {method} on {url}")
         result = requests.request(
-            method, url, headers=self._headers(), params=params, json=data
+            method,
+            url,
+            headers=self._headers(),
+            params=params,
+            json=data,
+            timeout=self._timeout,
         )
         result.raise_for_status()
 

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/utils/retry.py

@@ -68,7 +68,8 @@ class Retry(BaseModel):
         self._retry_attempts += 1
         wait_ms = self.base() + self.jitter()
         wait_s = float(wait_ms) / MS_IN_SEC
-        logger.warning(f"Attempting a new call in {wait_s} seconds")
+        msg = f"Attempting a new call in {wait_s} seconds, {self._retry_attempts} attempt(s) / {self.max_retries} max retries"
+        logger.warning(msg)
         time.sleep(wait_s)
         return True
 
@@ -93,6 +94,7 @@ def retry(
             try:
                 return None, callable(*args, **kwargs)
             except exceptions_ as err:
+                logger.warning(f"Exception within {callable.__name__}")
                 return err, None
 
         def _func(*args, **kwargs) -> Any:

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/visualization/tableau_revamp/client/client.py

@@ -31,8 +31,11 @@ _TSC_ASSETS = (
 # increase the value when extraction is too slow
 # decrease the value when timeouts arise
 _CUSTOM_PAGE_SIZE: Dict[TableauRevampAsset, int] = {
-    # fields and columns are light but volumes are bigger
-    TableauRevampAsset.COLUMN: 200,
+    # for some clients, extraction of columns tend to hit the node limit
+    # https://community.tableau.com/s/question/0D54T00000YuK60SAF/metadata-query-nodelimitexceeded-error
+    # the workaround is to reduce pagination
+    TableauRevampAsset.COLUMN: 50,
+    # fields are light but volumes are bigger
     TableauRevampAsset.FIELD: 1000,
     TableauRevampAsset.TABLE: 50,
 }

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/visualization/tableau_revamp/client/gql_queries.py

@@ -63,12 +63,21 @@ downstreamWorkbooks { id }
 id
 name
 ... on DatabaseTable {
-    connectionType
     fullName
     schema
+    database {
+        connectionType
+        id
+        name
+    }
 }
 ... on CustomSQLTable {
     query
+    database {
+        connectionType
+        id
+        name
+    }
 }
 """
 

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/warehouse/abstract/__init__.py

@@ -1,6 +1,8 @@
 from .asset import (
+    ADDITIONAL_LINEAGE_ASSETS,
     CATALOG_ASSETS,
     EXTERNAL_LINEAGE_ASSETS,
+    FUNCTIONS_ASSETS,
     QUERIES_ASSETS,
     VIEWS_ASSETS,
     SupportedAssets,

{castor_extractor-0.16.11 → castor_extractor-0.17.0}/castor_extractor/warehouse/abstract/asset.py

@@ -7,6 +7,8 @@ from ...types import ExternalAsset, classproperty
 class WarehouseAsset(ExternalAsset):
     """Assets that can be extracted from warehouses"""
 
+    ADDITIONAL_COLUMN_LINEAGE = "additional_column_lineage"
+    ADDITIONAL_TABLE_LINEAGE = "additional_table_lineage"
     COLUMN = "column"
     COLUMN_LINEAGE = "column_lineage"  # specific to snowflake
     DATABASE = "database"
@@ -19,22 +21,28 @@ class WarehouseAsset(ExternalAsset):
     ROLE = "role"
     SCHEMA = "schema"
     TABLE = "table"
+    FUNCTION = "function"
     USER = "user"
     VIEW_DDL = "view_ddl"
 
     @classproperty
     def optional(cls) -> Set["WarehouseAsset"]:
         return {
+            WarehouseAsset.ADDITIONAL_COLUMN_LINEAGE,
+            WarehouseAsset.ADDITIONAL_TABLE_LINEAGE,
             WarehouseAsset.EXTERNAL_COLUMN_LINEAGE,
             WarehouseAsset.EXTERNAL_TABLE_LINEAGE,
+            WarehouseAsset.FUNCTION,
         }
 
 
 class WarehouseAssetGroup(Enum):
     """Groups of assets that can be extracted together"""
 
+    ADDITIONAL_LINEAGE = "additional_lineage"
     CATALOG = "catalog"
     EXTERNAL_LINEAGE = "external_lineage"
+    FUNCTION = "function"
     QUERY = "query"
     ROLE = "role"
     SNOWFLAKE_LINEAGE = "snowflake_lineage"
@@ -53,6 +61,7 @@ CATALOG_ASSETS = (
 )
 
 # shared by technologies supporting queries
+FUNCTIONS_ASSETS = (WarehouseAsset.FUNCTION,)
 QUERIES_ASSETS = (WarehouseAsset.QUERY,)
 VIEWS_ASSETS = (WarehouseAsset.VIEW_DDL,)
 
@@ -61,6 +70,11 @@ EXTERNAL_LINEAGE_ASSETS = (
     WarehouseAsset.EXTERNAL_TABLE_LINEAGE,
 )
 
+ADDITIONAL_LINEAGE_ASSETS = (
+    WarehouseAsset.ADDITIONAL_COLUMN_LINEAGE,
+    WarehouseAsset.ADDITIONAL_TABLE_LINEAGE,
+)
+
 NON_EXTRACTABLE_ASSETS = {WarehouseAssetGroup.EXTERNAL_LINEAGE}