omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/sql_database/callbacks.py

@@ -0,0 +1,66 @@
+from datetime import datetime
+
+from sqlalchemy import text
+from sqlalchemy import types as sa
+from sqlalchemy.dialects import mysql
+
+
+def type_adapter_callback(sql_type):
+    if isinstance(sql_type, mysql.SET):
+        return sa.JSON
+    return sql_type
+
+
+def chained_query_adapter_callback(query_adapters):
+    """
+    This function is used to chain multiple query adapters together,.
+    This gives us the flexibility to introduce various adapters based on the given command parameters.
+    """
+
+    def callback(query, table):
+        for adapter in query_adapters:
+            query = adapter(query, table)
+
+        return query
+
+    return callback
+
+
+def limit_callback(sql_limit: int, incremental_key: str):
+    def callback(query, table):
+        query = query.limit(sql_limit)
+        if incremental_key:
+            query = query.order_by(incremental_key)
+        return query
+
+    return callback
+
+
+def custom_query_variable_subsitution(query_value: str, kwargs: dict):
+    def callback(query, table, incremental=None, engine=None):
+        params = {}
+        if incremental:
+            params["interval_start"] = (
+                incremental.last_value
+                if incremental.last_value is not None
+                else datetime(year=1, month=1, day=1)
+            )
+            if incremental.end_value is not None:
+                params["interval_end"] = incremental.end_value
+        else:
+            if ":interval_start" in query_value:
+                params["interval_start"] = (
+                    datetime.min
+                    if kwargs.get("interval_start") is None
+                    else kwargs.get("interval_start")
+                )
+            if ":interval_end" in query_value:
+                params["interval_end"] = (
+                    datetime.max
+                    if kwargs.get("interval_end") is None
+                    else kwargs.get("interval_end")
+                )
+
+        return text(query_value).bindparams(**params)
+
+    return callback

omniload/src/stripe_analytics/__init__.py

@@ -0,0 +1,183 @@
+# 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.
+
+"""This source uses Stripe API and dlt to load data such as Customer, Subscription, Event etc. to the database and to calculate the MRR and churn rate."""
+
+from typing import Any, Dict, Generator, Iterable, Optional, Tuple
+
+import dlt
+import stripe
+from dlt.sources import DltResource
+from pendulum import DateTime
+
+from .helpers import (
+    async_parallel_pagination,
+    generate_date_ranges,
+    pagination,
+    transform_date,
+)
+
+
+@dlt.source(max_table_nesting=0)
+def stripe_source(
+    endpoints: Tuple[str, ...],
+    stripe_secret_key: str = dlt.secrets.value,
+    start_date: Optional[DateTime] = None,
+    end_date: Optional[DateTime] = None,
+) -> Iterable[DltResource]:
+    """
+    Retrieves data from the Stripe API for the specified endpoints.
+
+    For all endpoints, Stripe API responses do not provide the key "updated",
+    so in most cases, we are forced to load the data in 'replace' mode.
+    This source is suitable for all types of endpoints, including 'Events', 'Invoice', etc.
+    but these endpoints can also be loaded in incremental mode (see source incremental_stripe_source).
+
+    Args:
+        endpoints (Tuple[str, ...]): A tuple of endpoint names to retrieve data from. Defaults to most popular Stripe API endpoints.
+        stripe_secret_key (str): The API access token for authentication. Defaults to the value in the `dlt.secrets` object.
+        start_date (Optional[DateTime]): An optional start date to limit the data retrieved. Format: datetime(YYYY, MM, DD). Defaults to None.
+        end_date (Optional[DateTime]): An optional end date to limit the data retrieved. Format: datetime(YYYY, MM, DD). Defaults to None.
+
+    Returns:
+        Iterable[DltResource]: Resources with data that was created during the period greater than or equal to 'start_date' and less than 'end_date'.
+    """
+    stripe.api_key = stripe_secret_key
+    stripe.api_version = "2022-11-15"
+
+    def stripe_resource(
+        endpoint: str,
+    ) -> Generator[Dict[Any, Any], Any, None]:
+        yield from pagination(endpoint, start_date, end_date)
+
+    for endpoint in endpoints:
+        yield dlt.resource(
+            stripe_resource,
+            name=endpoint,
+            write_disposition="replace",
+        )(endpoint)
+
+
+@dlt.source(max_table_nesting=0)
+def async_stripe_source(
+    endpoints: Tuple[str, ...],
+    stripe_secret_key: str = dlt.secrets.value,
+    start_date: Optional[DateTime] = None,
+    end_date: Optional[DateTime] = None,
+    max_workers: int = 4,
+    rate_limit_delay: float = 0.03,
+) -> Iterable[DltResource]:
+    """
+    ULTRA-FAST async Stripe source optimized for maximum speed and throughput.
+
+    WARNING: Returns data in RANDOM ORDER for maximum performance.
+    Uses aggressive concurrency and minimal delays to maximize API throughput.
+
+    Args:
+        endpoints (Tuple[str, ...]): A tuple of endpoint names to retrieve data from.
+        stripe_secret_key (str): The API access token for authentication. Defaults to the value in the `dlt.secrets` object.
+        start_date (Optional[DateTime]): An optional start date to limit the data retrieved. Format: datetime(YYYY, MM, DD). Defaults to 2010-01-01.
+        end_date (Optional[DateTime]): An optional end date to limit the data retrieved. Format: datetime(YYYY, MM, DD). Defaults to today.
+        max_workers (int): Maximum number of concurrent async tasks. Defaults to 40 for maximum speed.
+        rate_limit_delay (float): Minimal delay between requests. Defaults to 0.03 seconds.
+
+    Returns:
+        Iterable[DltResource]: Resources with data in RANDOM ORDER (optimized for speed).
+    """
+    stripe.api_key = stripe_secret_key
+    stripe.api_version = "2022-11-15"
+
+    async def async_stripe_resource(endpoint: str):
+        yield async_parallel_pagination(endpoint, max_workers, rate_limit_delay)
+
+    for endpoint in endpoints:
+        yield dlt.resource(
+            async_stripe_resource,
+            name=endpoint,
+            write_disposition="replace",
+        )(endpoint)
+
+
+@dlt.source(max_table_nesting=0)
+def incremental_stripe_source(
+    endpoints: Tuple[str, ...],
+    stripe_secret_key: str = dlt.secrets.value,
+    initial_start_date: Optional[DateTime] = None,
+    end_date: Optional[DateTime] = None,
+) -> Iterable[DltResource]:
+    stripe.api_key = stripe_secret_key
+    stripe.api_version = "2022-11-15"
+    start_date_unix = (
+        transform_date(initial_start_date) if initial_start_date is not None else -1
+    )
+
+    for endpoint in endpoints:
+
+        def date_range_resource(
+            endpoint: str = endpoint,
+            created: Optional[Any] = dlt.sources.incremental(
+                "created",
+                initial_value=start_date_unix,
+                end_value=transform_date(end_date) if end_date is not None else None,
+                range_end="closed",
+                range_start="closed",
+            ),
+        ) -> Generator[Dict[str, Any], None, None]:
+            from dlt.common import pendulum
+
+            # Use 2010-01-01 as default start (Stripe founding year) to avoid
+            # generating hundreds of thousands of hourly ranges from 1969
+            default_start_ts = int(pendulum.datetime(2010, 1, 1).timestamp())
+            start_ts = (
+                created.last_value
+                if created.last_value is not None
+                else start_date_unix
+            )
+            if start_ts < 0:
+                start_ts = default_start_ts
+            end_ts = (
+                created.end_value
+                if created.end_value is not None
+                else int(pendulum.now().timestamp())
+            )
+            for date_range in generate_date_ranges(start_ts, end_ts):
+                date_range["endpoint"] = endpoint
+                date_range["created"] = date_range["end_ts"]
+                yield date_range
+
+        def fetch_date_range(
+            date_range: Dict[str, int],
+        ) -> Generator[Dict[Any, Any], Any, None]:
+            """Transformer that fetches data for a given date range."""
+            yield from pagination(
+                date_range["endpoint"],
+                start_date=date_range["start_ts"],
+                end_date=date_range["end_ts"],
+            )
+
+        date_ranges = dlt.resource(
+            date_range_resource,
+            name=f"{endpoint}_date_ranges",
+        )()
+
+        yield (
+            date_ranges
+            | dlt.transformer(
+                fetch_date_range,
+                name=endpoint,
+                write_disposition="merge",
+                primary_key="id",
+                parallelized=True,
+            )
+        )

omniload/src/stripe_analytics/helpers.py

@@ -0,0 +1,386 @@
+# 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.
+
+"""Stripe analytics source helpers"""
+
+import asyncio
+import math
+from datetime import datetime, timedelta
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+import stripe
+from dlt.common import pendulum
+from dlt.common.typing import TDataItem
+from pendulum import DateTime
+
+
+def pagination(
+    endpoint: str, start_date: Optional[Any] = None, end_date: Optional[Any] = None
+) -> Iterable[TDataItem]:
+    """
+    Retrieves data from an endpoint with pagination.
+
+    Args:
+        endpoint (str): The endpoint to retrieve data from.
+        start_date (Optional[Any]): An optional start date to limit the data retrieved. Defaults to None.
+        end_date (Optional[Any]): An optional end date to limit the data retrieved. Defaults to None.
+
+    Returns:
+        Iterable[TDataItem]: Data items retrieved from the endpoint.
+    """
+    starting_after = None
+    while True:
+        response = stripe_get_data(
+            endpoint,
+            start_date=start_date,
+            end_date=end_date,
+            starting_after=starting_after,
+        )
+
+        if len(response["data"]) > 0:
+            starting_after = response["data"][-1]["id"]
+        yield response["data"]
+
+        if not response["has_more"]:
+            break
+
+
+def _create_time_chunks(start_ts: int, end_ts: int, num_chunks: int) -> List[tuple]:
+    """
+    Divide a time range into equal chunks for parallel processing.
+
+    Args:
+        start_ts (int): Start timestamp
+        end_ts (int): End timestamp
+        num_chunks (int): Number of chunks to create
+
+    Returns:
+        List[tuple]: List of (chunk_start, chunk_end) timestamp pairs
+    """
+    total_duration = end_ts - start_ts
+    chunk_duration = math.ceil(total_duration / num_chunks)
+
+    chunks = []
+    current_start = start_ts
+
+    for i in range(num_chunks):
+        current_end = min(current_start + chunk_duration, end_ts)
+        if current_start < end_ts:
+            chunks.append((current_start, current_end))
+        current_start = current_end
+
+        if current_start >= end_ts:
+            break
+
+    return chunks
+
+
+def _create_adaptive_time_chunks(
+    start_ts: int, end_ts: int, max_workers: int
+) -> List[tuple]:
+    """
+    Create time chunks with adaptive sizing - larger chunks for 2010s (less data expected).
+
+    Args:
+        start_ts (int): Start timestamp
+        end_ts (int): End timestamp
+        max_workers (int): Maximum number of workers
+
+    Returns:
+        List[tuple]: List of (chunk_start, chunk_end) timestamp pairs
+    """
+    chunks = []
+
+    # Key timestamps
+    year_2020_ts = int(pendulum.datetime(2020, 1, 1).timestamp())
+    year_2015_ts = int(pendulum.datetime(2015, 1, 1).timestamp())
+
+    current_start = start_ts
+
+    # Handle 2010-2015: Large chunks (2-3 year periods)
+    if current_start < year_2015_ts:
+        chunk_end = min(year_2015_ts, end_ts)
+        if current_start < chunk_end:
+            # Split 2010-2015 into 2-3 chunks max
+            pre_2015_chunks = _create_time_chunks(
+                current_start, chunk_end, min(3, max_workers)
+            )
+            chunks.extend(pre_2015_chunks)
+        current_start = chunk_end
+
+    # Handle 2015-2020: Medium chunks (6 month to 1 year periods)
+    if current_start < year_2020_ts and current_start < end_ts:
+        chunk_end = min(year_2020_ts, end_ts)
+        if current_start < chunk_end:
+            # Split 2015-2020 into smaller chunks
+            duration_2015_2020 = chunk_end - current_start
+            years_2015_2020 = duration_2015_2020 / (365 * 24 * 60 * 60)
+            num_chunks_2015_2020 = min(
+                max_workers, max(2, int(years_2015_2020 * 2))
+            )  # ~6 months per chunk
+
+            pre_2020_chunks = _create_time_chunks(
+                current_start, chunk_end, num_chunks_2015_2020
+            )
+            chunks.extend(pre_2020_chunks)
+        current_start = chunk_end
+
+    if current_start < end_ts:
+        # Split post-2020 data into daily chunks for maximum granularity
+        current_chunk_start = current_start
+        while current_chunk_start < end_ts:
+            # Calculate end of current day
+            current_date = datetime.fromtimestamp(current_chunk_start)
+            next_day = current_date + timedelta(days=1)
+            chunk_end = min(int(next_day.timestamp()), end_ts)
+
+            chunks.append((current_chunk_start, chunk_end))
+            current_chunk_start = chunk_end
+
+    return chunks
+
+
+def _fetch_chunk_data_streaming(
+    endpoint: str, start_ts: int, end_ts: int
+) -> List[List[TDataItem]]:
+    """
+    Fetch data for a specific time chunk using sequential pagination with memory-efficient approach.
+
+    Args:
+        endpoint (str): The Stripe endpoint to fetch from
+        start_ts (int): Start timestamp for this chunk
+        end_ts (int): End timestamp for this chunk
+
+    Returns:
+        List[List[TDataItem]]: List of batches of data items
+    """
+    # For streaming, we still need to collect the chunk data to maintain order
+    # but we can optimize by not holding all data in memory at once
+    print(
+        f"Fetching chunk {datetime.fromtimestamp(start_ts).strftime('%Y-%m-%d')}-{datetime.fromtimestamp(end_ts).strftime('%Y-%m-%d')}"
+    )
+    chunk_data = []
+    batch_count = 0
+
+    for batch in pagination(endpoint, start_ts, end_ts):
+        chunk_data.append(batch)
+        print(
+            f"Processed {batch_count} batches for chunk {datetime.fromtimestamp(start_ts).strftime('%Y-%m-%d')}-{datetime.fromtimestamp(end_ts).strftime('%Y-%m-%d')}"
+        )
+        batch_count += 1
+
+    return chunk_data
+
+
+async def async_pagination(
+    endpoint: str, start_date: Optional[Any] = None, end_date: Optional[Any] = None
+) -> Iterable[TDataItem]:
+    """
+    Async version of pagination that retrieves data from an endpoint with pagination.
+
+    Args:
+        endpoint (str): The endpoint to retrieve data from.
+        start_date (Optional[Any]): An optional start date to limit the data retrieved. Defaults to None.
+        end_date (Optional[Any]): An optional end date to limit the data retrieved. Defaults to None.
+
+    Returns:
+        Iterable[TDataItem]: Data items retrieved from the endpoint.
+    """
+    starting_after = None
+    while True:
+        response = await stripe_get_data_async(
+            endpoint,
+            start_date=start_date,
+            end_date=end_date,
+            starting_after=starting_after,
+        )
+
+        if len(response["data"]) > 0:
+            starting_after = response["data"][-1]["id"]
+        yield response["data"]
+
+        if not response["has_more"]:
+            break
+
+
+async def async_parallel_pagination(
+    endpoint: str,
+    max_workers: int = 8,
+    rate_limit_delay: float = 5,
+) -> Iterable[TDataItem]:
+    """
+    ULTRA-FAST async parallel pagination - yields data in random order for maximum speed.
+    No ordering constraints - pure performance optimization.
+
+    Args:
+        endpoint (str): The endpoint to retrieve data from.
+        start_date (Optional[Any]): An optional start date to limit the data retrieved. Defaults to 2010-01-01 if None.
+        end_date (Optional[Any]): An optional end date to limit the data retrieved. Defaults to today if None.
+        max_workers (int): Maximum number of concurrent async tasks. Defaults to 8 for balanced speed/rate limit respect.
+        rate_limit_delay (float): Minimal delay between requests. Defaults to 5 seconds.
+
+    Returns:
+        Iterable[TDataItem]: Data items retrieved from the endpoint (RANDOM ORDER FOR SPEED).
+    """
+
+    start_date = pendulum.datetime(2010, 1, 1)
+    end_date = pendulum.now()
+    start_ts = transform_date(start_date)
+    end_ts = transform_date(end_date)
+
+    # Create time chunks with larger chunks for 2010s (less data expected)
+    time_chunks = _create_adaptive_time_chunks(start_ts, end_ts, max_workers)
+
+    # Use asyncio semaphore to control concurrency and respect rate limits
+    semaphore = asyncio.Semaphore(max_workers)
+
+    async def fetch_chunk_with_semaphore(chunk_start: int, chunk_end: int):
+        async with semaphore:
+            return await _fetch_chunk_data_async_fast(endpoint, chunk_start, chunk_end)
+
+    # Create all tasks
+    tasks = [
+        fetch_chunk_with_semaphore(chunk_start, chunk_end)
+        for chunk_start, chunk_end in time_chunks
+    ]
+
+    for coro in asyncio.as_completed(tasks):
+        try:
+            chunk_data = await coro
+
+            for batch in chunk_data:
+                yield batch
+
+        except Exception as exc:
+            print(f"Async chunk processing generated an exception: {exc}")
+            raise exc
+
+
+async def _fetch_chunk_data_async_fast(
+    endpoint: str, start_ts: int, end_ts: int
+) -> List[List[TDataItem]]:
+    """
+    ULTRA-FAST async chunk fetcher - no metadata overhead, direct data return.
+
+    Args:
+        endpoint (str): The Stripe endpoint to fetch from
+        start_ts (int): Start timestamp for this chunk
+        end_ts (int): End timestamp for this chunk
+
+    Returns:
+        List[List[TDataItem]]: Raw batches with zero overhead
+    """
+    chunk_data = []
+    async for batch in async_pagination(endpoint, start_ts, end_ts):
+        chunk_data.append(batch)
+
+    return chunk_data
+
+
+def generate_date_ranges(start_ts: int, end_ts: int) -> Iterable[Dict[str, int]]:
+    """Generate hourly date range dicts for parallel processing.
+
+    Args:
+        start_ts (int): Start timestamp (unix)
+        end_ts (int): End timestamp (unix)
+
+    Yields:
+        Dict[str, int]: Dictionary with 'start_ts' and 'end_ts' keys for each hour
+    """
+    current_ts = start_ts
+
+    while current_ts < end_ts:
+        next_hour = (current_ts // 3600 + 1) * 3600
+        next_ts = min(next_hour, end_ts)
+        yield {"start_ts": current_ts, "end_ts": next_ts}
+        current_ts = next_ts
+
+
+def transform_date(date: Union[str, DateTime, int]) -> int:
+    if isinstance(date, str):
+        date = pendulum.from_format(date, "%Y-%m-%dT%H:%M:%SZ")
+    if isinstance(date, DateTime):
+        # convert to unix timestamp
+        date = int(date.timestamp())
+    return date
+
+
+def stripe_get_data(
+    resource: str,
+    start_date: Optional[Any] = None,
+    end_date: Optional[Any] = None,
+    **kwargs: Any,
+) -> Dict[Any, Any]:
+    if start_date:
+        start_date = transform_date(start_date)
+    if end_date:
+        end_date = transform_date(end_date)
+
+    if resource == "Subscription":
+        kwargs.update({"status": "all"})
+
+    resource_dict = getattr(stripe, resource).list(
+        created={"gte": start_date, "lt": end_date}, limit=100, **kwargs
+    )
+    return dict(resource_dict)
+
+
+async def stripe_get_data_async(
+    resource: str,
+    start_date: Optional[Any] = None,
+    end_date: Optional[Any] = None,
+    **kwargs: Any,
+) -> Dict[Any, Any]:
+    """Async version of stripe_get_data"""
+    if start_date:
+        start_date = transform_date(start_date)
+    if end_date:
+        end_date = transform_date(end_date)
+
+    if resource == "Subscription":
+        kwargs.update({"status": "all"})
+
+    import asyncio
+
+    from stripe import RateLimitError
+
+    max_retries = 50
+    retry_count = 0
+    max_wait_time_ms = 10000
+
+    while retry_count < max_retries:
+        # print(
+        #     f"Fetching {resource} from {datetime.fromtimestamp(start_date).strftime('%Y-%m-%d %H:%M:%S') if start_date else 'None'} to {datetime.fromtimestamp(end_date).strftime('%Y-%m-%d %H:%M:%S') if end_date else 'None'}, retry {retry_count} of {max_retries}",
+        #     flush=True,
+        # )
+        try:
+            resource_dict = await getattr(stripe, resource).list_async(
+                created={"gte": start_date, "lt": end_date}, limit=100, **kwargs
+            )
+            return dict(resource_dict)
+        except RateLimitError:
+            retry_count += 1
+            if retry_count < max_retries:
+                wait_time = min(2**retry_count * 0.001, max_wait_time_ms)
+                print(
+                    f"Got rate limited, sleeping {wait_time} seconds before retrying...",
+                    flush=True,
+                )
+                await asyncio.sleep(wait_time)
+            else:
+                # Re-raise the last exception if we've exhausted retries
+                print(f"✗ Failed to fetch {resource} after {max_retries} retries")
+                raise
+
+    return dict(resource_dict)