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

CHANGELOG.md

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

castor_extractor/uploader/upload.py

@@ -83,7 +83,7 @@ def _upload(
             response = requests.post(
                 url=url,
                 headers=headers,
-                data=file_content,
+                files={"file": file_content},
                 timeout=timeout,
             )
             response.raise_for_status()

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

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

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

@@ -1,4 +1,4 @@
-CHANGELOG.md,sha256=EVZ9vhIVN7HLn5PYkRyBWyT3hk72Nt3i1SghwSipfR4,10957
+CHANGELOG.md,sha256=qR-os6GsyPBBkE3MwBCB4nYlMULY-D0vqKLXDcamyMU,11078
 Dockerfile,sha256=HcX5z8OpeSvkScQsN-Y7CNMUig_UB6vTMDl7uqzuLGE,303
 LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
 README.md,sha256=uF6PXm9ocPITlKVSh9afTakHmpLx3TvawLf-CbMP3wM,3578
@@ -41,13 +41,13 @@ castor_extractor/uploader/__init__.py,sha256=SSRtwjg-dNoxME-RJy9G1flASiUKAC5bH1h
 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=bTWD1_-hmJ6q1qcEosjZ96wsBtWDnWoCt692NYX_Nko,3228
+castor_extractor/uploader/upload.py,sha256=W1TGqO8_PtFdR661qNlh6v-LOIRvoJoda65-5OujFXs,3239
 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
-castor_extractor/utils/client/api.py,sha256=z1o4fteWx1HxNTqCYihl9sGkIgSQTbd8lW_B9Y2wyeQ,1742
+castor_extractor/utils/client/api.py,sha256=AGDj2JH__Q_x7RQdodoVazGvjGQJ9TzNqs-XEX6Hrms,1840
 castor_extractor/utils/client/api_test.py,sha256=NSMdXg1FLc37erqHp2FZsIsogWVv6lFSs7rDXHikr-E,542
 castor_extractor/utils/client/postgres.py,sha256=n6ulaT222WWPY0_6qAZ0MHF0m91HtI9mMqL71nyygo0,866
 castor_extractor/utils/client/query.py,sha256=O6D5EjD1KmBlwa786Uw4D4kzxx97_HH50xIIeSWt0B8,205
@@ -277,12 +277,12 @@ castor_extractor/warehouse/bigquery/queries/view_ddl.sql,sha256=obCm-IN9V8_YSZTw
 castor_extractor/warehouse/bigquery/query.py,sha256=hrFfjd5jW2oQnZ6ozlkn-gDe6sCIzu5zSX19T9W6fIk,4162
 castor_extractor/warehouse/bigquery/types.py,sha256=LZVWSmE57lOemNbB5hBRyYmDk9bFAU4nbRaJWALl6N8,140
 castor_extractor/warehouse/databricks/__init__.py,sha256=bTvDxjGQGM2J3hOnVhfNmFP1y8DK0tySiD_EXe5_xWE,200
-castor_extractor/warehouse/databricks/client.py,sha256=oHR_htE25p5tiAAFZKbF48efo7tqIENW4dAGA7yEqHg,16895
+castor_extractor/warehouse/databricks/client.py,sha256=sNY-7FDg9nLaqw2zk_aoGAhSGi8KST9QpHAHn46700w,17439
 castor_extractor/warehouse/databricks/client_test.py,sha256=KNp4Hi_CC6GwiW2QDJQQwqALfUebuT9D_qL6FuP_8tY,5246
 castor_extractor/warehouse/databricks/credentials.py,sha256=PpGv5_GP320UQjV_gvaxSpOw58AmqSznmjGhGfe6bdU,655
 castor_extractor/warehouse/databricks/extract.py,sha256=VX-3uo5dZucenrg-wnPur3CxOgpC5H7Ds92TO7OTAjc,7379
-castor_extractor/warehouse/databricks/format.py,sha256=2bRy2fa45NW3uk030rmyba4n2Em-NnyZPBurUslEbcw,5522
-castor_extractor/warehouse/databricks/format_test.py,sha256=iPmdJof43fBYL1Sa_fBrCWDQHCHgm7IWCZag1kWkj9E,1970
+castor_extractor/warehouse/databricks/format.py,sha256=zSO3Cm-vpidzNA07W81I506u-ToQzkjXVwKDmS-tfiE,6088
+castor_extractor/warehouse/databricks/format_test.py,sha256=HZvJjcB7sj7LF1kIxAeaf_KdD3XOKn9nfeQLRT39G3s,2804
 castor_extractor/warehouse/databricks/test_constants.py,sha256=Hm96yq_ltVAKv7WYhYz637r4Cuj-1cCdyOuxMEe3J-Q,2246
 castor_extractor/warehouse/databricks/types.py,sha256=hD6gC8oiT3QSWEvbtgUOGK_lLzzz36sEauB3lS_wxlE,218
 castor_extractor/warehouse/mysql/__init__.py,sha256=2KFDogo9GNbApHqw3Vm5t_uNmIRjdp76nmP_WQQMfQY,116
@@ -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.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,,
+castor_extractor-0.17.2.dist-info/LICENCE,sha256=sL-IGa4hweyya1HgzMskrRdybbIa2cktzxb5qmUgDg8,8254
+castor_extractor-0.17.2.dist-info/METADATA,sha256=V-1vK-HPqqZxUKCRncvmp__Yg4UA3B1Aza7Ac4uVFrA,6582
+castor_extractor-0.17.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+castor_extractor-0.17.2.dist-info/entry_points.txt,sha256=SbyPk58Gh-FRztfCNnUZQ6w7SatzNJFZ6GIJLNsy7tI,1427
+castor_extractor-0.17.2.dist-info/RECORD,,