castor-extractor 0.16.15__tar.gz → 0.17.2__tar.gz

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## 0.17.2 - 2024-06-14
+
+* Uploader: support multipart
+
+## 0.17.1 - 2024-06-12
+
+* Databricks: extract table source links
+
+## 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

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/PKG-INFO

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

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/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.15 → castor_extractor-0.17.2}/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.2/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,
+                files={"file": 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.15 → castor_extractor-0.17.2}/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.15 → castor_extractor-0.17.2}/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.15 → castor_extractor-0.17.2}/castor_extractor/utils/client/api.py

@@ -58,7 +58,12 @@ class APIClient:
 
         return result.json()
 
-    def get(self, path: str, payload: Optional[dict] = None) -> dict:
+    def get(
+        self,
+        path: str,
+        payload: Optional[dict] = None,
+        processor: Optional[Callable] = None,
+    ) -> dict:
         """path: REST API operation path, such as /api/2.0/clusters/get"""
         url = self.build_url(self._host, path)
-        return self._call(url=url, data=payload)
+        return self._call(url=url, data=payload, processor=processor)

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/castor_extractor/warehouse/abstract/asset.py

@@ -32,6 +32,7 @@ class WarehouseAsset(ExternalAsset):
             WarehouseAsset.ADDITIONAL_TABLE_LINEAGE,
             WarehouseAsset.EXTERNAL_COLUMN_LINEAGE,
             WarehouseAsset.EXTERNAL_TABLE_LINEAGE,
+            WarehouseAsset.FUNCTION,
         }
 
 

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/castor_extractor/warehouse/databricks/client.py

@@ -5,6 +5,7 @@ from functools import partial
 from typing import Any, Dict, List, Optional, Set, Tuple, cast
 
 import requests
+from requests import Response
 
 from ...utils import (
     SafeMode,
@@ -30,6 +31,7 @@ _RETRY_BASE_MS = 1000
 _RETRY_EXCEPTIONS = [
     requests.exceptions.ConnectTimeout,
 ]
+_WORKSPACE_ID_HEADER = "X-Databricks-Org-Id"
 
 safe_params = SafeMode((BaseException,), _MAX_NUMBER_OF_LINEAGE_ERRORS)
 
@@ -120,15 +122,28 @@ class DatabricksClient(APIClient):
             for schema in self._schemas_of_database(database)
         ]
 
+    @staticmethod
+    def _process_table_response(response: Response) -> Tuple[dict, str]:
+        """
+        Returns both the JSON content and the Workspace ID, which is found
+        in the response's headers.
+        """
+        return response.json(), response.headers[_WORKSPACE_ID_HEADER]
+
     def _tables_columns_of_schema(self, schema: dict) -> TablesColumns:
         path = "api/2.1/unity-catalog/tables"
         payload = {
             "catalog_name": schema["database_id"],
             "schema_name": schema["schema_name"],
         }
-        content = self.get(path=path, payload=payload)
+        content, workspace_id = self.get(
+            path=path,
+            payload=payload,
+            processor=self._process_table_response,
+        )
+        host = self.build_url(self._host, path="")
         return self.formatter.format_table_column(
-            content.get("tables", []), schema
+            content.get("tables", []), schema, host, workspace_id
         )
 
     @staticmethod

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/castor_extractor/warehouse/databricks/format.py

@@ -9,6 +9,8 @@ logger = logging.getLogger(__name__)
 EXCLUDED_DATABASES = {"system"}
 EXCLUDED_SCHEMAS = {"information_schema", "default"}
 
+TABLE_URL_TPL = "{host}explore/data/{catalog_name}/{schema_name}/{table_name}?o={workspace_id}"
+
 
 def _to_datetime_or_none(time_ms: Optional[int]) -> Optional[datetime]:
     """return time in ms as datetime or None"""
@@ -17,7 +19,23 @@ def _to_datetime_or_none(time_ms: Optional[int]) -> Optional[datetime]:
     return datetime.fromtimestamp(time_ms / 1000.0)
 
 
-def _table_payload(schema: dict, table: dict) -> dict:
+def _table_payload(
+    schema: dict,
+    table: dict,
+    host: str,
+    workspace_id: str,
+) -> dict:
+    """
+    Prepares the table payload. This also includes a source link which is built
+    here using the host and workspace_id.
+    """
+    url = TABLE_URL_TPL.format(
+        host=host,
+        catalog_name=table["catalog_name"],
+        schema_name=table["schema_name"],
+        table_name=table["name"],
+        workspace_id=workspace_id,
+    )
     return {
         "description": table.get("comment"),
         "id": table["table_id"],
@@ -26,6 +44,7 @@ def _table_payload(schema: dict, table: dict) -> dict:
         "table_name": table["name"],
         "tags": [],
         "type": table.get("table_type"),
+        "url": url,
     }
 
 
@@ -78,14 +97,14 @@ class DatabricksFormatter:
 
     @staticmethod
     def format_table_column(
-        raw_tables: List[dict], schema: dict
+        raw_tables: List[dict], schema: dict, host: str, workspace_id: str
     ) -> TablesColumns:
         tables = []
         columns = []
         if not raw_tables:
             return [], []
         for table in raw_tables:
-            t = _table_payload(schema, table)
+            t = _table_payload(schema, table, host, workspace_id)
             tables.append(t)
             if not table.get("columns"):
                 continue

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/castor_extractor/warehouse/databricks/format_test.py

@@ -1,6 +1,11 @@
 from datetime import datetime
 
-from .format import DatabricksFormatter, _column_payload, _to_datetime_or_none
+from .format import (
+    DatabricksFormatter,
+    _column_payload,
+    _table_payload,
+    _to_datetime_or_none,
+)
 
 
 def test__to_datetime_or_none():
@@ -23,6 +28,35 @@ def test_DatabricksFormatter__primary():
     assert DatabricksFormatter._primary([]) is None
 
 
+def test__table_payload():
+    schema = {"id": "id123"}
+
+    table = {
+        "name": "baz",
+        "catalog_name": "foo",
+        "schema_name": "bar",
+        "table_type": "MANAGED",
+        "owner": "pot@ato.com",
+        "table_id": "732pot5e-8ato-4c27-b701-9fa51febc192",
+    }
+    host = "https://some.cloud.databricks.net/"
+    workspace_id = "123456"
+
+    payload = _table_payload(schema, table, host, workspace_id)
+
+    expected = {
+        "description": None,
+        "id": "732pot5e-8ato-4c27-b701-9fa51febc192",
+        "owner_email": "pot@ato.com",
+        "schema_id": "id123",
+        "table_name": "baz",
+        "tags": [],
+        "type": "MANAGED",
+        "url": "https://some.cloud.databricks.net/explore/data/foo/bar/baz?o=123456",
+    }
+    assert payload == expected
+
+
 def test__column_payload():
     table = {
         "id": "18175cd5-9b9b-4d78-9d28-caaa12c21ce0",

{castor_extractor-0.16.15 → castor_extractor-0.17.2}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.16.15"
+version = "0.17.2"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"

castor_extractor-0.16.15/castor_extractor/uploader/upload.py

@@ -1,116 +0,0 @@
-#!/usr/bin/env python3
-import json
-import logging
-import ntpath
-from datetime import datetime
-from typing import Iterable, Optional, Union
-from uuid import UUID
-
-from google.cloud import storage  # type: ignore
-
-from .constant import EXTRACTION_BUCKET, PATH_TEMPLATES, FileType
-from .env import get_blob_env
-from .utils import file_exist, iter_files
-
-logger = logging.getLogger(__name__)
-
-
-def _client(credentials: Union[str, dict]) -> storage.Client:
-    """supports dict, string or path to the JSON file"""
-    if isinstance(credentials, dict):
-        return storage.Client.from_service_account_info(credentials)
-    if file_exist(credentials):
-        return storage.Client.from_service_account_json(credentials)
-    if isinstance(credentials, str):
-        credentials = json.loads(credentials)
-        return storage.Client.from_service_account_info(credentials)
-    raise ValueError("needs path or dict for credentials")
-
-
-def _path(source_id: UUID, file_type: FileType, file_path: str) -> str:
-    now = datetime.utcnow()
-    timestamp = int(now.timestamp())
-    filename = ntpath.basename(file_path)
-
-    path_template = PATH_TEMPLATES[file_type]
-    return path_template.format(
-        timestamp=timestamp,
-        source_id=source_id,
-        filename=filename,
-    )
-
-
-def _get_blob(
-    credentials: Union[str, dict],
-    source_id: UUID,
-    file_path: str,
-    file_type: FileType,
-) -> storage.Blob:
-    """get the target blob to upload to"""
-    client = _client(credentials)
-    path = _path(source_id, file_type, file_path)
-
-    bucket = client.bucket(EXTRACTION_BUCKET)
-    return bucket.blob(path)
-
-
-def _upload(
-    credentials: Union[str, dict],
-    source_id: UUID,
-    file_path: str,
-    file_type: FileType,
-) -> None:
-    """
-    credentials: path to file or dict
-    source_id: id for the source
-    file_type: type of file to upload
-    file_path: path to the local file to upload
-    """
-    timeout, retries = get_blob_env()
-
-    blob = _get_blob(credentials, source_id, file_path, file_type)
-    with open(file_path, "rb") as f:
-        blob.upload_from_file(f, timeout=timeout, num_retries=retries)
-    logger.info(
-        f"uploaded {file_path} as {file_type.value} to {blob.public_url}",
-    )
-
-
-def upload_manifest(
-    credentials: Union[str, dict],
-    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(credentials, source_id, file_path, FileType.DBT)
-
-
-def upload(
-    credentials: Union[str, dict],
-    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(credentials, source_id, file_, file_type)