ingestr 0.2.6__py3-none-any.whl → 0.3.1__py3-none-any.whl

ingestr/main.py

@@ -5,7 +5,7 @@ from typing import Optional
 import dlt
 import humanize
 import typer
-from dlt.common.runtime.collector import Collector
+from dlt.common.runtime.collector import Collector, LogCollector
 from rich.console import Console
 from rich.status import Status
 from typing_extensions import Annotated
@@ -32,10 +32,17 @@ DATE_FORMATS = [
     "%Y-%m-%dT%H:%M:%S.%f%z",
 ]
 
+# https://dlthub.com/docs/dlt-ecosystem/file-formats/parquet#supported-destinations
+PARQUET_SUPPORTED_DESTINATIONS = [
+    "bigquery",
+    "duckdb",
+    "snowflake",
+    "databricks",
+    "synapse",
+]
 
-class SpinnerCollector(Collector):
-    """A Collector that shows progress with `tqdm` progress bars"""
 
+class SpinnerCollector(Collector):
     status: Status
     current_step: str
     started: bool
@@ -150,6 +157,27 @@ def ingest(
             envvar="FULL_REFRESH",
         ),
     ] = False,  # type: ignore
+    progress: Annotated[
+        Optional[str],
+        typer.Option(
+            help="The progress display type, must be one of 'interactive', 'log'",
+            envvar="PROGRESS",
+        ),
+    ] = "interactive",  # type: ignore
+    sql_backend: Annotated[
+        Optional[str],
+        typer.Option(
+            help="The SQL backend to use, must be one of 'sqlalchemy', 'pyarrow'",
+            envvar="SQL_BACKEND",
+        ),
+    ] = "pyarrow",  # type: ignore
+    loader_file_format: Annotated[
+        Optional[str],
+        typer.Option(
+            help="The file format to use when loading data, must be one of 'jsonl', 'parquet', 'default'",
+            envvar="LOADER_FILE_FORMAT",
+        ),
+    ] = "default",  # type: ignore
 ):
     track(
         "command_triggered",
@@ -186,12 +214,16 @@ def ingest(
         m = hashlib.sha256()
         m.update(dest_table.encode("utf-8"))
 
+        progressInstance: Collector = SpinnerCollector()
+        if progress == "log":
+            progressInstance = LogCollector()
+
         pipeline = dlt.pipeline(
             pipeline_name=m.hexdigest(),
             destination=destination.dlt_dest(
                 uri=dest_uri,
             ),
-            progress=SpinnerCollector(),
+            progress=progressInstance,
             pipelines_dir="pipeline_data",
             full_refresh=full_refresh,
         )
@@ -231,11 +263,20 @@ def ingest(
             merge_key=merge_key,
             interval_start=interval_start,
             interval_end=interval_end,
+            sql_backend=sql_backend,
         )
 
         if original_incremental_strategy == "delete+insert":
             dlt_source.incremental.primary_key = ()
 
+        if (
+            factory.destination_scheme in PARQUET_SUPPORTED_DESTINATIONS
+            and loader_file_format == "default"
+        ):
+            loader_file_format = "parquet"
+        elif loader_file_format == "default":
+            loader_file_format = "jsonl"
+
         run_info = pipeline.run(
             dlt_source,
             **destination.dlt_run_params(
@@ -244,6 +285,7 @@ def ingest(
             ),
             write_disposition=incremental_strategy,  # type: ignore
             primary_key=(primary_key if primary_key and len(primary_key) > 0 else None),  # type: ignore
+            loader_file_format=loader_file_format,  # type: ignore
         )
 
         destination.post_load()

ingestr/main_test.py

@@ -24,6 +24,8 @@ def invoke_ingest_command(
     merge_key=None,
     interval_start=None,
     interval_end=None,
+    sql_backend=None,
+    loader_file_format=None,
 ):
     args = [
         "ingest",
@@ -61,6 +63,14 @@ def invoke_ingest_command(
         args.append("--interval-end")
         args.append(interval_end)
 
+    if sql_backend:
+        args.append("--sql-backend")
+        args.append(sql_backend)
+
+    if loader_file_format:
+        args.append("--loader-file-format")
+        args.append(loader_file_format)
+
     result = runner.invoke(
         app,
         args,
@@ -137,6 +147,7 @@ def test_append():
             "testschema_append.output",
             "append",
             "updated_at",
+            sql_backend="sqlalchemy",
         )
         assert res.exit_code == 0
 
@@ -193,6 +204,7 @@ def test_merge_with_primary_key():
             "merge",
             "updated_at",
             "id",
+            sql_backend="sqlalchemy",
         )
         assert res.exit_code == 0
         return res
@@ -333,10 +345,10 @@ def test_delete_insert_without_primary_key():
         "CREATE TABLE testschema_delete_insert.input (id INTEGER, val VARCHAR, updated_at TIMESTAMP WITH TIME ZONE)"
     )
     conn.execute(
-        "INSERT INTO testschema_delete_insert.input VALUES (1, 'val1', '2022-01-01')"
+        "INSERT INTO testschema_delete_insert.input VALUES (1, 'val1', '2022-01-01 00:00:00+00:00')"
     )
     conn.execute(
-        "INSERT INTO testschema_delete_insert.input VALUES (2, 'val2', '2022-02-01')"
+        "INSERT INTO testschema_delete_insert.input VALUES (2, 'val2', '2022-02-01 00:00:00+00:00')"
     )
 
     res = conn.sql("select count(*) from testschema_delete_insert.input").fetchall()
@@ -350,6 +362,8 @@ def test_delete_insert_without_primary_key():
             "testschema_delete_insert.output",
             inc_strategy="delete+insert",
             inc_key="updated_at",
+            sql_backend="sqlalchemy",
+            loader_file_format="jsonl",
         )
         assert res.exit_code == 0
         return res
@@ -357,7 +371,7 @@ def test_delete_insert_without_primary_key():
     def get_output_rows():
         conn.execute("CHECKPOINT")
         return conn.sql(
-            "select id, val, strftime(updated_at, '%Y-%m-%d') as updated_at from testschema_delete_insert.output order by id asc"
+            "select id, val, strftime(CAST(updated_at AT TIME ZONE 'UTC' AS TIMESTAMP), '%Y-%m-%d %H:%M:%S') from testschema_delete_insert.output order by id asc"
         ).fetchall()
 
     def assert_output_equals(expected):
@@ -367,7 +381,9 @@ def test_delete_insert_without_primary_key():
             assert res[i] == row
 
     run()
-    assert_output_equals([(1, "val1", "2022-01-01"), (2, "val2", "2022-02-01")])
+    assert_output_equals(
+        [(1, "val1", "2022-01-01 00:00:00"), (2, "val2", "2022-02-01 00:00:00")]
+    )
 
     first_run_id = conn.sql(
         "select _dlt_load_id from testschema_delete_insert.output limit 1"
@@ -375,8 +391,10 @@ def test_delete_insert_without_primary_key():
 
     ##############################
     # we'll run again, since this is a delete+insert, we expect the run ID to change for the last one
-    run()
-    assert_output_equals([(1, "val1", "2022-01-01"), (2, "val2", "2022-02-01")])
+    res = run()
+    assert_output_equals(
+        [(1, "val1", "2022-01-01 00:00:00"), (2, "val2", "2022-02-01 00:00:00")]
+    )
 
     # we ensure that one of the rows is updated with a new run
     count_by_run_id = conn.sql(
@@ -392,17 +410,17 @@ def test_delete_insert_without_primary_key():
     ##############################
     # now we'll insert a few more lines for the same day, the new rows should show up
     conn.execute(
-        "INSERT INTO testschema_delete_insert.input VALUES (3, 'val3', '2022-02-01'), (4, 'val4', '2022-02-01')"
+        "INSERT INTO testschema_delete_insert.input VALUES (3, 'val3', '2022-02-01 00:00:00+00:00'), (4, 'val4', '2022-02-01 00:00:00+00:00')"
     )
     conn.execute("CHECKPOINT")
 
     run()
     assert_output_equals(
         [
-            (1, "val1", "2022-01-01"),
-            (2, "val2", "2022-02-01"),
-            (3, "val3", "2022-02-01"),
-            (4, "val4", "2022-02-01"),
+            (1, "val1", "2022-01-01 00:00:00"),
+            (2, "val2", "2022-02-01 00:00:00"),
+            (3, "val3", "2022-02-01 00:00:00"),
+            (4, "val4", "2022-02-01 00:00:00"),
         ]
     )
 
@@ -460,6 +478,8 @@ def test_delete_insert_with_timerange():
             inc_key="updated_at",
             interval_start=start_date,
             interval_end=end_date,
+            sql_backend="sqlalchemy",
+            loader_file_format="jsonl",
         )
         assert res.exit_code == 0
         return res

ingestr/src/destinations.py

@@ -177,7 +177,8 @@ class CsvDestination(GenericSqlDestination):
         )
 
         output_path = self.uri.split("://")[1]
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        if output_path.count("/") > 1:
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
 
         with gzip.open(first_file_path, "rt", encoding="utf-8") as jsonl_file:  # type: ignore
             with open(output_path, "w", newline="") as csv_file:

ingestr/src/factory.py

@@ -14,7 +14,7 @@ from ingestr.src.destinations import (
     SnowflakeDestination,
     SynapseDestination,
 )
-from ingestr.src.sources import LocalCsvSource, MongoDbSource, SqlSource
+from ingestr.src.sources import LocalCsvSource, MongoDbSource, NotionSource, SqlSource
 
 SQL_SOURCE_SCHEMES = [
     "bigquery",
@@ -80,6 +80,8 @@ class SourceDestinationFactory:
             return LocalCsvSource()
         elif self.source_scheme == "mongodb":
             return MongoDbSource()
+        elif self.source_scheme == "notion":
+            return NotionSource()
         else:
             raise ValueError(f"Unsupported source scheme: {self.source_scheme}")
 

ingestr/src/mongodb/__init__.py

@@ -70,7 +70,7 @@ def mongodb_collection(
     collection: str = dlt.config.value,
     incremental: Optional[dlt.sources.incremental] = None,  # type: ignore[type-arg]
     write_disposition: Optional[str] = dlt.config.value,
-    parallel: Optional[bool] = dlt.config.value,
+    parallel: Optional[bool] = False,
 ) -> Any:
     """
     A DLT source which loads a collection from a mongo database using PyMongo.

ingestr/src/mongodb/helpers.py

@@ -83,7 +83,7 @@ class CollectionLoaderParallell(CollectionLoader):
     def _get_cursor(self) -> TCursor:
         cursor = self.collection.find(filter=self._filter_op)
         if self._sort_op:
-            cursor = cursor.sort(self._sort_op)  # type: ignore
+            cursor = cursor.sort(self._sort_op)
         return cursor
 
     @dlt.defer
@@ -155,11 +155,11 @@ class MongoDbCollectionConfiguration(BaseConfiguration):
 
 @configspec
 class MongoDbCollectionResourceConfiguration(BaseConfiguration):
-    connection_url: str
-    database: Optional[str]
-    collection: str
+    connection_url: str = dlt.secrets.value
+    database: Optional[str] = dlt.config.value
+    collection: str = dlt.config.value
     incremental: Optional[dlt.sources.incremental] = None  # type: ignore[type-arg]
-    write_disposition: Optional[str] = None
+    write_disposition: Optional[str] = dlt.config.value
     parallel: Optional[bool] = False
 
 

ingestr/src/notion/__init__.py

@@ -0,0 +1,55 @@
+"""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
+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",
+        )

ingestr/src/notion/helpers/client.py

@@ -0,0 +1,164 @@
+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)

ingestr/src/notion/helpers/database.py

@@ -0,0 +1,78 @@
+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")

ingestr/src/notion/settings.py

@@ -0,0 +1,3 @@
+"""Notion source settings and constants"""
+
+API_URL = "https://api.notion.com/v1"

ingestr/src/sources.py

@@ -1,9 +1,11 @@
 import csv
 from typing import Callable
+from urllib.parse import parse_qs, urlparse
 
 import dlt
 
 from ingestr.src.mongodb import mongodb_collection
+from ingestr.src.notion import notion_databases
 from ingestr.src.sql_database import sql_table
 
 
@@ -39,6 +41,7 @@ class SqlSource:
             table=table_fields[-1],
             incremental=incremental,
             merge_key=kwargs.get("merge_key"),
+            backend=kwargs.get("sql_backend", "sqlalchemy"),
         )
 
         return table_instance
@@ -104,3 +107,25 @@ class LocalCsvSource:
             csv_file,
             merge_key=kwargs.get("merge_key"),  # type: ignore
         )
+
+
+class NotionSource:
+    table_builder: Callable
+
+    def __init__(self, table_builder=notion_databases) -> None:
+        self.table_builder = table_builder
+
+    def dlt_source(self, uri: str, table: str, **kwargs):
+        if kwargs.get("incremental_key"):
+            raise ValueError("Incremental loads are not supported for Notion")
+
+        source_fields = urlparse(uri)
+        source_params = parse_qs(source_fields.query)
+        api_key = source_params.get("api_key")
+        if not api_key:
+            raise ValueError("api_key in the URI is required to connect to Notion")
+
+        return self.table_builder(
+            database_ids=[{"id": table}],
+            api_key=api_key[0],
+        )

ingestr/src/sources_test.py

@@ -22,10 +22,11 @@ class SqlSourceTest(unittest.TestCase):
         table = "schema.table"
 
         # monkey patch the sql_table function
-        def sql_table(credentials, schema, table, incremental, merge_key):
+        def sql_table(credentials, schema, table, incremental, merge_key, backend):
             self.assertEqual(credentials, uri)
             self.assertEqual(schema, "schema")
             self.assertEqual(table, "table")
+            self.assertEqual(backend, "sqlalchemy")
             self.assertIsNone(incremental)
             self.assertIsNone(merge_key)
             return dlt.resource()
@@ -40,10 +41,11 @@ class SqlSourceTest(unittest.TestCase):
         incremental_key = "id"
 
         # monkey patch the sql_table function
-        def sql_table(credentials, schema, table, incremental, merge_key):
+        def sql_table(credentials, schema, table, incremental, merge_key, backend):
             self.assertEqual(credentials, uri)
             self.assertEqual(schema, "schema")
             self.assertEqual(table, "table")
+            self.assertEqual(backend, "sqlalchemy")
             self.assertIsInstance(incremental, dlt.sources.incremental)
             self.assertEqual(incremental.cursor_path, incremental_key)
             self.assertIsNone(merge_key)