castor-extractor 0.16.15__py3-none-any.whl → 0.17.0__py3-none-any.whl

CHANGELOG.md

@@ -1,5 +1,9 @@
 # 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

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/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/uploader/upload.py

@@ -1,83 +1,100 @@
 #!/usr/bin/env python3
-import json
 import logging
 import ntpath
 from datetime import datetime
-from typing import Iterable, Optional, Union
+from typing import Dict, Iterable, Optional, Tuple
 from uuid import UUID
 
-from google.cloud import storage  # type: ignore
-
-from .constant import EXTRACTION_BUCKET, PATH_TEMPLATES, FileType
+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 file_exist, iter_files
+from .utils import iter_files
 
 logger = logging.getLogger(__name__)
 
+_EXCEPTIONS = (
+    requests.exceptions.Timeout,
+    requests.exceptions.ConnectTimeout,
+)
 
-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_and_url(
+    source_id: UUID,
+    file_type: FileType,
+    file_path: str,
+) -> Tuple[str, str]:
 
-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(
+    path = path_template.format(
         timestamp=timestamp,
         source_id=source_id,
         filename=filename,
     )
 
+    url = f"{INGEST_URL}/{path}"
 
-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)
+    return path, url
 
-    bucket = client.bucket(EXTRACTION_BUCKET)
-    return bucket.blob(path)
+
+def _headers(token: str) -> Dict:
+    return {
+        "Authorization": f"Token {token}",
+        "Accept": "text/csv, application/json",
+    }
 
 
 def _upload(
-    credentials: Union[str, dict],
+    token: str,
     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
+    Upload the given file to Google Cloud Storage (GCS)
+    - Don't call GCS API directly
+    - Call the ingestion proxy which handles authorisation and uploading
     """
-    timeout, retries = get_blob_env()
+    path, url = _path_and_url(source_id, file_type, file_path)
+    headers = _headers(token)
+    timeout, max_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}",
-    )
+    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(
-    credentials: Union[str, dict],
+    token: str,
     source_id: UUID,
     file_path: str,
 ) -> None:
@@ -86,11 +103,11 @@ def upload_manifest(
     source_id: id for the source
     file_path: path to the local manifest to upload
     """
-    _upload(credentials, source_id, file_path, FileType.DBT)
+    _upload(token, source_id, file_path, FileType.DBT)
 
 
 def upload(
-    credentials: Union[str, dict],
+    token: str,
     source_id: UUID,
     file_type: FileType,
     file_path: Optional[str] = None,
@@ -113,4 +130,4 @@ def upload(
         raise ValueError(message)
 
     for file_ in files:
-        _upload(credentials, source_id, file_, file_type)
+        _upload(token, source_id, file_, file_type)

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/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/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.dist-info → castor_extractor-0.17.0.dist-info}/METADATA

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

{castor_extractor-0.16.15.dist-info → castor_extractor-0.17.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-CHANGELOG.md,sha256=QYFobUPMbdi6cidq_yU-oMbXWoAr1BjTE6thfdZ9tA4,10866
+CHANGELOG.md,sha256=EVZ9vhIVN7HLn5PYkRyBWyT3hk72Nt3i1SghwSipfR4,10957
 Dockerfile,sha256=HcX5z8OpeSvkScQsN-Y7CNMUig_UB6vTMDl7uqzuLGE,303
 LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
 README.md,sha256=uF6PXm9ocPITlKVSh9afTakHmpLx3TvawLf-CbMP3wM,3578
@@ -23,7 +23,7 @@ castor_extractor/commands/extract_snowflake.py,sha256=vYiruxRoo--GeMemOGsSE1w9kc
 castor_extractor/commands/extract_sqlserver.py,sha256=lwhbcNChaXHZgMgSOch3faVr7WJw-sDU6GHl3lzBt_0,1141
 castor_extractor/commands/extract_tableau.py,sha256=u-6UCd-kfXwyhNWYxZusqtgTTYkf4gAJS1vRIYWsAVU,1415
 castor_extractor/commands/file_check.py,sha256=PkXFK-kPoG8PpvBP-PCvVdreqwYw6Y1fTh2QzDxACsM,2684
-castor_extractor/commands/upload.py,sha256=tWN9hUn1aKJrGEmNHI_qjSciFiAoL9x7YolbIsYbg8Y,1956
+castor_extractor/commands/upload.py,sha256=tAyHfIDOdUpD0yMJe2f64nXpaVnIbgYXi4bYx2nNvNU,1922
 castor_extractor/file_checker/__init__.py,sha256=OSt6YLhUT42U_Cp3LCLHMVruwDkksL75Ij13X2UPnVk,119
 castor_extractor/file_checker/column.py,sha256=fMchy5v-Sd-0xuYS0V9mob7wnljslzWLhQGqrKGybdk,3097
 castor_extractor/file_checker/column_test.py,sha256=1j8PxvmvmJgpd-mk30iMYOme32ovPSIn4yCXywFoXrg,1935
@@ -38,12 +38,12 @@ castor_extractor/file_checker/templates/generic_warehouse.py,sha256=zvnWnYB8FNvh
 castor_extractor/logger.py,sha256=ovf1mBEKwbJBskBXoqHbcAomBrp58mUwSrCWtEMlYPM,1197
 castor_extractor/types.py,sha256=-QgiOaq--nXUsYLy_oESDrYbRMxs353-YiQnG1blJvU,1303
 castor_extractor/uploader/__init__.py,sha256=SSRtwjg-dNoxME-RJy9G1flASiUKAC5bH1htq3CURQg,75
-castor_extractor/uploader/constant.py,sha256=hEJlWYx0dyBzgo59XUBKCYIKEODpIc2DyzwAZIiNO8g,718
+castor_extractor/uploader/constant.py,sha256=yTigLHDlYwoRr6CpFIl7ReElFsQd4H-qkluMZJPWSx0,865
 castor_extractor/uploader/env.py,sha256=5HSniVSOYVg4u38O4k8TB_qaJq9s8yJ1hjedkq_gdVg,878
 castor_extractor/uploader/env_test.py,sha256=ClCWWtwd2N-5ClIDUxVMeKkWfhhOTxpppsXUDmdjxSg,472
-castor_extractor/uploader/upload.py,sha256=5Aj3UOx8cpSVvzjYRz7S6nLk249IqUiCia70utU_970,3363
-castor_extractor/uploader/upload_test.py,sha256=BfGjAYEEDBmEcUS6_b3SlKyiQNR1iRf6-qmADDirTJI,328
-castor_extractor/uploader/utils.py,sha256=NCe0tkB28BVhqzOaDhDjaSfODjjcPWB17X6chnvyCWs,478
+castor_extractor/uploader/upload.py,sha256=bTWD1_-hmJ6q1qcEosjZ96wsBtWDnWoCt692NYX_Nko,3228
+castor_extractor/uploader/upload_test.py,sha256=7fwstdQe7FjuwGilsCdFpEQr1qLoR2WTRUzyy93fISw,402
+castor_extractor/uploader/utils.py,sha256=Tx_i875L2vJ8btOLV3-L0UMEFiyhH8E5n0XXRyLjO0Y,793
 castor_extractor/utils/__init__.py,sha256=bmzAOc-PKsVreMJtF7DGpPQeHrVqxWel_BblRftt6Ag,1186
 castor_extractor/utils/client/__init__.py,sha256=CRE-xJKm6fVV9dB8ljzB5YoOxX4I1sCD1KSgqs3Y8_Y,161
 castor_extractor/utils/client/abstract.py,sha256=aA5Qcb9TwWDSMq8WpXbGkOB20hehwX2VTpqQAwV76wk,2048
@@ -253,7 +253,7 @@ castor_extractor/visualization/tableau_revamp/constants.py,sha256=PcdudAogQhi3e-
 castor_extractor/visualization/tableau_revamp/extract.py,sha256=2SLUxp5okM4AcEJJ61ZgcC2ikfZZl9MH17CEXMXmgl0,1450
 castor_extractor/warehouse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 castor_extractor/warehouse/abstract/__init__.py,sha256=Fdfa026tgOo64MvzVRLHM_F2G-JmcehrF0mh3dHgb7s,419
-castor_extractor/warehouse/abstract/asset.py,sha256=9nHL4WKUU_vRgj7u3sUdIzgI4rRpdS7YrfwNku4Gz9Q,2652
+castor_extractor/warehouse/abstract/asset.py,sha256=Qs7T2Iw7KHgWVT2aAoBfCQ8tB143cUZY-DRUSkpgvGU,2689
 castor_extractor/warehouse/abstract/asset_test.py,sha256=_kd4ybNlWSAdSdEgJKC-jhJTa1nMRa9i8RO3YbqKLM4,758
 castor_extractor/warehouse/abstract/extract.py,sha256=fVBhdE-yMI_g6RBYZcr7q-ZVW7jK7WVkO_GO_KfkRqg,2908
 castor_extractor/warehouse/abstract/query.py,sha256=GAgeISCmAdrkTKzFGO79hQDf6SA6EFrrlW43w-LiXKo,2632
@@ -370,8 +370,8 @@ castor_extractor/warehouse/synapse/queries/schema.sql,sha256=aX9xNrBD_ydwl-znGSF
 castor_extractor/warehouse/synapse/queries/table.sql,sha256=mCE8bR1Vb7j7SwZW2gafcXidQ2fo1HwxcybA8wP2Kfs,1049
 castor_extractor/warehouse/synapse/queries/user.sql,sha256=sTb_SS7Zj3AXW1SggKPLNMCd0qoTpL7XI_BJRMaEpBg,67
 castor_extractor/warehouse/synapse/queries/view_ddl.sql,sha256=3EVbp5_yTgdByHFIPLHmnoOnqqLE77SrjAwFDvu4e54,249
-castor_extractor-0.16.15.dist-info/LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
-castor_extractor-0.16.15.dist-info/METADATA,sha256=CsdtS6LQFjsgi0A7tj0sMwtkQVYBye4Savn2DFGBHso,6583
-castor_extractor-0.16.15.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-castor_extractor-0.16.15.dist-info/entry_points.txt,sha256=SbyPk58Gh-FRztfCNnUZQ6w7SatzNJFZ6GIJLNsy7tI,1427
-castor_extractor-0.16.15.dist-info/RECORD,,
+castor_extractor-0.17.0.dist-info/LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
+castor_extractor-0.17.0.dist-info/METADATA,sha256=mPiUyxCqXFifcPbhcOPFsnkPAV4OcWXoYzGeUKlbkoo,6582
+castor_extractor-0.17.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+castor_extractor-0.17.0.dist-info/entry_points.txt,sha256=SbyPk58Gh-FRztfCNnUZQ6w7SatzNJFZ6GIJLNsy7tI,1427
+castor_extractor-0.17.0.dist-info/RECORD,,