castor-extractor 0.18.7__py3-none-any.whl → 0.19.0__py3-none-any.whl

CHANGELOG.md

@@ -1,5 +1,44 @@
 # Changelog
 
+## 0.19.0 - 2024-08-21
+
+* Breaking change Looker CLI: 
+
+`-u` and `--username` changed to `-c` `--client-id`
+
+`-p` and `--password` changed to `-s` `--client-secret`
+
+* Breaking change Metabase CLI: 
+
+`-u` and `--username` changed to `-u` `--user`
+
+
+* Note: if you use environment variables you shouldn't be impacted
+
+## 0.18.13 - 2024-08-19
+
+* Qlik: improve measures extraction
+
+## 0.18.12 - 2024-08-19
+
+* Bumps: looker-sdk: 24, setuptools: 72
+
+## 0.18.11 - 2023-08-14
+
+* Linting and formatting with Ruff
+
+## 0.18.10 - 2024-08-13
+
+* Soda: Added Soda extractor (Castor-Managed integration only)
+
+## 0.18.9 - 2024-08-12
+
+* Notion: Added Notion extractor
+
+## 0.18.8 - 2024-08-02
+
+* Databricks: more reliable extraction of queries
+
 ## 0.18.7 - 2024-08-01
 
 * Salesforce: extract table descriptions
@@ -420,7 +459,7 @@
 
 ## 0.2.2 - 2023-03-13
 
-* Constrain `setuptools` explicitly added
+* Constraint `setuptools` explicitly added
 
 ## 0.2.1 - 2023-02-23
 

castor_extractor/commands/extract_looker.py

@@ -7,8 +7,8 @@ from castor_extractor.visualization import looker  # type: ignore
 def main():
     parser = ArgumentParser()
     parser.add_argument("-b", "--base-url", help="Looker base url")
-    parser.add_argument("-u", "--username", help="Looker client id")
-    parser.add_argument("-p", "--password", help="Looker client secret")
+    parser.add_argument("-c", "--client-id", help="Looker client id")
+    parser.add_argument("-s", "--client-secret", help="Looker client secret")
     parser.add_argument("-o", "--output", help="Directory to write to")
     parser.add_argument("-t", "--timeout", type=int, help="Timeout in seconds")
     parser.add_argument(
@@ -18,7 +18,7 @@ def main():
     )
     parser.add_argument(
         "--safe-mode",
-        "-s",
+        "-S",
         help="Looker safe mode",
         action="store_true",
     )

castor_extractor/commands/extract_metabase_api.py

@@ -11,7 +11,7 @@ def main():
     parser = ArgumentParser()
 
     parser.add_argument("-b", "--base-url", help="Metabase base url")
-    parser.add_argument("-u", "--username", help="Metabase username")
+    parser.add_argument("-u", "--user", help="Metabase username")
     parser.add_argument("-p", "--password", help="Metabase password")
 
     parser.add_argument("-o", "--output", help="Directory to write to")

castor_extractor/commands/extract_metabase_db.py

@@ -19,7 +19,7 @@ def main():
     parser.add_argument("-P", "--port", help="TCP/IP port number")
     parser.add_argument("-d", "--database", help="Database name")
     parser.add_argument("-s", "--schema", help="Schema name")
-    parser.add_argument("-u", "--username", help="Username")
+    parser.add_argument("-u", "--user", help="Username")
     parser.add_argument("-p", "--password", help="Password")
     parser.add_argument(
         "-k",

castor_extractor/commands/extract_notion.py

@@ -0,0 +1,16 @@
+import logging
+from argparse import ArgumentParser
+
+from castor_extractor.knowledge import notion  # type: ignore
+from castor_extractor.utils import parse_filled_arguments  # type: ignore
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+
+
+def main():
+    parser = ArgumentParser()
+
+    parser.add_argument("-t", "--token", help="Notion token")
+    parser.add_argument("-o", "--output", help="Directory to write to")
+
+    notion.extract_all(**parse_filled_arguments(parser))

castor_extractor/commands/file_check.py

@@ -4,8 +4,11 @@ from argparse import ArgumentParser
 from typing import Set
 
 from castor_extractor import file_checker  # type: ignore
-from castor_extractor.utils import LocalStorage  # type: ignore
-from castor_extractor.utils import explode, search_files  # type: ignore
+from castor_extractor.utils import (  # type: ignore
+    LocalStorage,  # type: ignore
+    explode,
+    search_files,
+)
 
 logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)

castor_extractor/commands/upload.py

@@ -1,9 +1,11 @@
 import argparse
 import logging
 
-from castor_extractor.uploader import FileType  # type: ignore
-from castor_extractor.uploader import upload  # type: ignore
-from castor_extractor.uploader import upload_manifest  # type: ignore
+from castor_extractor.uploader import (  # type: ignore
+    FileType,
+    upload,
+    upload_manifest,
+)
 
 FILE_TYPES = {FileType.QUALITY, FileType.VIZ, FileType.WAREHOUSE}
 

castor_extractor/knowledge/notion/__init__.py

@@ -0,0 +1,3 @@
+from .assets import NotionAsset
+from .client import NotionClient, NotionCredentials
+from .extract import extract_all

castor_extractor/knowledge/notion/assets.py

@@ -0,0 +1,9 @@
+from ...types import ExternalAsset
+
+
+class NotionAsset(ExternalAsset):
+    """Notion assets"""
+
+    DATABASES = "databases"
+    PAGES = "pages"
+    USERS = "users"

castor_extractor/knowledge/notion/client/__init__.py

@@ -0,0 +1,2 @@
+from .client import NotionClient
+from .credentials import NotionCredentials

castor_extractor/knowledge/notion/client/client.py

@@ -0,0 +1,145 @@
+import logging
+from http import HTTPStatus
+from typing import Dict, Iterator
+
+from pydantic import ValidationError
+
+from ....utils import RequestSafeMode, empty_iterator
+from ....utils.client.api import APIClient, HttpMethod
+from ..assets import NotionAsset
+from .constants import (
+    CASTOR_NOTION_USER_AGENT,
+    NOTION_BASE_URL,
+    NOTION_TIMEOUT_MS,
+)
+from .credentials import NotionCredentials
+from .endpoints import EndpointFactory
+from .pagination import PaginatedResponse
+
+logger = logging.getLogger("__name__")
+
+NOTION_VERSION = "2021-08-16"
+
+VOLUME_IGNORED = 10
+IGNORED_ERROR_CODES = (HTTPStatus.BAD_GATEWAY,)
+NOTION_SAFE_MODE = RequestSafeMode(
+    max_errors=VOLUME_IGNORED,
+    status_codes=IGNORED_ERROR_CODES,
+)
+
+
+def _search_filter(asset: str) -> Dict[str, Dict[str, str]]:
+    return {"filter": {"value": asset, "property": "object"}}
+
+
+class NotionClient(APIClient):
+    """Client fetching data from Notion"""
+
+    def __init__(
+        self,
+        credentials: NotionCredentials,
+        timeout: int = NOTION_TIMEOUT_MS,
+        base_url: str = NOTION_BASE_URL,
+        user_agent: str = CASTOR_NOTION_USER_AGENT,
+        safe_mode: RequestSafeMode = NOTION_SAFE_MODE,
+    ) -> None:
+        base_headers = {
+            "Notion-Version": NOTION_VERSION,
+            "User-Agent": user_agent,
+        }
+
+        super().__init__(
+            host=base_url,
+            token=credentials.token,
+            headers=base_headers,
+            timeout=timeout,
+            safe_mode=safe_mode,
+        )
+
+    def _request(
+        self, method: HttpMethod, endpoint: str, params: dict
+    ) -> Iterator[dict]:
+        """
+        API call to Notion:
+        - If result is paginated, yield all pages
+        - If not, yield only the response payload
+        """
+        built_url = self.build_url(self._host, endpoint)
+        response_payload = self._call(
+            method=method,
+            url=built_url,
+            params=params if params and method == "GET" else None,
+            data=params if params and method == "POST" else None,
+        )
+        try:
+            paginated_response = PaginatedResponse(**response_payload)
+            yield from paginated_response.results
+
+            if not paginated_response.has_more:
+                return empty_iterator()
+
+            yield from self._request(
+                method=method,
+                endpoint=endpoint,
+                params=paginated_response.start_cursor,
+            )
+
+        except ValidationError:
+            yield response_payload
+
+    def _page_listing(self) -> Iterator[dict]:
+        return self._request(
+            method="POST",
+            endpoint=EndpointFactory.search(),
+            params=_search_filter("page"),
+        )
+
+    def _blocks(self, block_id: str) -> Iterator[dict]:
+        return self._request("GET", EndpointFactory.blocks(block_id), {})
+
+    def databases(self) -> Iterator[dict]:
+        return self._request(
+            method="POST",
+            endpoint=EndpointFactory.search(),
+            params=_search_filter("database"),
+        )
+
+    def recursive_blocks(self, block_id: str) -> Iterator[dict]:
+        """Fetch recursively all children blocks of a given block or page"""
+        blocks = self._blocks(block_id)
+        for block in blocks:
+            if block["has_children"] and block.get("type") != "child_page":
+                children = self.recursive_blocks(block["id"])
+                block["child_blocks"] = list(children)
+
+            yield block
+
+    def pages(self) -> Iterator[dict]:
+        """Fetch all pages with its whole content"""
+        pages = list(self._page_listing())
+        logger.info(f"Extracting {len(pages)} pages ...")
+        for page in pages:
+            if page.get("object") == "database":
+                # Notion Search API filter for page doesn't work
+                continue
+            content = list(self.recursive_blocks(page["id"]))
+            page["child_blocks"] = content
+            yield page
+
+    def users(self) -> Iterator[dict]:
+        """Fetch all users"""
+        return self._request("GET", EndpointFactory.users(), {})
+
+    def fetch(self, asset: NotionAsset) -> Iterator[dict]:
+        """Returns the needed metadata for the queried asset"""
+        if asset == NotionAsset.PAGES:
+            yield from self.pages()
+
+        elif asset == NotionAsset.DATABASES:
+            yield from self.databases()
+
+        elif asset == NotionAsset.USERS:
+            yield from self.users()
+
+        else:
+            raise ValueError(f"This asset {asset} is unknown")

castor_extractor/knowledge/notion/client/client_test.py

@@ -0,0 +1,67 @@
+from unittest.mock import patch
+
+from .client import NotionClient
+from .credentials import NotionCredentials
+
+MOCK_PAGINATED_RESPONSE = {
+    "results": [{"result_id": 1}],
+    "next_cursor": "2",
+    "has_more": True,
+}
+MOCK_PAGINATED_RESPONSE_2 = {
+    "results": [{"result_id": 2}],
+    "next_cursor": None,
+    "has_more": False,
+}
+MOCK_RESPONSE = {"result_id": 3}
+
+
+@patch.object(NotionClient, "_call")
+def test_NotionClient__request(mock_call):
+    mock_call.side_effect = [
+        MOCK_PAGINATED_RESPONSE,
+        MOCK_PAGINATED_RESPONSE_2,
+        MOCK_RESPONSE,
+    ]
+
+    client = NotionClient(NotionCredentials(token="MockToken"))
+    response = list(client._request("GET", "fake_endpoint", {}))
+    assert response == [{"result_id": 1}, {"result_id": 2}]
+
+    response = list(client._request("GET", "fake_endpoint", {}))
+    assert response == [{"result_id": 3}]
+
+
+MOCK_BLOCK = [{"object": "block", "id": "1", "has_children": True}]
+MOCK_BLOCK_CHILDREN = [
+    {"object": "block", "id": "2", "has_children": False},
+    {"object": "block", "id": "3", "has_children": True, "type": "child_page"},
+]
+
+
+@patch.object(NotionClient, "_blocks")
+def test_NotionClient__recursive_blocks(mock_blocks):
+    mock_blocks.side_effect = [
+        MOCK_BLOCK,
+        MOCK_BLOCK_CHILDREN,
+    ]
+
+    client = NotionClient(NotionCredentials(token="MockToken"))
+    response = list(client.recursive_blocks("1"))
+
+    assert response == [
+        {
+            "object": "block",
+            "id": "1",
+            "has_children": True,
+            "child_blocks": [
+                {"object": "block", "id": "2", "has_children": False},
+                {
+                    "object": "block",
+                    "id": "3",
+                    "has_children": True,
+                    "type": "child_page",
+                },
+            ],
+        }
+    ]

castor_extractor/knowledge/notion/client/constants.py

@@ -0,0 +1,3 @@
+NOTION_BASE_URL = "https://api.notion.com/v1/"
+CASTOR_NOTION_USER_AGENT = "castor-extractor"
+NOTION_TIMEOUT_MS = 300

castor_extractor/knowledge/notion/client/credentials.py

@@ -0,0 +1,16 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+CASTOR_ENV_PREFIX = "CASTOR_NOTION_"
+
+
+class NotionCredentials(BaseSettings):
+    """Class to handle Notion rest API permissions"""
+
+    model_config = SettingsConfigDict(
+        env_prefix=CASTOR_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    token: str = Field(repr=False)

castor_extractor/knowledge/notion/client/endpoints.py

@@ -0,0 +1,18 @@
+class EndpointFactory:
+    """Wrapper class around all endpoints we're using"""
+
+    BLOCKS = "blocks"
+    SEARCH = "search"
+    USERS = "users"
+
+    @classmethod
+    def users(cls) -> str:
+        return cls.USERS
+
+    @classmethod
+    def blocks(cls, block_id: str) -> str:
+        return f"{cls.BLOCKS}/{block_id}/children"
+
+    @classmethod
+    def search(cls) -> str:
+        return cls.SEARCH

castor_extractor/knowledge/notion/client/pagination.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+from pydantic.dataclasses import dataclass
+
+
+@dataclass
+class PaginatedResponse:
+    """Class to handle paginated results"""
+
+    results: list
+    next_cursor: Optional[str]
+    has_more: bool
+
+    @property
+    def start_cursor(self) -> dict:
+        return {"start_cursor": self.next_cursor}

castor_extractor/knowledge/notion/extract.py

@@ -0,0 +1,59 @@
+import logging
+from typing import Iterable, Iterator, Tuple, Union
+
+from ...utils import (
+    OUTPUT_DIR,
+    current_timestamp,
+    deep_serialize,
+    from_env,
+    get_output_filename,
+    write_json,
+    write_summary,
+)
+from .assets import NotionAsset
+from .client import NotionClient, NotionCredentials
+
+logger = logging.getLogger(__name__)
+
+
+def iterate_all_data(
+    client: NotionClient,
+) -> Iterable[Tuple[NotionAsset, Union[list, Iterator, dict]]]:
+    """Iterate over the extracted data from Notion"""
+
+    logger.info("Extracting USERS from API")
+    users = list(deep_serialize(client.users()))
+    yield NotionAsset.USERS, users
+    logger.info(f"Extracted {len(users)} users from API")
+
+    logger.info("Extracting PAGES from API")
+    pages = list(deep_serialize(client.pages()))
+    yield NotionAsset.PAGES, pages
+    logger.info(f"Extracted {len(pages)} pages from API")
+
+    logger.info("Extracting DATABASES from API")
+    databases = list(deep_serialize(client.databases()))
+    yield NotionAsset.DATABASES, databases
+    logger.info(f"Extracted {len(databases)} databases from API")
+
+
+def extract_all(**kwargs) -> None:
+    """
+    Extract data from Notion API
+    Store the output files locally under the given output_directory
+    """
+
+    output_directory = kwargs.get("output") or from_env(OUTPUT_DIR)
+
+    credentials = NotionCredentials(**kwargs)
+    client = NotionClient(credentials=credentials)
+
+    ts = current_timestamp()
+
+    for key, data in iterate_all_data(client):
+        filename = get_output_filename(key.name.lower(), output_directory, ts)
+        write_json(filename, data)
+
+    es = current_timestamp()
+
+    write_summary(output_directory, ts, duration_second=es - ts)

castor_extractor/quality/soda/__init__.py

@@ -0,0 +1,2 @@
+from .assets import SodaAsset
+from .client import SodaClient, SodaCredentials

castor_extractor/quality/soda/assets.py

@@ -0,0 +1,8 @@
+from ...types import ExternalAsset
+
+
+class SodaAsset(ExternalAsset):
+    """Soda assets"""
+
+    CHECK_RESULTS = "check_results"
+    DATASETS = "datasets"

castor_extractor/quality/soda/client/__init__.py

@@ -0,0 +1 @@
+from .client import SodaClient, SodaCredentials

castor_extractor/quality/soda/client/client.py

@@ -0,0 +1,99 @@
+from time import sleep
+from typing import Any, Iterator, Optional
+
+import requests
+
+from ....utils import format_date, yesterday
+from ....utils.client.api import DEFAULT_TIMEOUT_S, HttpMethod
+from ..assets import SodaAsset
+from .credentials import SodaCredentials
+from .endpoints import SodaEndpointFactory
+
+_REPORTING_PAGE_SIZE = 400
+_CLOUD_PAGE_SIZE = 100
+_REQUESTS_PER_MINUTE = 10
+_SECONDS_PER_MINUTE = 60
+_RATE_LIMIT_S = (_SECONDS_PER_MINUTE // _REQUESTS_PER_MINUTE) + 1
+
+
+class SodaClient:
+    def __init__(
+        self,
+        credentials: SodaCredentials,
+    ):
+        self._timeout = DEFAULT_TIMEOUT_S
+        self._reporting_headers = credentials.reporting_headers
+        self._auth = (credentials.api_key, credentials.secret)
+
+    def _call(
+        self,
+        url: str,
+        headers: dict,
+        method: HttpMethod = "GET",
+        *,
+        params: Optional[dict] = None,
+        data: Optional[dict] = None,
+        auth: Optional[tuple] = None,
+    ) -> Any:
+        response = requests.request(
+            method,
+            url,
+            headers=headers,
+            params=params,
+            json=data,
+            timeout=self._timeout,
+            auth=auth,
+        )
+        response.raise_for_status()
+
+        return response.json()
+
+    def _get_results_paginated(self, url: str, additional: dict) -> Iterator:
+        page_number = 1
+        next_page = True
+        while next_page:
+            json_data = {
+                **{"page": page_number, "size": _REPORTING_PAGE_SIZE},
+                **additional,
+            }
+            _check_results_page = self._call(
+                url=url,
+                method="POST",
+                data=json_data,
+                headers=self._reporting_headers,
+            )
+            yield from _check_results_page["data"]
+
+            next_page = len(_check_results_page["data"]) == _REPORTING_PAGE_SIZE
+            page_number += 1
+
+    def _datasets(self) -> Iterator[dict]:
+        url = SodaEndpointFactory.datasets()
+        next_page = True
+        page_number = 0
+        while next_page:
+            data = self._call(
+                url=url,
+                method="GET",
+                headers={},
+                params={"size": _CLOUD_PAGE_SIZE, "page": page_number},
+                auth=self._auth,
+            )
+            yield from data["content"]
+            next_page = not data["last"]
+            page_number += 1
+            sleep(_RATE_LIMIT_S)
+
+    def _check_results(self) -> Iterator:
+        url = SodaEndpointFactory.check_results()
+        _date = format_date(timestamp=yesterday())
+        return self._get_results_paginated(
+            url, additional={"from_datetime": _date}
+        )
+
+    def fetch(self, asset: SodaAsset) -> Iterator[dict]:
+        if asset == SodaAsset.DATASETS:
+            yield from self._datasets()
+        if asset == SodaAsset.CHECK_RESULTS:
+            yield from self._check_results()
+        raise ValueError(f"The asset {asset}, is not supported")

castor_extractor/quality/soda/client/credentials.py

@@ -0,0 +1,28 @@
+from typing import Dict
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from requests.auth import HTTPBasicAuth
+
+SODA_ENV_PREFIX = "CASTOR_SODA_"
+
+
+class SodaCredentials(BaseSettings):
+    """Class to handle Soda rest API permissions"""
+
+    model_config = SettingsConfigDict(
+        env_prefix=SODA_ENV_PREFIX,
+        extra="ignore",
+        populate_by_name=True,
+    )
+
+    api_key: str = Field(repr=False)
+    secret: str = Field(repr=False)
+
+    @property
+    def reporting_headers(self) -> Dict[str, str]:
+        return {
+            "Content-Type": "application/json",
+            "API_KEY_ID": self.api_key,
+            "API_KEY_SECRET": self.secret,
+        }

castor_extractor/quality/soda/client/endpoints.py

@@ -0,0 +1,13 @@
+class SodaEndpointFactory:
+    """Wrapper class around all endpoints we're using"""
+
+    CLOUD_API = "https://cloud.soda.io/api/v1"
+    REPORTING_API = "https://reporting.cloud.soda.io/v1"
+
+    @classmethod
+    def datasets(cls) -> str:
+        return f"{cls.CLOUD_API}/datasets"
+
+    @classmethod
+    def check_results(cls) -> str:
+        return f"{cls.REPORTING_API}/quality/check_results"

castor_extractor/types.py

@@ -50,6 +50,4 @@ class ExternalAsset(Enum):
         """
         Returns the assets that must always be provided.
         """
-        return {
-            item for item in cls if item not in cls.optional  # type: ignore
-        }
+        return {item for item in cls if item not in cls.optional}  # type: ignore

castor_extractor/uploader/upload.py

@@ -32,7 +32,6 @@ def _path_and_url(
     file_type: FileType,
     file_path: str,
 ) -> Tuple[str, str]:
-
     now = datetime.utcnow()
     timestamp = int(now.timestamp())
     filename = ntpath.basename(file_path)