omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/notion/__init__.py

@@ -0,0 +1,69 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A source that extracts data from Notion API"""
+
+from typing import Dict, Iterator, List, Optional
+
+import dlt
+from dlt.sources import DltResource
+
+from .helpers.client import NotionClient
+from .helpers.database import NotionDatabase
+
+
+@dlt.source(max_table_nesting=1)
+def notion_databases(
+    database_ids: Optional[List[Dict[str, str]]] = None,
+    api_key: str = dlt.secrets.value,
+) -> Iterator[DltResource]:
+    """
+    Retrieves data from Notion databases.
+
+    Args:
+        database_ids (List[Dict[str, str]], optional): A list of dictionaries
+            each containing a database id and a name.
+            Defaults to None. If None, the function will generate all databases
+            in the workspace that are accessible to the integration.
+        api_key (str): The Notion API secret key.
+
+    Yields:
+        DltResource: Data resources from Notion databases.
+    """
+    notion_client = NotionClient(api_key)
+
+    if database_ids is None:
+        search_results = notion_client.search(
+            filter_criteria={"value": "database", "property": "object"}
+        )
+        database_ids = [
+            {"id": result["id"], "use_name": result["title"][0]["plain_text"]}
+            for result in search_results
+        ]
+
+    for database in database_ids:
+        if "use_name" not in database:
+            # Fetch the database details from Notion
+            details = notion_client.get_database(database["id"])
+
+            # Extract the name/title from the details
+            database["use_name"] = details["title"][0]["plain_text"]
+
+        notion_database = NotionDatabase(database["id"], notion_client)
+        yield dlt.resource(  # type: ignore
+            notion_database.query(),
+            primary_key="id",
+            name=database["use_name"],
+            write_disposition="replace",
+        )

omniload/src/notion/helpers/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+

omniload/src/notion/helpers/client.py

@@ -0,0 +1,178 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Iterator, Optional
+
+from dlt.sources.helpers import requests
+
+from ..settings import API_URL
+
+
+class NotionClient:
+    """A client to interact with the Notion API.
+
+    Attributes:
+        api_key (str): The Notion API secret key.
+    """
+
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key
+
+    def _create_headers(self) -> Dict[str, str]:
+        headers = {
+            "accept": "application/json",
+            "Notion-Version": "2022-06-28",
+            "Authorization": f"Bearer {self.api_key}",
+        }
+        return headers
+
+    def _filter_out_none_values(self, dict_in: Dict[str, Any]) -> Dict[str, Any]:
+        return {k: v for k, v in dict_in.items() if v is not None}
+
+    def get_endpoint(
+        self, resource: str, resource_id: str, subresource: Optional[str] = None
+    ) -> str:
+        """Returns the endpoint for a given resource.
+
+        Args:
+            resource (str): The resource to get the endpoint for.
+            resource_id (str): The id of the resource.
+            subresource (str, optional): The subresource to get the endpoint for.
+
+        Returns:
+            str: The endpoint for the resource.
+        """
+        url = f"{API_URL}/{resource}/{resource_id}"
+        if subresource:
+            url += f"/{subresource}"
+        return url
+
+    def fetch_resource(
+        self, resource: str, resource_id: str, subresource: Optional[str] = None
+    ) -> Any:
+        """Fetches a resource from the Notion API.
+
+        Args:
+            resource (str): The resource to fetch.
+            resource_id (str): The id of the resource.
+            subresource (str, optional): The subresource to fetch. Defaults to None.
+
+        Returns:
+            Any: The resource from the Notion API.
+        """
+        url = self.get_endpoint(resource, resource_id, subresource)
+        headers = self._create_headers()
+        response = requests.get(url, headers=headers)
+        response.raise_for_status()
+        return response.json()
+
+    def send_payload(
+        self,
+        resource: str,
+        resource_id: str,
+        subresource: Optional[str] = None,
+        query_params: Optional[Dict[str, Any]] = None,
+        payload: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """Sends a payload to the Notion API using the POST method.
+
+        Args:
+            resource (str): The resource to send the payload to.
+            resource_id (str): The id of the resource.
+            subresource (str, optional): The subresource to send the payload to.
+                Defaults to None.
+            query_params (Dict[str, Any], optional): The query parameters to send
+                with the payload. Defaults to None.
+            payload (Dict[str, Any], optional): The payload to send. Defaults to None.
+
+        Returns:
+            Any: The response from the Notion API.
+
+        Raises:
+            requests.HTTPError: If the response from the Notion API is not 200.
+        """
+
+        url = self.get_endpoint(resource, resource_id, subresource)
+        headers = self._create_headers()
+
+        if payload is None:
+            payload = {}
+
+        filtered_payload = self._filter_out_none_values(payload)
+
+        response = requests.post(
+            url, headers=headers, params=query_params, json=filtered_payload
+        )
+        response.raise_for_status()
+        return response.json()
+
+    def search(
+        self,
+        query: Optional[str] = None,
+        filter_criteria: Optional[Dict[str, Any]] = None,
+        sort: Optional[Dict[str, Any]] = None,
+        start_cursor: Optional[str] = None,
+        page_size: Optional[int] = None,
+    ) -> Iterator[Dict[str, Any]]:
+        """Searches all parent or child pages and databases that have been
+        shared with an integration.
+
+        Notion API Reference. Search:
+            https://developers.notion.com/reference/post-search
+
+        Args:
+            query (str, optional): The string to search for. Defaults to None.
+            filter_criteria (Dict[str, Any], optional): The filter to apply to
+                the results.
+            sort (Dict[str, Any], optional): The sort to apply to the results.
+            start_cursor (str, optional): The cursor to start the query at.
+                Defaults to None.
+            page_size (int, optional): The number of results to return.
+                Defaults to None.
+
+        Yields:
+            Dict[str, Any]: A result from the search.
+        """
+        has_more = True
+
+        while has_more:
+            payload = {
+                "query": query,
+                "sort": sort,
+                "filter": filter_criteria,
+                "start_cursor": start_cursor,
+                "page_size": page_size,
+            }
+
+            filtered_payload = self._filter_out_none_values(payload)
+
+            response = self.send_payload("search", "", payload=filtered_payload)
+
+            for result in response.get("results", []):
+                yield result
+
+            next_cursor = response.get("next_cursor")
+            has_more = next_cursor is not None
+            start_cursor = next_cursor
+
+    def get_database(self, database_id: str) -> Any:
+        """Fetches the details of a specific database by its ID.
+
+        Args:
+            database_id (str): The ID of the database to fetch.
+
+        Returns:
+            Any: The details of the database.
+        """
+        return self.fetch_resource("databases", database_id)

omniload/src/notion/helpers/database.py

@@ -0,0 +1,92 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Iterable, Optional
+
+from dlt.common.typing import TDataItem
+
+from .client import NotionClient
+
+
+class NotionDatabase:
+    """
+    A class to represent a Notion database.
+
+    Attributes:
+        database_id (str): The ID of the Notion database.
+        notion_client (NotionClient): A client to interact with the Notion API.
+    """
+
+    def __init__(self, database_id: str, notion_client: NotionClient):
+        self.database_id = database_id
+        self.notion_client = notion_client
+
+    def get_structure(self) -> Any:
+        """Retrieves the structure of the database.
+
+        Notion API Reference. Retrieve a database:
+            https://developers.notion.com/reference/retrieve-a-database
+
+        Returns:
+            Any: The structure of the database.
+        """
+        return self.notion_client.fetch_resource("databases", self.database_id)
+
+    def query(
+        self,
+        filter_properties: Optional[Dict[str, Any]] = None,
+        filter_criteria: Optional[Dict[str, Any]] = None,
+        sorts: Optional[Dict[str, Any]] = None,
+        start_cursor: Optional[str] = None,
+        page_size: Optional[int] = None,
+    ) -> Iterable[TDataItem]:
+        """Queries the database for records.
+
+        Notion API Reference. Query a database:
+            https://developers.notion.com/reference/post-database-query
+
+        Args:
+            filter_properties (Dict[str, Any], optional): A dictionary of
+                properties to filter the records by. Defaults to None.
+            filter_criteria (Dict[str, Any], optional): A dictionary of filters
+                to apply to the records. Defaults to None.
+            sorts (Dict[str, Any], optional): A dictionary of sorts to apply
+                to the records. Defaults to None.
+            start_cursor (str, optional): The cursor to start the query at.
+                Defaults to None.
+            page_size (int, optional): The number of records to return.
+                Defaults to None.
+
+        Yields:
+            List[Dict[str, Any]]: A record from the database.
+        """
+        while True:
+            payload = {
+                "filter": filter_criteria,
+                "sorts": sorts,
+                "start_cursor": start_cursor,
+                "page_size": page_size,
+            }
+            response = self.notion_client.send_payload(
+                "databases",
+                self.database_id,
+                subresource="query",
+                query_params=filter_properties,
+                payload=payload,
+            )
+
+            yield response.get("results", [])
+            if not response.get("has_more"):
+                break
+            start_cursor = response.get("next_cursor")

omniload/src/notion/settings.py

@@ -0,0 +1,17 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Notion source settings and constants"""
+
+API_URL = "https://api.notion.com/v1"

omniload/src/partition.py

@@ -0,0 +1,32 @@
+from typing import Dict
+
+from dlt.common.schema.typing import TColumnSchema
+from dlt.sources import DltResource, DltSource
+
+import omniload.src.resource as resource
+
+
+def apply_athena_hints(
+    source: DltSource | DltResource,
+    partition_column: str,
+    additional_hints: Dict[str, TColumnSchema] = {},
+) -> None:
+    from dlt.destinations.adapters import athena_adapter, athena_partition
+
+    def _apply_partition_hint(resource: DltResource) -> None:
+        columns = resource.columns if resource.columns else {}
+
+        partition_hint = (
+            columns.get(partition_column)  # type: ignore
+            or additional_hints.get(partition_column)
+        )
+
+        athena_adapter(
+            resource,
+            athena_partition.day(partition_column)
+            if partition_hint
+            and partition_hint.get("data_type") in ("timestamp", "date")
+            else partition_column,
+        )
+
+    resource.for_each(source, _apply_partition_hint)