castor-extractor 0.16.1__tar.gz → 0.16.3__tar.gz

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 0.16.3 - 2024-04-24
+
+* Databricks: Extract table owners
+
+## 0.16.2 - 2024-04-09
+
+* PowerBI: Extract pages from report
+
 ## 0.16.1 - 2024-04-02
 
 * Systematically escape nul bytes on CSV write

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.16.1
+Version: 0.16.3
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -27,10 +27,9 @@ Provides-Extra: redshift
 Provides-Extra: snowflake
 Provides-Extra: sqlserver
 Provides-Extra: tableau
-Requires-Dist: click (>=8.1,<8.2)
 Requires-Dist: cryptography (>=41.0.5) ; extra == "snowflake"
 Requires-Dist: google-api-core (>=2.1.1,<3.0.0)
-Requires-Dist: google-auth (>=1.6.3,<3.0.0)
+Requires-Dist: google-auth (>=2,<3)
 Requires-Dist: google-cloud-core (>=2.1.0,<3.0.0)
 Requires-Dist: google-cloud-storage (>=2,<3)
 Requires-Dist: google-resumable-media (>=2.0.3,<3.0.0)

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/visualization/domo/client/client.py

@@ -14,7 +14,7 @@ RawData = Iterator[dict]
 
 DOMO_PUBLIC_URL = "https://api.domo.com"
 FORMAT = "%Y-%m-%d %I:%M:%S %p"
-DEFAULT_TIMEOUT = 300
+DEFAULT_TIMEOUT = 500
 TOKEN_EXPIRATION_SECONDS = timedelta(seconds=3000)  # auth token lasts 1 hour
 
 IGNORED_ERROR_CODES = (

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/visualization/powerbi/client/rest.py

@@ -64,19 +64,15 @@ class Client:
 
     def __init__(self, credentials: Credentials):
         self.creds = credentials
-
-    def _access_token(self) -> dict:
         client_app = f"{Urls.CLIENT_APP_BASE}{self.creds.tenant_id}"
-        app = msal.ConfidentialClientApplication(
+        self.app = msal.ConfidentialClientApplication(
             client_id=self.creds.client_id,
             authority=client_app,
             client_credential=self.creds.secret,
         )
 
-        token = app.acquire_token_silent(self.creds.scopes, account=None)
-
-        if not token:
-            token = app.acquire_token_for_client(scopes=self.creds.scopes)
+    def _access_token(self) -> dict:
+        token = self.app.acquire_token_for_client(scopes=self.creds.scopes)
 
         if Keys.ACCESS_TOKEN not in token:
             raise ValueError(f"No access token in token response: {token}")
@@ -248,7 +244,17 @@ class Client:
         Returns a list of reports for the organization.
         https://learn.microsoft.com/en-us/rest/api/power-bi/admin/reports-get-reports-as-admin
         """
-        return self._get(Urls.REPORTS)[Keys.VALUE]
+        reports = self._get(Urls.REPORTS)[Keys.VALUE]
+        for report in reports:
+            report_id = report.get("id")
+            try:
+                url = Urls.REPORTS + f"/{report_id}/pages"
+                pages = self._get(url)[Keys.VALUE]
+                report["pages"] = pages
+            except (requests.HTTPError, requests.exceptions.Timeout) as e:
+                logger.debug(e)
+                continue
+        return reports
 
     def _dashboards(self) -> List[Dict]:
         """

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/visualization/powerbi/client/rest_test.py

@@ -31,7 +31,6 @@ def test__access_token(mock_app):
     # init mocks
     valid_response = {"access_token": "mock_token"}
     returning_valid_token = Mock(return_value=valid_response)
-    mock_app.return_value.acquire_token_silent = Mock(return_value=None)
     mock_app.return_value.acquire_token_for_client = returning_valid_token
 
     # init client
@@ -40,30 +39,29 @@ def test__access_token(mock_app):
     # generated token
     assert client._access_token() == valid_response
 
-    # via silent endpoint token
-    mock_app.return_value.acquire_token_silent = returning_valid_token
-    mock_app.return_value.acquire_token_for_client = None
-    assert client._access_token() == valid_response
-
     # token missing in response
     invalid_response = {"not_access_token": "666"}
     returning_invalid_token = Mock(return_value=invalid_response)
-    mock_app.return_value.acquire_token_silent = returning_invalid_token
+    mock_app.return_value.acquire_token_for_client = returning_invalid_token
 
     with pytest.raises(ValueError):
         client._access_token()
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch.object(Client, "_access_token")
-def test__headers(mock_access_token):
+def test__headers(mock_access_token, mock_app):
+    mock_app.return_value = None
     client = _client()
     mock_access_token.return_value = {Keys.ACCESS_TOKEN: "666"}
     assert client._header() == {"Authorization": "Bearer 666"}
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__get(mocked_access_token, mocked_request):
+def test__get(mocked_access_token, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_access_token.return_value = {Keys.ACCESS_TOKEN: "666"}
     fact = {"fact": "Approximately 24 cat skins can make a coat.", "length": 43}
@@ -81,9 +79,11 @@ def test__get(mocked_access_token, mocked_request):
         result = client._get("https/whatev.er")
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__workspace_ids(_, mocked_request):
+def test__workspace_ids(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_request.return_value = Mock(
         json=lambda: [{"id": 1000}, {"id": 1001}, {"id": 1003}],
@@ -112,9 +112,11 @@ def test__workspace_ids(_, mocked_request):
     )
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__post_default(_, mocked_request):
+def test__post_default(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     url = "https://estcequecestbientotleweekend.fr/"
     params = QueryParams.METADATA_SCAN
@@ -129,9 +131,11 @@ def test__post_default(_, mocked_request):
     )
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__post_with_processor(_, mocked_request):
+def test__post_with_processor(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     url = "https://estcequecestbientotleweekend.fr/"
     params = QueryParams.METADATA_SCAN
@@ -146,9 +150,11 @@ def test__post_with_processor(_, mocked_request):
     assert result == 1000
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__datasets(_, mocked_request):
+def test__datasets(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_request.return_value = Mock(
         json=lambda: {"value": [{"id": 1, "type": "dataset"}]},
@@ -164,27 +170,50 @@ def test__datasets(_, mocked_request):
     assert datasets == [{"id": 1, "type": "dataset"}]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__reports(_, mocked_request):
+def test__reports(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
-    mocked_request.return_value = Mock(
-        json=lambda: {"value": [{"id": 1, "type": "report"}]},
-    )
+    page_url = f"{Urls.REPORTS}/1/pages"
+    calls = [
+        call(GET, Urls.REPORTS, data=None, headers=ANY, params=None),
+        call(
+            GET,
+            page_url,
+            data=None,
+            headers=ANY,
+            params=None,
+        ),
+    ]
+    mocked_request.side_effect = [
+        Mock(json=lambda: {"value": [{"id": 1, "type": "report"}]}),
+        Mock(
+            json=lambda: {
+                "value": [
+                    {"name": "page_name", "displayName": "page", "order": 0}
+                ]
+            }
+        ),
+    ]
     reports = client._reports()
-    mocked_request.assert_called_with(
-        GET,
-        Urls.REPORTS,
-        data=None,
-        headers=ANY,
-        params=None,
-    )
-    assert reports == [{"id": 1, "type": "report"}]
+    mocked_request.assert_has_calls(calls)
+
+    assert reports == [
+        {
+            "id": 1,
+            "type": "report",
+            "pages": [{"name": "page_name", "displayName": "page", "order": 0}],
+        }
+    ]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch("requests.request")
 @patch.object(Client, "_access_token")
-def test__dashboards(_, mocked_request):
+def test__dashboards(_, mocked_request, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked_request.return_value = Mock(
         json=lambda: {"value": [{"id": 1, "type": "dashboard"}]},
@@ -200,6 +229,7 @@ def test__dashboards(_, mocked_request):
     assert dashboards == [{"id": 1, "type": "dashboard"}]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch.object(Client, "_workspace_ids")
 @patch.object(Client, "_create_scan")
 @patch.object(Client, "_wait_for_scan_result")
@@ -209,7 +239,9 @@ def test__metadata(
     mocked_wait,
     mocked_create_scan,
     mocked_workspace_ids,
+    mock_app,
 ):
+    mock_app.return_value = None
     mocked_workspace_ids.return_value = list(range(200))
     mocked_create_scan.return_value = 314
     mocked_wait.return_value = True
@@ -240,8 +272,10 @@ _CALLS = [
 ]
 
 
+@patch.object(msal, "ConfidentialClientApplication")
 @patch.object(Client, "_call")
-def test__activity_events(mocked):
+def test__activity_events(mocked, mock_app):
+    mock_app.return_value = None
     client = _client()
     mocked.side_effect = _CALLS
 

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/visualization/tableau/assets.py

@@ -12,10 +12,12 @@ class TableauAsset(ExternalAsset):
     CUSTOM_SQL_TABLE = "custom_sql_tables"
     CUSTOM_SQL_QUERY = "custom_sql_queries"
     DASHBOARD = "dashboards"
+    DASHBOARD_SHEET = "dashboards_sheets"
     DATASOURCE = "datasources"
     FIELD = "fields"
     PROJECT = "projects"
     PUBLISHED_DATASOURCE = "published_datasources"
+    SHEET = "sheets"
     USAGE = "views"
     USER = "users"
     WORKBOOK = "workbooks"
@@ -25,7 +27,9 @@ class TableauAsset(ExternalAsset):
     def optional(cls) -> Set["TableauAsset"]:
         return {
             TableauAsset.DASHBOARD,
+            TableauAsset.DASHBOARD_SHEET,
             TableauAsset.FIELD,
+            TableauAsset.SHEET,
             TableauAsset.PUBLISHED_DATASOURCE,
         }
 
@@ -42,4 +46,5 @@ class TableauGraphqlAsset(Enum):
     DASHBOARD = "dashboards"
     DATASOURCE = "datasources"
     GROUP_FIELD = "groupFields"
+    SHEETS = "sheets"
     WORKBOOK_TO_DATASOURCE = "workbooks"

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/visualization/tableau/client/client.py

@@ -173,6 +173,13 @@ class ApiClient:
             TableauAsset.DASHBOARD,
         )
 
+    def _fetch_sheets(self) -> SerializedAsset:
+        """Fetches sheets"""
+
+        return self._fetch_paginated_objects(
+            TableauAsset.SHEET,
+        )
+
     def _fetch_paginated_objects(self, asset: TableauAsset) -> SerializedAsset:
         """Fetches paginated objects"""
 
@@ -203,6 +210,9 @@ class ApiClient:
         if asset == TableauAsset.PUBLISHED_DATASOURCE:
             assets = self._fetch_published_datasources()
 
+        if asset == TableauAsset.SHEET:
+            assets = self._fetch_sheets()
+
         if asset == TableauAsset.USAGE:
             assets = self._fetch_usages(self._safe_mode)
 

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/visualization/tableau/gql_fields.py

@@ -111,15 +111,15 @@ class GQLQueryFields(Enum):
     """
 
     DASHBOARDS: str = """
-    id
-    name
-    path
-    tags {
-        name
-    }
-    workbook {
-       luid # to retrieve the parent
-    }
+            id
+            name
+            path
+            tags {
+                name
+            }
+            workbook {
+               luid # to retrieve the parent
+            }
     """
 
     DATASOURCE: str = """
@@ -160,6 +160,21 @@ class GQLQueryFields(Enum):
            role
     """
 
+    SHEET: str = """
+            containedInDashboards {
+                id
+            }
+            id
+            index
+            name
+            upstreamFields{
+                name
+            }
+            workbook {
+                luid
+            }
+    """
+
     WORKBOOK_TO_DATASOURCE: str = """
            luid
            id
@@ -219,6 +234,12 @@ QUERY_FIELDS: Dict[TableauAsset, QueryInfo] = {
             OBJECT_TYPE: TableauGraphqlAsset.GROUP_FIELD,
         },
     ],
+    TableauAsset.SHEET: [
+        {
+            FIELDS: GQLQueryFields.SHEET,
+            OBJECT_TYPE: TableauGraphqlAsset.SHEETS,
+        },
+    ],
     TableauAsset.WORKBOOK_TO_DATASOURCE: [
         {
             FIELDS: GQLQueryFields.WORKBOOK_TO_DATASOURCE,

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/warehouse/databricks/client.py

@@ -87,15 +87,32 @@ class DatabricksClient(APIClient):
             content.get("tables", []), schema
         )
 
-    def tables_and_columns(self, schemas: List[dict]) -> TablesColumns:
+    @staticmethod
+    def _match_table_with_user(table: dict, user_id_by_email: dict) -> dict:
+        table_owner_email = table.get("owner_email")
+        if not table_owner_email:
+            return table
+        owner_external_id = user_id_by_email.get(table_owner_email)
+        if not owner_external_id:
+            return table
+        return {**table, "owner_external_id": owner_external_id}
+
+    def tables_and_columns(
+        self, schemas: List[dict], users: List[dict]
+    ) -> TablesColumns:
         """
         Get the databricks tables & columns leveraging the unity catalog API
         """
         tables: List[dict] = []
         columns: List[dict] = []
+        user_id_by_email = {user.get("email"): user.get("id") for user in users}
         for schema in schemas:
             t_to_add, c_to_add = self._tables_columns_of_schema(schema)
-            tables.extend(t_to_add)
+            t_with_owner = [
+                self._match_table_with_user(table, user_id_by_email)
+                for table in t_to_add
+            ]
+            tables.extend(t_with_owner)
             columns.extend(c_to_add)
         return tables, columns
 

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/warehouse/databricks/client_test.py

@@ -64,3 +64,17 @@ def test_DatabricksClient__keep_catalog():
     assert client._keep_catalog("staging")
     assert not client._keep_catalog("dev")
     assert not client._keep_catalog("something_unknown")
+
+
+def test_DatabricksClient__match_table_with_user():
+    client = MockDatabricksClient()
+    users_by_email = {"bob@castordoc.com": 3}
+
+    table = {"id": 1, "owner_email": "bob@castordoc.com"}
+    table_with_owner = client._match_table_with_user(table, users_by_email)
+
+    assert table_with_owner == {**table, "owner_external_id": 3}
+
+    table_without_owner = {"id": 1, "owner_email": None}
+    actual = client._match_table_with_user(table_without_owner, users_by_email)
+    assert actual == table_without_owner

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/warehouse/databricks/extract.py

@@ -82,7 +82,8 @@ class DatabricksExtractionProcessor:
 
         del databases
 
-        tables, columns = self._client.tables_and_columns(schemas)
+        users = self._client.users()
+        tables, columns = self._client.tables_and_columns(schemas, users)
 
         location = self._storage.put(WarehouseAsset.TABLE.value, tables)
         catalog_locations[WarehouseAsset.TABLE.value] = location

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/castor_extractor/warehouse/databricks/format.py

@@ -19,10 +19,11 @@ def _to_datetime_or_none(time_ms: Optional[int]) -> Optional[datetime]:
 
 def _table_payload(schema: dict, table: dict) -> dict:
     return {
+        "description": table.get("comment"),
         "id": table["table_id"],
+        "owner_email": table.get("owner"),
         "schema_id": f"{schema['id']}",
         "table_name": table["name"],
-        "description": table.get("comment"),
         "tags": [],
         "type": table.get("table_type"),
     }
@@ -30,12 +31,12 @@ def _table_payload(schema: dict, table: dict) -> dict:
 
 def _column_payload(table: dict, column: dict) -> dict:
     return {
-        "id": f"`{table['id']}`.`{column['name']}`",
         "column_name": column["name"],
-        "table_id": table["id"],
-        "description": column.get("comment"),
         "data_type": column["type_name"],
+        "description": column.get("comment"),
+        "id": f"`{table['id']}`.`{column['name']}`",
         "ordinal_position": column["position"],
+        "table_id": table["id"],
     }
 
 

{castor_extractor-0.16.1 → castor_extractor-0.16.3}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.16.1"
+version = "0.16.3"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"
@@ -14,16 +14,22 @@ classifiers = [
 readme = "README.md"
 homepage = "https://www.castordoc.com/"
 documentation = "https://docs.castordoc.com/castor-package/castor-extractor"
-include = ["README.md", "CHANGELOG.md", "LICENCE", "*/*/*/*.sql", "Dockerfile", "DockerfileUsage.md"]
+include = [
+  "README.md",
+  "CHANGELOG.md",
+  "LICENCE",
+  "*/*/*/*.sql",
+  "Dockerfile",
+  "DockerfileUsage.md",
+]
 
 
 [tool.poetry.dependencies]
 python = ">=3.8, <3.12"
 
-click = "~8.1"
 cryptography = { version = ">= 41.0.5", optional = true }
 google-api-core = ">=2.1.1, <3.0.0"
-google-auth = ">=1.6.3 , <3.0.0"
+google-auth = "^2"
 google-cloud-core = ">=2.1.0, <3.0.0"
 google-cloud-storage = "^2"
 google-resumable-media = "^2.0.3"
@@ -49,7 +55,7 @@ websocket-client = { version = "^1", optional = true }
 pydantic = "^2.6"
 pydantic-settings = "^2.2"
 pymssql = { version = "^2.2.11", optional = true }
-pymysql = {extras = ["rsa"], version = "^1.1.0", optional = true}
+pymysql = { extras = ["rsa"], version = "^1.1.0", optional = true }
 
 
 [tool.poetry.extras]
@@ -62,7 +68,11 @@ powerbi = ["msal"]
 qlik = ["websocket-client"]
 postgres = ["psycopg2-binary"]
 redshift = ["psycopg2-binary", "sqlalchemy-redshift"]
-snowflake = ["snowflake-connector-python", "snowflake-sqlalchemy", "cryptography"]
+snowflake = [
+  "snowflake-connector-python",
+  "snowflake-sqlalchemy",
+  "cryptography",
+]
 sqlserver = ["pymssql"]
 tableau = ["tableauserverclient"]
 all = [
@@ -89,7 +99,7 @@ pytest-freezegun = "^0.4"
 requests-mock = "^1.11"
 
 [tool.poetry.group.dev.dependencies]
-deptry = "^0.15.0"
+deptry = "^0.16"
 
 [tool.poetry.scripts]
 castor-extract-bigquery = "castor_extractor.commands.extract_bigquery:main"
@@ -115,16 +125,15 @@ castor-upload = "castor_extractor.commands.upload:main"
 [tool.poetry.group.poetry.dependencies]
 pip = "*"
 
+[[tool.poetry.source]]
+name = "PyPI"
+priority = "primary"
+
 [[tool.poetry.source]]
 name = "test-pypi"
 url = "https://test.pypi.org/legacy/"
 priority = "explicit"
 
-
-[[tool.poetry.source]]
-name = "PyPI"
-priority = "default"
-
 [tool.black]
 line-length = 80
 target-version = ["py38", "py39", "py310"]
@@ -218,7 +227,6 @@ pytest-freezegun = "freezegun"
 DEP004 = ["pytest", "requests_mock", "freezegun"]
 # SQL Alchemy needs some dependencies to be installed
 DEP002 = [
-  "click",
   "psycopg2-binary",
   "snowflake-connector-python",
   "snowflake-sqlalchemy",