diff --git a/omniload/source/airtable/adapter.py b/omniload/source/airtable/adapter.py index 160df4d3c..3e794837c 100644 --- a/omniload/source/airtable/adapter.py +++ b/omniload/source/airtable/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,7 +23,7 @@ from dlt.sources import DltResource -@dlt.source(max_table_nesting=1) +@dlt.source def airtable_source( base_id: str = dlt.config.value, table_names: Optional[List[str]] = dlt.config.value, @@ -64,13 +64,12 @@ def airtable_resource( It starts with "app". See https://support.airtable.com/docs/finding-airtable-ids table (Dict[str, Any]): Metadata about an airtable, does not contain the actual records """ - primary_key_id = table["primaryFieldId"] primary_key_field = [ field for field in table["fields"] if field["id"] == primary_key_id ][0] table_name: str = table["name"] - primary_key: List[str] = [f"fields__{primary_key_field['name']}".lower()] + primary_key: List[str] = [primary_key_field["name"]] air_table = api.table(base_id, table["id"]) # Table.iterate() supports rich customization options, such as chunk size, fields, cell format, timezone, locale, and view diff --git a/omniload/source/asana/adapter.py b/omniload/source/asana/adapter.py index 1cd98b001..f2254bc7a 100644 --- a/omniload/source/asana/adapter.py +++ b/omniload/source/asana/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -72,7 +72,7 @@ def workspaces( Yields: dict: The workspace data. """ - yield from get_client(access_token).workspaces.find_all(opt_fields=",".join(fields)) # ty: ignore[unresolved-attribute] + yield from get_client(access_token).workspaces.find_all(opt_fields=",".join(fields)) @dlt.transformer( @@ -95,7 +95,7 @@ def projects( list[dict]: The project data for the given workspace. """ return list( - get_client(access_token).projects.find_all( # ty: ignore[unresolved-attribute] + get_client(access_token).projects.find_all( workspace=workspace["gid"], timeout=REQUEST_TIMEOUT, opt_fields=",".join(fields), @@ -125,7 +125,7 @@ def sections( return [ section for project in project_array - for section in get_client(access_token).sections.get_sections_for_project( # ty: ignore[unresolved-attribute] + for section in get_client(access_token).sections.get_sections_for_project( project_gid=project["gid"], timeout=REQUEST_TIMEOUT, opt_fields=",".join(fields), @@ -151,7 +151,7 @@ def tags( """ return [ tag - for tag in get_client(access_token).tags.find_all( # ty: ignore[unresolved-attribute] + for tag in get_client(access_token).tags.find_all( workspace=workspace["gid"], timeout=REQUEST_TIMEOUT, opt_fields=",".join(fields), @@ -164,10 +164,7 @@ def tasks( project_array: t.List[TDataItem], access_token: str = dlt.secrets.value, modified_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "modified_at", - initial_value=DEFAULT_START_DATE, - range_end="closed", - range_start="closed", + "modified_at", initial_value=DEFAULT_START_DATE ), fields: Iterable[str] = TASK_FIELDS, ) -> Iterable[TDataItem]: @@ -185,7 +182,7 @@ def tasks( yield from ( task for project in project_array - for task in get_client(access_token).tasks.find_all( # ty: ignore[unresolved-attribute] + for task in get_client(access_token).tasks.find_all( project=project["gid"], timeout=REQUEST_TIMEOUT, modified_since=modified_at.start_value, @@ -196,7 +193,7 @@ def tasks( @dlt.transformer( data_from=tasks, - write_disposition="replace", + write_disposition="append", ) @dlt.defer def stories( @@ -215,7 +212,7 @@ def stories( """ return [ story - for story in get_client(access_token).stories.get_stories_for_task( # ty: ignore[unresolved-attribute] + for story in get_client(access_token).stories.get_stories_for_task( task_gid=task["gid"], timeout=REQUEST_TIMEOUT, opt_fields=",".join(fields), @@ -244,7 +241,7 @@ def teams( """ return [ team - for team in get_client(access_token).teams.find_by_organization( # ty: ignore[unresolved-attribute] + for team in get_client(access_token).teams.find_by_organization( organization=workspace["gid"], timeout=REQUEST_TIMEOUT, opt_fields=",".join(fields), @@ -273,7 +270,7 @@ def users( """ return [ user - for user in get_client(access_token).users.find_all( # ty: ignore[unresolved-attribute] + for user in get_client(access_token).users.find_all( workspace=workspace["gid"], timeout=REQUEST_TIMEOUT, opt_fields=",".join(fields), diff --git a/omniload/source/asana/helpers.py b/omniload/source/asana/helpers.py index eacb0e1b8..9f71cfe2e 100644 --- a/omniload/source/asana/helpers.py +++ b/omniload/source/asana/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,8 +14,7 @@ """Asana source helpers""" -from asana import ApiClient as AsanaClient -from asana import Configuration +from asana import Client as AsanaClient def get_client( @@ -28,6 +27,5 @@ def get_client( Returns: AsanaClient: The Asana API client. """ - configuration = Configuration() - configuration.access_token = access_token - return AsanaClient(configuration=configuration) + asana = AsanaClient.access_token(access_token) + return asana diff --git a/omniload/source/asana/settings.py b/omniload/source/asana/settings.py index dc0c35c4e..6fbcf84b8 100644 --- a/omniload/source/asana/settings.py +++ b/omniload/source/asana/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/bing_webmaster/adapter.py b/omniload/source/bing_webmaster/adapter.py new file mode 100644 index 000000000..17b2249af --- /dev/null +++ b/omniload/source/bing_webmaster/adapter.py @@ -0,0 +1,110 @@ +# Copyright 2022-2026 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 loading history of organic search traffic from Bing Webmaster API +See documentation: https://learn.microsoft.com/en-us/dotnet/api/microsoft.bing.webmaster.api.interfaces.iwebmasterapi?view=bing-webmaster-dotnet +The API returns aggregated weekly statistics for the entire history of up to 26 weeks. +The dates are always Fridays and during tests, the data up to the latest Friday has been available on the following Monday. +""" + +import time +from typing import Iterable, Iterator, List, Sequence + +import dlt +from dlt.common import logger +from dlt.common.typing import DictStrAny, DictStrStr +from dlt.sources import DltResource + +from .helpers import get_stats_with_retry, parse_response + + +@dlt.source(name="bing_webmaster") +def source( + site_urls: List[str] = None, site_url_pages: Iterable[DictStrStr] = None +) -> Sequence[DltResource]: + """ + A dlt source for the Bing Webmaster api. + It groups resources for the APIs which return organic search traffic statistics + Args: + site_urls: List[str]: A list of site_urls, e.g, ["dlthub.com", "dlthub.de"]. Use this if you need the weekly traffic per site_url and page + site_url_pages: Iterable[DictStrStr]: A list of pairs of site_url and page. Use this if you need the weekly traffic per site_url, page, and query + Returns: + Sequence[DltResource]: A sequence of resources that can be selected from including page_stats and page_query_stats. + """ + return ( + page_stats(site_urls), + page_query_stats(site_url_pages), + ) + + +@dlt.resource( + write_disposition="merge", + merge_key=("date", "page", "site_url"), + primary_key=("date", "page", "site_url"), + table_name="bing_page_stats", +) +def page_stats( + site_urls: List[str], api_key: str = dlt.secrets.value +) -> Iterator[Iterator[DictStrAny]]: + """ + Yields detailed traffic statistics for top pages belonging to a site_url + Contains the entire available history of up to 26 weeks. Thus, we recommend to use write_disposition="merge" + API documentation: + https://learn.microsoft.com/en-us/dotnet/api/microsoft.bing.webmaster.api.interfaces.iwebmasterapi.getpagestats + Args: + site_urls (List[str]): List of site_urls to retrieve statistics for. + Yields: + Iterator[Dict[str, Any]]: An iterator over list of organic traffic statistics. + """ + api_path = "GetPageStats" + for site_url in site_urls: + params = {"siteUrl": site_url, "apikey": api_key} + logger.info(f"Fetching for site_url: {site_url}") + response = get_stats_with_retry(api_path, params) + if len(response) > 0: + yield parse_response(response, site_url) + + +@dlt.resource( + write_disposition="merge", + merge_key=("date", "page", "site_url", "query"), + primary_key=("date", "page", "site_url", "query"), + table_name="bing_page_query_stats", +) +def page_query_stats( + site_url_pages: Iterable[DictStrStr], + api_key: str = dlt.secrets.value, +) -> Iterator[Iterator[DictStrAny]]: + """ + Yields weekly statistics and queries for each pair of page and site_url. + Contains the entire available history of up to 26 weeks. Thus, we recommend to use write_disposition="merge" + API documentation: + https://learn.microsoft.com/en-us/dotnet/api/microsoft.bing.webmaster.api.interfaces.iwebmasterapi.getpagequerystats + + Args: + site_url_page (Iterable[DictStrStr]): Iterable of site_url and pages to retrieve statistics for. Can be result of a SQL query, a parsed sitemap, etc. + Yields: + Iterator[Dict[str, Any]]: An iterator over list of organic traffic statistics. + """ + api_path = "GetPageQueryStats" + for record in site_url_pages: + time.sleep(0.5) # this avoids rate limit observed after dozens of requests + site_url = record.get("site_url") + page = record.get("page") + params = {"siteUrl": site_url, "page": page, "apikey": api_key} + logger.info(f"Fetching for site_url: {site_url}, page: {page}") + response = get_stats_with_retry(api_path, params) + if len(response) > 0: + yield parse_response(response, site_url, page) diff --git a/omniload/source/bing_webmaster/helpers.py b/omniload/source/bing_webmaster/helpers.py new file mode 100644 index 000000000..3e68aa56b --- /dev/null +++ b/omniload/source/bing_webmaster/helpers.py @@ -0,0 +1,74 @@ +# Copyright 2022-2026 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. + +"""Bing Webmaster source helpers""" + +import re +from typing import Iterator, List +from urllib.parse import urljoin + +from dlt.common import logger, pendulum +from dlt.common.typing import DictStrAny, DictStrStr +from dlt.sources.helpers import requests + +from .settings import BASE_URL, HEADERS + + +def get_url_with_retry(url: str, params: DictStrStr) -> DictStrAny: + try: + r = requests.get(url, headers=HEADERS, params=params) + return r.json() # type: ignore + except requests.HTTPError as e: + if e.response.status_code == 400: + logger.warning( + f"""HTTP Error {e.response.status_code}. + Is your API key authorized to fetch data about the domain + '{params.get("siteUrl")}'?""" + ) + e.response.raise_for_status() + return e.response.json() # type: ignore + + +def get_stats_with_retry(api_path: str, params: DictStrStr) -> List[DictStrAny]: + url = urljoin(BASE_URL, api_path) + response = get_url_with_retry(url, params) + return response.get("d") # type: ignore + + +def parse_response( + response: List[DictStrAny], site_url: str, page: str = None +) -> Iterator[DictStrAny]: + """ + Adds site_url from the request to the response. + Otherwise, we would not know to which site_url a page and its statistics belong. + Further, corrects that what the API returns as 'Query' is actually the 'page'. + """ + for r in response: + if page is None: + # in GetPageStats endpoint the page is under the key "Query" + r.update({"page": r.get("Query")}) + del r["Query"] + else: + r.update({"page": page}) + r.update({"site_url": site_url, "Date": _parse_date(r)}) + del r["__type"] + yield r + + +def _parse_date(record: DictStrStr) -> pendulum.Date: + """Parses Microsoft's date format into a date. The number is a unix timestamp""" + match = re.findall(r"\d+", record.get("Date")) # extract the digits + timestamp_in_seconds = int(match[0]) // 1000 + d: pendulum.Date = pendulum.Date.fromtimestamp(timestamp_in_seconds) + return d diff --git a/omniload/source/bing_webmaster/settings.py b/omniload/source/bing_webmaster/settings.py new file mode 100644 index 000000000..a16561b28 --- /dev/null +++ b/omniload/source/bing_webmaster/settings.py @@ -0,0 +1,18 @@ +# Copyright 2022-2026 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. + +"""Bing Webmaster source settings and constants""" + +BASE_URL = "https://ssl.bing.com/webmaster/api.svc/json/" +HEADERS = {"Content-Type": "application/json", "charset": "utf-8"} diff --git a/omniload/source/chess/adapter.py b/omniload/source/chess/adapter.py index a6f9a9b1a..3e4797a16 100644 --- a/omniload/source/chess/adapter.py +++ b/omniload/source/chess/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,7 +14,7 @@ """A source loading player profiles and games from chess.com api""" -from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence +from typing import Any, Callable, Dict, Iterator, List, Sequence import dlt from dlt.common import pendulum @@ -26,11 +26,9 @@ from .settings import UNOFFICIAL_CHESS_API_URL -@dlt.source(name="chess", max_table_nesting=0) +@dlt.source(name="chess") def source( - players: List[str], - start_month: Optional[str] = None, - end_month: Optional[str] = None, + players: List[str], start_month: str = None, end_month: str = None ) -> Sequence[DltResource]: """ A dlt source for the chess.com api. It groups several resources (in this case chess.com API endpoints) containing @@ -91,12 +89,10 @@ def players_archives(players: List[str]) -> Iterator[List[TDataItem]]: @dlt.resource( - write_disposition="replace", columns={"end_time": {"data_type": "timestamp"}} + write_disposition="append", columns={"end_time": {"data_type": "timestamp"}} ) def players_games( - players: List[str], - start_month: Optional[str] = None, - end_month: Optional[str] = None, + players: List[str], start_month: str = None, end_month: str = None ) -> Iterator[Callable[[], List[TDataItem]]]: """ Yields `players` games that happened between `start_month` and `end_month`. @@ -119,9 +115,10 @@ def players_games( # get archives in parallel by decorating the http request with defer @dlt.defer def _get_archive(url: str) -> List[TDataItem]: + print(f"Getting archive from {url}") try: games = get_url_with_retry(url).get("games", []) - return games + return games # type: ignore except requests.HTTPError as http_err: # sometimes archives are not available and the error seems to be permanent if http_err.response.status_code == 404: diff --git a/omniload/source/chess/helpers.py b/omniload/source/chess/helpers.py index ba9dc7fa7..6e72a7590 100644 --- a/omniload/source/chess/helpers.py +++ b/omniload/source/chess/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,8 +14,6 @@ """Chess source helpers""" -from typing import Optional - from dlt.common.typing import StrAny from dlt.sources.helpers import requests @@ -24,14 +22,14 @@ def get_url_with_retry(url: str) -> StrAny: r = requests.get(url) - return r.json() + return r.json() # type: ignore def get_path_with_retry(path: str) -> StrAny: return get_url_with_retry(f"{OFFICIAL_CHESS_API_URL}{path}") -def validate_month_string(string: Optional[str] = None) -> None: +def validate_month_string(string: str) -> None: """Validates that the string is in YYYY/MM format""" if string and string[4] != "/": raise ValueError(string) diff --git a/omniload/source/chess/settings.py b/omniload/source/chess/settings.py index 0573873e6..1051e0efe 100644 --- a/omniload/source/chess/settings.py +++ b/omniload/source/chess/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/facebook_ads/adapter.py b/omniload/source/facebook_ads/adapter.py index 84b7e0fbd..0d67dadbe 100644 --- a/omniload/source/facebook_ads/adapter.py +++ b/omniload/source/facebook_ads/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,129 +14,60 @@ """Loads campaigns, ads sets, ads, leads and insight data from Facebook Marketing API""" -from typing import Iterator, Sequence, cast +from typing import Iterator, Sequence import dlt from dlt.common import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc -from dlt.common.typing import TDataItems +from dlt.common.typing import DictStrAny, TDataItem, TDataItems from dlt.sources import DltResource -from facebook_business.adobjects.ad import Ad - -from omniload.error import MissingValueError +from facebook_business import FacebookAdsApi +from facebook_business.api import FacebookResponse from .helpers import ( + enrich_ad_objects, execute_job, get_ads_account, get_data_chunked, + get_start_date, process_report_item, ) from .settings import ( ALL_ACTION_ATTRIBUTION_WINDOWS, - ALL_ACTION_BREAKDOWNS, DEFAULT_AD_FIELDS, DEFAULT_ADCREATIVE_FIELDS, DEFAULT_ADSET_FIELDS, DEFAULT_CAMPAIGN_FIELDS, + DEFAULT_INSIGHT_FIELDS, DEFAULT_LEAD_FIELDS, + FACEBOOK_INSIGHTS_RETENTION_PERIOD, INSIGHT_FIELDS_TYPES, + INSIGHTS_BREAKDOWNS_OPTIONS, + INSIGHTS_PRIMARY_KEY, + INVALID_INSIGHTS_FIELDS, + TFbMethod, + TInsightsBreakdownOptions, TInsightsLevels, ) +from .utils import ( + AbstractCrudObject, + AbstractObject, + Ad, + AdCreative, + AdSet, + Campaign, + Lead, + debug_access_token, + get_long_lived_token, +) -def _create_facebook_insights_resource( - accounts: list, - start_date, - end_date, - dimensions: Sequence[str], - fields: Sequence[str], - level, - action_breakdowns: Sequence[str], - batch_size: int, - time_increment_days: int, - action_attribution_windows: Sequence[str], - insights_max_async_sleep_seconds: int, - insights_max_wait_to_finish_seconds: int, - insights_max_wait_to_start_seconds: int, -): - """Create a facebook_insights resource for the given accounts.""" - columns = {} - for field in fields: - if field in INSIGHT_FIELDS_TYPES: - columns[field] = INSIGHT_FIELDS_TYPES[field] - - @dlt.resource( - write_disposition="merge", - merge_key="date_start", - columns=columns, - ) - def facebook_insights( - date_start: dlt.sources.incremental[pendulum.Date] = dlt.sources.incremental( - "date_start", - initial_value=ensure_pendulum_datetime_utc(start_date) - .start_of("day") - .date(), - end_value=ensure_pendulum_datetime_utc(end_date).end_of("day").date() - if end_date - else None, - range_end="closed", - range_start="closed", - ), - ) -> Iterator[TDataItems]: - if date_start.last_value is None: - raise MissingValueError("date_start.last_value", "Facebook Ads") - current_start_date = ensure_pendulum_datetime_utc(date_start.last_value).date() - if date_start.end_value: - end_date_val = pendulum.instance(date_start.end_value) - current_end_date = ( - end_date_val - if isinstance(end_date_val, pendulum.Date) - else end_date_val.date() - ) - else: - current_end_date = pendulum.now().date() - - while current_start_date <= current_end_date: - query = { - "level": level, - "action_breakdowns": list(action_breakdowns), - "breakdowns": dimensions, - "limit": batch_size, - "fields": fields, - "time_increment": time_increment_days, - "action_attribution_windows": list(action_attribution_windows), - "time_ranges": [ - { - "since": current_start_date.to_date_string(), - "until": current_start_date.add( - days=time_increment_days - 1 - ).to_date_string(), - } - ], - } - for account in accounts: - job = execute_job( - account.get_insights(params=query, is_async=True), - insights_max_async_sleep_seconds=insights_max_async_sleep_seconds, - insights_max_wait_to_finish_seconds=insights_max_wait_to_finish_seconds, - insights_max_wait_to_start_seconds=insights_max_wait_to_start_seconds, - ) - output = list(map(process_report_item, job.get_result())) # ty: ignore[unresolved-attribute] - yield output - current_start_date = current_start_date.add(days=time_increment_days) - - return facebook_insights - - -@dlt.source(name="facebook_ads", max_table_nesting=0) +@dlt.source(name="facebook_ads") def facebook_ads_source( - account_id: str | list[str] = dlt.config.value, + account_id: str = dlt.config.value, access_token: str = dlt.secrets.value, chunk_size: int = 50, request_timeout: float = 300.0, - app_api_version: str = "v20.0", - interval_start=None, - interval_end=None, + app_api_version: str = None, ) -> Sequence[DltResource]: """Returns a list of resources to load campaigns, ad sets, ads, creatives and ad leads data from Facebook Marketing API. @@ -148,7 +79,7 @@ def facebook_ads_source( We also provide a transformation `enrich_ad_objects` that you can add to any of the resources to get additional data per object via `object.get_api` Args: - account_id (str | list[str], optional): Account id(s) associated with ad manager. Can be a single ID or a list of IDs. See README.md + account_id (str, optional): Account id associated with add manager. See README.md access_token (str, optional): Access token associated with the Business Facebook App. See README.md chunk_size (int, optional): A size of the page and batch request. You may need to decrease it if you request a lot of fields. Defaults to 50. request_timeout (float, optional): Connection timeout. Defaults to 300.0. @@ -157,134 +88,62 @@ def facebook_ads_source( Returns: Sequence[DltResource]: campaigns, ads, ad_sets, ad_creatives, leads """ - # Convert interval dates to strings if they are datetime objects - if interval_start is not None and hasattr(interval_start, "strftime"): - interval_start = interval_start.strftime("%Y-%m-%d") - if interval_end is not None and hasattr(interval_end, "strftime"): - interval_end = interval_end.strftime("%Y-%m-%d") - - account_ids = account_id if isinstance(account_id, list) else [account_id] - accounts = [ - get_ads_account(acc_id, access_token, request_timeout, app_api_version) - for acc_id in account_ids - ] - - def filter_by_end_date(records, time_field): - if not interval_end: - return records - return [ - r - for r in records - if r.get(time_field) and r[time_field][:10] <= interval_end - ] + account = get_ads_account( + account_id, access_token, request_timeout, app_api_version + ) - @dlt.resource(primary_key="id", write_disposition="merge") + @dlt.resource(primary_key="id", write_disposition="replace") def campaigns( - fields: Sequence[str] = DEFAULT_CAMPAIGN_FIELDS, - states: Sequence[str] | None = None, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "updated_time", initial_value=interval_start - ), + fields: Sequence[str] = DEFAULT_CAMPAIGN_FIELDS, states: Sequence[str] = None ) -> Iterator[TDataItems]: - for account in accounts: - for chunk in get_data_chunked( - account.get_campaigns, fields, states, chunk_size - ): - filtered = filter_by_end_date(chunk, "updated_time") - if filtered: - yield filtered + yield get_data_chunked(account.get_campaigns, fields, states, chunk_size) - @dlt.resource(primary_key="id", write_disposition="merge") + @dlt.resource(primary_key="id", write_disposition="replace") def ads( - fields: Sequence[str] = DEFAULT_AD_FIELDS, - states: Sequence[str] | None = None, + fields: Sequence[str] = DEFAULT_AD_FIELDS, states: Sequence[str] = None ) -> Iterator[TDataItems]: - updated_since = None - if interval_start: - updated_since = int( - cast(pendulum.DateTime, pendulum.parse(interval_start)).timestamp() - ) - for account in accounts: - for chunk in get_data_chunked( - account.get_ads, fields, states, chunk_size, updated_since=updated_since - ): - filtered = filter_by_end_date(chunk, "updated_time") - if filtered: - yield filtered + yield get_data_chunked(account.get_ads, fields, states, chunk_size) - @dlt.resource(primary_key="id", write_disposition="merge") + @dlt.resource(primary_key="id", write_disposition="replace") def ad_sets( - fields: Sequence[str] = DEFAULT_ADSET_FIELDS, - states: Sequence[str] | None = None, + fields: Sequence[str] = DEFAULT_ADSET_FIELDS, states: Sequence[str] = None ) -> Iterator[TDataItems]: - updated_since = None - if interval_start: - updated_since = int( - cast(pendulum.DateTime, pendulum.parse(interval_start)).timestamp() - ) - for account in accounts: - for chunk in get_data_chunked( - account.get_ad_sets, - fields, - states, - chunk_size, - updated_since=updated_since, - ): - filtered = filter_by_end_date(chunk, "updated_time") - if filtered: - yield filtered + yield get_data_chunked(account.get_ad_sets, fields, states, chunk_size) - @dlt.transformer( - primary_key=["id", "created_time"], write_disposition="merge", selected=True - ) + @dlt.transformer(primary_key="id", write_disposition="replace", selected=True) def leads( items: TDataItems, fields: Sequence[str] = DEFAULT_LEAD_FIELDS, - states: Sequence[str] | None = None, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "created_time", initial_value=None - ), + states: Sequence[str] = None, ) -> Iterator[TDataItems]: for item in items: ad = Ad(item["id"]) - for chunk in get_data_chunked(ad.get_leads, fields, states, chunk_size): - filtered = filter_by_end_date(chunk, "created_time") - if filtered: - yield filtered + yield get_data_chunked(ad.get_leads, fields, states, chunk_size) @dlt.resource(primary_key="id", write_disposition="replace") def ad_creatives( - fields: Sequence[str] = DEFAULT_ADCREATIVE_FIELDS, - states: Sequence[str] | None = None, + fields: Sequence[str] = DEFAULT_ADCREATIVE_FIELDS, states: Sequence[str] = None ) -> Iterator[TDataItems]: - for account in accounts: - for chunk in get_data_chunked( - account.get_ad_creatives, fields, states, chunk_size - ): - yield chunk + yield get_data_chunked(account.get_ad_creatives, fields, states, chunk_size) return campaigns, ads, ad_sets, ad_creatives, ads | leads -@dlt.source(name="facebook_ads", max_table_nesting=0) +@dlt.source(name="facebook_ads") def facebook_insights_source( account_id: str = dlt.config.value, access_token: str = dlt.secrets.value, - initial_load_past_days: int = 1, - dimensions: Sequence[str] | None = None, - fields: Sequence[str] | None = None, + initial_load_past_days: int = 30, + fields: Sequence[str] = DEFAULT_INSIGHT_FIELDS, + attribution_window_days_lag: int = 7, time_increment_days: int = 1, - action_breakdowns: Sequence[str] = ALL_ACTION_BREAKDOWNS, - level: TInsightsLevels | None = "ad", + breakdowns: TInsightsBreakdownOptions = None, + action_breakdowns: Sequence[str] = None, + level: TInsightsLevels = "ad", action_attribution_windows: Sequence[str] = ALL_ACTION_ATTRIBUTION_WINDOWS, batch_size: int = 50, request_timeout: int = 300, - app_api_version: str | None = None, - start_date: pendulum.DateTime | None = None, - end_date: pendulum.DateTime | None = None, - insights_max_wait_to_finish_seconds: int = 60 * 60 * 4, - insights_max_wait_to_start_seconds: int = 60 * 30, - insights_max_async_sleep_seconds: int = 20, + app_api_version: str = None, ) -> DltResource: """Incrementally loads insight reports with defined granularity level, fields, breakdowns etc. @@ -301,8 +160,8 @@ def facebook_insights_source( fields (Sequence[str], optional): A list of fields to include in each reports. Note that `breakdowns` option adds fields automatically. Defaults to DEFAULT_INSIGHT_FIELDS. attribution_window_days_lag (int, optional): Attribution window in days. The reports in attribution window are refreshed on each run.. Defaults to 7. time_increment_days (int, optional): The report aggregation window in days. use 7 for weekly aggregation. Defaults to 1. - breakdowns (TInsightsBreakdownOptions, optional): A presents with common aggregations. See settings.py for details. Defaults to "ads_insights_age_and_gender". - action_breakdowns (Sequence[str], optional): Action aggregation types. See settings.py for details. Defaults to ALL_ACTION_BREAKDOWNS. + breakdowns (TInsightsBreakdownOptions, optional): A presents with common aggregations. See settings.py for details. Defaults to None (no breakdowns). + action_breakdowns (Sequence[str], optional): Action aggregation types. See settings.py for details. Defaults to None (no action breakdowns). level (TInsightsLevels, optional): The granularity level. Defaults to "ad". action_attribution_windows (Sequence[str], optional): Attribution windows for actions. Defaults to ALL_ACTION_ATTRIBUTION_WINDOWS. batch_size (int, optional): Page size when reading data from particular report. Defaults to 50. @@ -317,97 +176,57 @@ def facebook_insights_source( account_id, access_token, request_timeout, app_api_version ) - if start_date is None: - start_date = pendulum.today().subtract(days=initial_load_past_days) - - if dimensions is None: - dimensions = [] - if fields is None: - fields = [] + # we load with a defined lag + initial_load_start_date = pendulum.today().subtract(days=initial_load_past_days) + initial_load_start_date_str = initial_load_start_date.isoformat() - return _create_facebook_insights_resource( - accounts=[account], - start_date=start_date, - end_date=end_date, - dimensions=dimensions, - fields=fields, - level=level, - action_breakdowns=action_breakdowns, - batch_size=batch_size, - time_increment_days=time_increment_days, - action_attribution_windows=action_attribution_windows, - insights_max_async_sleep_seconds=insights_max_async_sleep_seconds, - insights_max_wait_to_finish_seconds=insights_max_wait_to_finish_seconds, - insights_max_wait_to_start_seconds=insights_max_wait_to_start_seconds, + @dlt.resource( + primary_key=INSIGHTS_PRIMARY_KEY, + write_disposition="merge", + columns=INSIGHT_FIELDS_TYPES, ) + def facebook_insights( + date_start: dlt.sources.incremental[str] = dlt.sources.incremental( + "date_start", initial_value=initial_load_start_date_str + ), + ) -> Iterator[TDataItems]: + start_date = get_start_date(date_start, attribution_window_days_lag) + end_date = pendulum.now() + # fetch insights in incremental day steps + while start_date <= end_date: + query = { + "level": level, + "limit": batch_size, + "time_increment": time_increment_days, + "action_attribution_windows": list(action_attribution_windows), + "time_ranges": [ + { + "since": start_date.to_date_string(), + "until": start_date.add( + days=time_increment_days - 1 + ).to_date_string(), + } + ], + } -@dlt.source(name="facebook_ads", max_table_nesting=0) -def facebook_insights_with_account_ids_source( - account_ids: list[str], - access_token: str = dlt.secrets.value, - initial_load_past_days: int = 1, - dimensions: Sequence[str] | None = None, - fields: Sequence[str] | None = None, - time_increment_days: int = 1, - action_breakdowns: Sequence[str] = ALL_ACTION_BREAKDOWNS, - level: TInsightsLevels | None = "ad", - action_attribution_windows: Sequence[str] = ALL_ACTION_ATTRIBUTION_WINDOWS, - batch_size: int = 50, - request_timeout: int = 300, - app_api_version: str | None = None, - start_date: pendulum.DateTime | None = None, - end_date: pendulum.DateTime | None = None, - insights_max_wait_to_finish_seconds: int = 60 * 60 * 4, - insights_max_wait_to_start_seconds: int = 60 * 30, - insights_max_async_sleep_seconds: int = 20, -) -> DltResource: - """Incrementally loads insight reports for multiple account IDs. - - Args: - account_ids (list[str]): List of account IDs to fetch insights for. - access_token (str): Access token associated with the Business Facebook App. - initial_load_past_days (int, optional): How many past days to initially load. Defaults to 1. - dimensions (Sequence[str], optional): Breakdown dimensions. - fields (Sequence[str], optional): Fields to include in reports. - time_increment_days (int, optional): Report aggregation window in days. Defaults to 1. - action_breakdowns (Sequence[str], optional): Action aggregation types. - level (TInsightsLevels, optional): Granularity level. Defaults to "ad". - action_attribution_windows (Sequence[str], optional): Attribution windows for actions. - batch_size (int, optional): Page size. Defaults to 50. - request_timeout (int, optional): Connection timeout. Defaults to 300. - app_api_version (str, optional): Facebook API version. - start_date: Start date for insights. - end_date: End date for insights. - - Returns: - DltResource: facebook_insights - """ - accounts = [ - get_ads_account(acc_id, access_token, request_timeout, app_api_version) - for acc_id in account_ids - ] + fields_to_use = set(fields) + # Only add breakdowns if explicitly provided + if breakdowns is not None: + query["breakdowns"] = list( + INSIGHTS_BREAKDOWNS_OPTIONS[breakdowns]["breakdowns"] + ) + fields_to_use = fields_to_use.union( + INSIGHTS_BREAKDOWNS_OPTIONS[breakdowns]["fields"] + ) + query["fields"] = list(fields_to_use.difference(INVALID_INSIGHTS_FIELDS)) - if start_date is None: - start_date = pendulum.today().subtract(days=initial_load_past_days) + # Only add action_breakdowns if explicitly provided + if action_breakdowns is not None: + query["action_breakdowns"] = list(action_breakdowns) - if dimensions is None: - dimensions = [] - if fields is None: - fields = [] + job = execute_job(account.get_insights(params=query, is_async=True)) + yield list(map(process_report_item, job.get_result())) + start_date = start_date.add(days=time_increment_days) - return _create_facebook_insights_resource( - accounts=accounts, - start_date=start_date, - end_date=end_date, - dimensions=dimensions, - fields=fields, - level=level, - action_breakdowns=action_breakdowns, - batch_size=batch_size, - time_increment_days=time_increment_days, - action_attribution_windows=action_attribution_windows, - insights_max_async_sleep_seconds=insights_max_async_sleep_seconds, - insights_max_wait_to_finish_seconds=insights_max_wait_to_finish_seconds, - insights_max_wait_to_start_seconds=insights_max_wait_to_start_seconds, - ) + return facebook_insights diff --git a/omniload/source/facebook_ads/exceptions.py b/omniload/source/facebook_ads/exceptions.py index d88334613..9906934a1 100644 --- a/omniload/source/facebook_ads/exceptions.py +++ b/omniload/source/facebook_ads/exceptions.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/facebook_ads/helpers.py b/omniload/source/facebook_ads/helpers.py index 5676ece01..c32494bfd 100644 --- a/omniload/source/facebook_ads/helpers.py +++ b/omniload/source/facebook_ads/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,56 +17,78 @@ import functools import itertools import time -from datetime import datetime -from typing import Any, Iterator, Optional, Sequence, cast +from typing import Any, Iterator, Sequence +import dlt import humanize import pendulum from dlt.common import logger from dlt.common.configuration.inject import with_config +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import DictStrAny, TDataItem, TDataItems from dlt.sources.helpers import requests from dlt.sources.helpers.requests import Client from facebook_business import FacebookAdsApi -from facebook_business.adobjects.abstractcrudobject import AbstractCrudObject -from facebook_business.adobjects.abstractobject import AbstractObject from facebook_business.adobjects.adaccount import AdAccount from facebook_business.adobjects.user import User from facebook_business.api import FacebookResponse from .exceptions import InsightsJobTimeout from .settings import ( + FACEBOOK_INSIGHTS_RETENTION_PERIOD, INSIGHTS_PRIMARY_KEY, TFbMethod, ) +from .utils import AbstractCrudObject, AbstractObject -def process_report_item(item: AbstractObject) -> DictStrAny: - if "date_start" in item: - item["date_start"] = datetime.strptime(item["date_start"], "%Y-%m-%d").date() - if "date_stop" in item: - item["date_stop"] = datetime.strptime(item["date_stop"], "%Y-%m-%d").date() +def get_start_date( + incremental_start_date: dlt.sources.incremental[str], + attribution_window_days_lag: int = 7, +) -> pendulum.DateTime: + """ + Get the start date for incremental loading of Facebook Insights data. + """ + start_date: pendulum.DateTime = ensure_pendulum_datetime( + incremental_start_date.start_value + ).subtract(days=attribution_window_days_lag) + + # facebook forgets insights so trim the lag and warn + min_start_date = pendulum.today().subtract( + months=FACEBOOK_INSIGHTS_RETENTION_PERIOD + ) + if start_date < min_start_date: + logger.warning( + "%s: Start date is earlier than %s months ago, using %s instead. " + "For more information, see https://www.facebook.com/business/help/1695754927158071?id=354406972049255", + "facebook_insights", + FACEBOOK_INSIGHTS_RETENTION_PERIOD, + min_start_date, + ) + start_date = min_start_date + incremental_start_date.start_value = min_start_date + + # lag the incremental start date by attribution window lag + incremental_start_date.start_value = start_date.isoformat() + return start_date + +def process_report_item(item: AbstractObject) -> DictStrAny: d: DictStrAny = item.export_all_data() for pki in INSIGHTS_PRIMARY_KEY: if pki not in d: d[pki] = "no_" + pki + return d def get_data_chunked( - method: TFbMethod, - fields: Sequence[str], - states: Sequence[str] | None, - chunk_size: int, - updated_since: Optional[int] = None, + method: TFbMethod, fields: Sequence[str], states: Sequence[str], chunk_size: int ) -> Iterator[TDataItems]: # add pagination and chunk into lists params: DictStrAny = {"limit": chunk_size} if states: params.update({"effective_status": states}) - if updated_since: - params.update({"updated_since": updated_since}) it: map[DictStrAny] = map( lambda c: c.export_all_data(), method(fields=fields, params=params) ) @@ -104,8 +126,8 @@ def fail(resp: FacebookResponse) -> None: raise resp.error() for item in items: - o: AbstractCrudObject = fb_obj_type(item["id"]) # ty: ignore[call-non-callable] - o.api_get( # ty: ignore[unresolved-attribute] + o: AbstractCrudObject = fb_obj_type(item["id"]) + o.api_get( fields=fields, batch=api_batch, success=functools.partial(update_item, item=item), @@ -127,12 +149,12 @@ def execute_job( insights_max_wait_to_finish_seconds: int = 30 * 60, insights_max_async_sleep_seconds: int = 5 * 60, ) -> AbstractCrudObject: - status: str = "Unknown" + status: str = None time_start = time.time() - sleep_time = 3 + sleep_time = 10 while status != "Job Completed": duration = time.time() - time_start - job = job.api_get() # ty: ignore[unresolved-attribute] + job = job.api_get() status = job["async_status"] percent_complete = job["async_percent_completion"] @@ -159,7 +181,7 @@ def execute_job( raise InsightsJobTimeout( "facebook_insights", pretty_error_message.format( - job_id, insights_max_wait_to_finish_seconds + job_id, insights_max_wait_to_finish_seconds // 60 ), ) @@ -170,10 +192,9 @@ def execute_job( return job -def _init_facebook_api( - access_token: str, request_timeout: float, app_api_version: Optional[str] = None -) -> None: - """Initialize Facebook API with retry session.""" +def get_ads_account( + account_id: str, access_token: str, request_timeout: float, app_api_version: str +) -> AdAccount: notify_on_token_expiration() def retry_on_limit(response: requests.Response, exception: BaseException) -> bool: @@ -206,40 +227,34 @@ def retry_on_limit(response: requests.Response, exception: BaseException) -> boo retry_session = Client( request_timeout=request_timeout, raise_for_status=False, - retry_condition=retry_on_limit, # ty: ignore[invalid-argument-type] + retry_condition=retry_on_limit, request_max_attempts=12, request_backoff_factor=2, ).session retry_session.params.update({"access_token": access_token}) # type: ignore + # patch dlt requests session with retries API = FacebookAdsApi.init( + account_id="act_" + account_id, access_token=access_token, api_version=app_api_version, ) API._session.requests = retry_session + user = User(fbid="me") + accounts = user.get_ad_accounts() + account: AdAccount = None + for acc in accounts: + if acc["account_id"] == account_id: + account = acc -def get_ads_account( - account_id: str, - access_token: str, - request_timeout: float, - app_api_version: Optional[str] = None, -) -> AdAccount: - """Get a specific ad account by ID.""" - _init_facebook_api(access_token, request_timeout, app_api_version) - return AdAccount(f"act_{account_id}") - + if not account: + raise ValueError("Couldn't find account with id {}".format(account_id)) -def get_all_ad_accounts( - access_token: str, request_timeout: float, app_api_version: str -) -> list[AdAccount]: - """Get all ad accounts for the authenticated user.""" - _init_facebook_api(access_token, request_timeout, app_api_version) - user = User(fbid="me") - return list(user.get_ad_accounts()) + return account @with_config(sections=("sources", "facebook_ads")) -def notify_on_token_expiration(access_token_expires_at: Optional[int] = None) -> None: +def notify_on_token_expiration(access_token_expires_at: int = None) -> None: """Notifies (currently via logger) if access token expires in less than 7 days. Needs `access_token_expires_at` to be configured.""" if not access_token_expires_at: logger.warning( @@ -251,49 +266,3 @@ def notify_on_token_expiration(access_token_expires_at: Optional[int] = None) -> logger.error( f"Access Token expires in {humanize.precisedelta(pendulum.now() - expires_at)}. Replace the token now!" ) - - -def parse_insights_table_to_source_kwargs(table: str) -> DictStrAny: - import typing - - from omniload.source.facebook_ads.settings import ( - INSIGHTS_BREAKDOWNS_OPTIONS, - TInsightsBreakdownOptions, - TInsightsLevels, - ) - - parts = table.split(":") - - source_kwargs = {} - - breakdown_type = cast(TInsightsBreakdownOptions, parts[1]) - - valid_breakdowns = list(typing.get_args(TInsightsBreakdownOptions)) - if breakdown_type in valid_breakdowns: - dimensions = INSIGHTS_BREAKDOWNS_OPTIONS[breakdown_type]["breakdowns"] - fields = INSIGHTS_BREAKDOWNS_OPTIONS[breakdown_type]["fields"] - source_kwargs["dimensions"] = dimensions - source_kwargs["fields"] = fields - else: - dimensions = breakdown_type.split(",") - valid_levels = list(typing.get_args(TInsightsLevels)) - level = None - for valid_level in reversed(valid_levels): - if valid_level in dimensions: - level = valid_level - dimensions.remove(valid_level) - break - - source_kwargs["level"] = level - source_kwargs["dimensions"] = dimensions - - # If custom metrics are provided, parse them - if len(parts) == 3: - fields = [f.strip() for f in parts[2].split(",") if f.strip()] - if not fields: - raise ValueError( - "Custom metrics must be provided after the second colon in format: facebook_insights:breakdown_type:metric1,metric2..." - ) - source_kwargs["fields"] = fields - - return source_kwargs diff --git a/omniload/source/facebook_ads/settings.py b/omniload/source/facebook_ads/settings.py index 4baf8a6b0..e917dd45e 100644 --- a/omniload/source/facebook_ads/settings.py +++ b/omniload/source/facebook_ads/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -111,23 +111,6 @@ "action_values", "cost_per_action_type", "website_ctr", - "account_currency", - "ad_click_actions", - "ad_name", - "adset_name", - "campaign_name", - "country", - "dma", - "full_view_impressions", - "full_view_reach", - "inline_link_click_ctr", - "outbound_clicks", - "reach", - "social_spend", - "spend", - "website_ctr", - "conversions", - "video_thruplay_watched_actions", ) TInsightsLevels = Literal["account", "campaign", "adset", "ad"] diff --git a/omniload/source/facebook_ads/utils.py b/omniload/source/facebook_ads/utils.py index 7b78b1ab7..a0d432bd2 100644 --- a/omniload/source/facebook_ads/utils.py +++ b/omniload/source/facebook_ads/utils.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,8 +15,17 @@ from typing import Dict import dlt +from dlt.common import logger, pendulum from dlt.common.configuration.inject import with_config from dlt.sources.helpers import requests +from facebook_business import FacebookAdsApi +from facebook_business.adobjects.abstractcrudobject import AbstractCrudObject +from facebook_business.adobjects.abstractobject import AbstractObject +from facebook_business.adobjects.ad import Ad +from facebook_business.adobjects.adcreative import AdCreative +from facebook_business.adobjects.adset import AdSet +from facebook_business.adobjects.campaign import Campaign +from facebook_business.adobjects.lead import Lead @with_config(sections=("sources", "facebook_ads")) diff --git a/omniload/source/freshdesk/adapter.py b/omniload/source/freshdesk/adapter.py index d4612a18f..380afab3a 100644 --- a/omniload/source/freshdesk/adapter.py +++ b/omniload/source/freshdesk/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,23 +18,18 @@ from typing import Any, Dict, Generator, Iterable, List, Optional import dlt -import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc from dlt.sources import DltResource -from .client import FreshdeskClient +from .freshdesk_client import FreshdeskClient from .settings import DEFAULT_ENDPOINTS @dlt.source() def freshdesk_source( - domain: str, - api_secret_key: str, - start_date: pendulum.DateTime, - end_date: Optional[pendulum.DateTime] = None, - per_page: int = 100, endpoints: Optional[List[str]] = None, - query: Optional[str] = None, + per_page: int = 100, + domain: str = dlt.secrets.value, + api_secret_key: str = dlt.secrets.value, ) -> Iterable[DltResource]: """ Retrieves data from specified Freshdesk API endpoints. @@ -58,11 +53,7 @@ def freshdesk_source( def incremental_resource( endpoint: str, updated_at: Optional[Any] = dlt.sources.incremental( - "updated_at", - initial_value=start_date.isoformat(), - end_value=end_date.isoformat() if end_date else None, - range_start="closed", - range_end="closed", + "updated_at", initial_value="2022-01-01T00:00:00Z" ), ) -> Generator[Dict[Any, Any], Any, None]: """ @@ -71,22 +62,19 @@ def incremental_resource( to ensure incremental loading. """ - effective_start_date = start_date - if updated_at is not None and updated_at.last_value is not None: - effective_start_date = ensure_pendulum_datetime_utc(updated_at.last_value) + # Retrieve the last updated timestamp to fetch only new or updated records. + updated_at = updated_at.last_value - if updated_at is not None and updated_at.end_value is not None: - end_date = ensure_pendulum_datetime_utc(updated_at.end_value) - else: - end_date = pendulum.now(tz="UTC") + # Endpoint and last updated timestamp. + print( + f"Fetching data from endpoint: {endpoint} and the last `updated_at` is: {updated_at}" + ) # Use the FreshdeskClient instance to fetch paginated responses yield from freshdesk.paginated_response( endpoint=endpoint, per_page=per_page, - start_date=effective_start_date, - end_date=end_date, - query=query, + updated_at=updated_at, ) # Set default endpoints if not provided diff --git a/omniload/source/freshdesk/client.py b/omniload/source/freshdesk/client.py index 82d11e29d..2f2dc8844 100644 --- a/omniload/source/freshdesk/client.py +++ b/omniload/source/freshdesk/client.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,16 +16,11 @@ import logging import time -from typing import Any, Dict, Iterable, Optional, cast +from typing import Any, Dict, Iterable, Optional -import pendulum from dlt.common.typing import TDataItem from dlt.sources.helpers import requests -from omniload.error import HTTPError - -TICKETS_QUERY_MAX_PAGE = 10 - class FreshdeskClient: """ @@ -80,15 +75,13 @@ def _request_with_rate_limit(self, url: str, **kwargs: Any) -> requests.Response else: # If the error is not a rate limit (429), raise the exception to be # handled elsewhere or stop execution - raise HTTPError(e) from e + raise def paginated_response( self, endpoint: str, per_page: int, - start_date: pendulum.DateTime, - end_date: pendulum.DateTime, - query: Optional[str] = None, + updated_at: Optional[str] = None, ) -> Iterable[TDataItem]: """ Fetches a paginated response from a specified endpoint. @@ -98,11 +91,6 @@ def paginated_response( updated at the specified timestamp. """ page = 1 - if query is not None: - query = query.replace('"', "").strip() - - is_tickets_query = query and endpoint == "tickets" - while True: # Construct the URL for the specific endpoint url = f"{self.base_url}/{endpoint}" @@ -114,39 +102,15 @@ def paginated_response( param_key = ( "updated_since" if endpoint == "tickets" else "_updated_since" ) - - params[param_key] = start_date.to_iso8601_string() - - if is_tickets_query: - url = f"{self.base_url}/search/tickets" - params = { - "query": f'"{query}"', - "page": page, - } + if updated_at: + params[param_key] = updated_at # Handle requests with rate-limiting # A maximum of 300 pages (30000 tickets) will be returned. response = self._request_with_rate_limit(url, params=params) data = response.json() - if query and endpoint == "tickets": - data = data["results"] - if not data: break # Stop if no data or max page limit reached - - filtered_data = [ - item - for item in data - if "updated_at" in item - and cast(pendulum.DateTime, pendulum.parse(item["updated_at"])) - <= end_date - ] - if not filtered_data: - break - yield filtered_data + yield data page += 1 - - # https://developers.freshdesk.com/api/#filter_tickets - if is_tickets_query and page > TICKETS_QUERY_MAX_PAGE: - break diff --git a/omniload/source/freshdesk/settings.py b/omniload/source/freshdesk/settings.py index b3d21d861..47eead4b2 100644 --- a/omniload/source/freshdesk/settings.py +++ b/omniload/source/freshdesk/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/github/adapter.py b/omniload/source/github/adapter.py index 042406d4b..1e16a9f3f 100644 --- a/omniload/source/github/adapter.py +++ b/omniload/source/github/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,23 +15,20 @@ """Source that load github issues, pull requests and reactions for a specific repository via customizable graphql query. Loads events incrementally.""" import urllib.parse -from typing import Iterator, Optional, Sequence, cast +from typing import Iterator, Optional, Sequence import dlt -import pendulum from dlt.common.typing import TDataItems from dlt.sources import DltResource -from omniload.error import MissingValueError - from .helpers import get_reactions_data, get_rest_pages, get_stargazers -@dlt.source(max_table_nesting=0) +@dlt.source def github_reactions( owner: str, name: str, - access_token: str, + access_token: str = dlt.secrets.value, items_per_page: int = 100, max_items: Optional[int] = None, ) -> Sequence[DltResource]: @@ -82,13 +79,9 @@ def github_reactions( ) -@dlt.source(max_table_nesting=0) +@dlt.source(max_table_nesting=2) def github_repo_events( - owner: str, - name: str, - access_token: str, - start_date: pendulum.DateTime, - end_date: Optional[pendulum.DateTime] = None, + owner: str, name: str, access_token: Optional[str] = None ) -> DltResource: """Gets events for repository `name` with owner `owner` incrementally. @@ -107,58 +100,18 @@ def github_repo_events( """ # use naming function in table name to generate separate tables for each event - @dlt.resource( - primary_key="id", table_name=lambda i: i["type"], write_disposition="merge" - ) + @dlt.resource(primary_key="id", table_name=lambda i: i["type"]) def repo_events( last_created_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "created_at", - initial_value=start_date.isoformat(), - end_value=end_date.isoformat() if end_date else None, - last_value_func=max, - range_end="closed", - range_start="closed", + "created_at", initial_value="1970-01-01T00:00:00Z", last_value_func=max ), ) -> Iterator[TDataItems]: repos_path = ( f"/repos/{urllib.parse.quote(owner)}/{urllib.parse.quote(name)}/events" ) - # Get the date range from the incremental state - start_filter_value = last_created_at.last_value or last_created_at.initial_value - if start_filter_value is None: - raise MissingValueError("start_filter_value", "GitHub") - start_filter = cast(pendulum.DateTime, pendulum.parse(start_filter_value)) - end_filter = cast( - pendulum.DateTime, - ( - pendulum.parse(last_created_at.end_value) - if last_created_at.end_value - else pendulum.now() - ), - ) - for page in get_rest_pages(access_token, repos_path + "?per_page=100"): - # Filter events by date range - filtered_events = [] - for event in page: - event_date = cast( - pendulum.DateTime, pendulum.parse(event["created_at"]) - ) - - # Check if event is within the date range - if event_date >= start_filter: - if end_filter is None or event_date <= end_filter: - filtered_events.append(event) - elif event_date > end_filter: - # Skip events that are newer than our end date - continue - else: - # Events are ordered by date desc, so if we hit an older event, we can stop - break - - if filtered_events: - yield filtered_events + yield page # stop requesting pages if the last element was already older than initial value # note: incremental will skip those items anyway, we just do not want to use the api limits @@ -171,11 +124,11 @@ def repo_events( return repo_events -@dlt.source(max_table_nesting=0) +@dlt.source def github_stargazers( owner: str, name: str, - access_token: str, + access_token: str = dlt.secrets.value, items_per_page: int = 100, max_items: Optional[int] = None, ) -> Sequence[DltResource]: diff --git a/omniload/source/github/helpers.py b/omniload/source/github/helpers.py index a68a655d2..59cc4117d 100644 --- a/omniload/source/github/helpers.py +++ b/omniload/source/github/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Iterator, List, Optional, Tuple, cast +from typing import Iterator, List, Optional, Tuple from dlt.common.typing import DictStrAny, StrAny from dlt.common.utils import chunks @@ -117,11 +117,11 @@ def get_reactions_data( def _extract_top_connection(data: StrAny, node_type: str) -> StrAny: - assert isinstance(data, dict) and len(data) == 1, ( # noqa: S101 + assert isinstance(data, dict) and len(data) == 1, ( f"The data with list of {node_type} must be a dictionary and contain only one element" ) - data = cast(StrAny, next(iter(data.values()))) - return data[node_type] + data = next(iter(data.values())) + return data[node_type] # type: ignore def _extract_nested_nodes(item: DictStrAny) -> DictStrAny: @@ -159,11 +159,7 @@ def _request() -> requests.Response: def _get_graphql_pages( - access_token: str, - query: str, - variables: DictStrAny, - node_type: str, - max_items: Optional[int] = None, + access_token: str, query: str, variables: DictStrAny, node_type: str, max_items: int ) -> Iterator[List[DictStrAny]]: items_count = 0 while True: @@ -186,7 +182,7 @@ def _get_graphql_pages( variables["page_after"] = _extract_top_connection(data, node_type)["pageInfo"][ "endCursor" ] - if max_items is not None and items_count >= max_items: + if max_items and items_count >= max_items: print(f"Max items limit reached: {items_count} >= {max_items}") return diff --git a/omniload/source/github/queries.py b/omniload/source/github/queries.py index 598f81a7f..1ef85bd5a 100644 --- a/omniload/source/github/queries.py +++ b/omniload/source/github/queries.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/github/settings.py b/omniload/source/github/settings.py index cdea84134..aaf246638 100644 --- a/omniload/source/github/settings.py +++ b/omniload/source/github/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/google_ads/adapter.py b/omniload/source/google_ads/adapter.py index ba47f1aa3..ff399e807 100644 --- a/omniload/source/google_ads/adapter.py +++ b/omniload/source/google_ads/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,187 +12,167 @@ # See the License for the specific language governing permissions and # limitations under the License. -from datetime import date, datetime -from typing import Any, Iterator, Optional +""" +Preliminary implementation of Google Ads pipeline. +""" + +import json +import tempfile +from typing import Iterator, List, Union import dlt -import proto -from dlt.common import json +from apiclient.discovery import Resource from dlt.common.exceptions import MissingDependencyException from dlt.common.typing import TDataItem from dlt.sources import DltResource -from flatten_json import flatten +from dlt.sources.credentials import GcpOAuthCredentials, GcpServiceAccountCredentials -from . import field -from .metrics import dlt_metrics_schema -from .predicates import date_predicate -from .reports import BUILTIN_REPORTS, Report +from .helpers.data_processing import to_dict try: - from google.ads.googleads.client import GoogleAdsClient + from google.ads.googleads.client import GoogleAdsClient # type: ignore except ImportError: raise MissingDependencyException("Requests-OAuthlib", ["google-ads"]) -@dlt.source +DIMENSION_TABLES = [ + "accounts", + "ad_group", + "ad_group_ad", + "ad_group_ad_label", + "ad_group_label", + "campaign_label", + "click_view", + "customer", + "keyword_view", + "geographic_view", +] + + +def get_client( + credentials: Union[GcpOAuthCredentials, GcpServiceAccountCredentials], + dev_token: str, + impersonated_email: str, +) -> GoogleAdsClient: + # generate access token for credentials if we are using OAuth2.0 + if isinstance(credentials, GcpOAuthCredentials): + credentials.auth("https://www.googleapis.com/auth/adwords") + conf = { + "developer_token": dev_token, + "use_proto_plus": True, + **json.loads(credentials.to_native_representation()), + } + return GoogleAdsClient.load_from_dict(config_dict=conf) + # use service account to authenticate if not using OAuth2.0 + else: + # google ads client requires the key to be in a file on the disc.. + with tempfile.NamedTemporaryFile() as f: + f.write(credentials.to_native_representation().encode()) + f.seek(0) + return GoogleAdsClient.load_from_dict( + config_dict={ + "json_key_file_path": f.name, + "impersonated_email": impersonated_email, + "use_proto_plus": True, + "developer_token": dev_token, + } + ) + + +@dlt.source(max_table_nesting=2) def google_ads( - client: GoogleAdsClient, - customer_ids: list[str], - report_spec: Optional[str] = None, - gaql_query: Optional[str] = None, - start_date: Optional[datetime] = None, - end_date: Optional[datetime] = None, -) -> Iterator[DltResource]: - date_range = dlt.sources.incremental( - "segments_date", - initial_value=start_date.date(), # type: ignore - end_value=end_date.date() if end_date is not None else None, - range_start="closed", - range_end="closed", + credentials: Union[ + GcpOAuthCredentials, GcpServiceAccountCredentials + ] = dlt.secrets.value, + impersonated_email: str = dlt.secrets.value, + dev_token: str = dlt.secrets.value, +) -> List[DltResource]: + """ + Loads default tables for google ads in the database. + :param credentials: + :param dev_token: + :return: + """ + client = get_client( + credentials=credentials, + dev_token=dev_token, + impersonated_email=impersonated_email, ) - if report_spec is not None: - custom_report, _ = Report.from_spec(report_spec) - yield dlt.resource( - daily_report, - name="daily_report", - write_disposition="merge", - primary_key=custom_report.primary_keys() + ["customer_id"], - columns=dlt_metrics_schema(custom_report.metrics), - )(client, customer_ids, custom_report, date_range) - - if gaql_query is not None: - yield dlt.resource( - run_gaql_query, - name="gaql_query", - write_disposition="append", - max_table_nesting=0, - )(client, customer_ids, gaql_query, start_date, end_date) - - for report_name, report in BUILTIN_REPORTS.items(): - yield dlt.resource( - daily_report, - name=report_name, - write_disposition="merge", - primary_key=report.primary_keys() + ["customer_id"], - columns=dlt_metrics_schema(report.metrics), - )(client, customer_ids, report, date_range) - - -def daily_report( - client: GoogleAdsClient, - customer_ids: list[str], - report: Report, - date: dlt.sources.incremental[date], + return [ + customers(client=client), + campaigns(client=client), + change_events(client=client), + customer_clients(client=client), + ] + + +@dlt.resource(write_disposition="replace") +def customers( + client: Resource, customer_id: str = dlt.secrets.value ) -> Iterator[TDataItem]: + """ + Dlt resource which loads dimensions. + :param client: + :return: + """ + # Issues a search request using streaming. ga_service = client.get_service("GoogleAdsService") - fields = report.dimensions + report.metrics + report.segments - criteria = date_predicate("segments.date", date.last_value, date.end_value) # type:ignore - # TODO: Review "Possible SQL injection vector through string-based query construction" - query = f""" - SELECT - {", ".join(fields)} - FROM - {report.resource} - WHERE - {criteria} - """ # noqa: S608 - if report.unfilterable is True: - i = query.index("WHERE", 0) - query = query[:i] - - allowed_keys = set([field.to_column(k) for k in fields]) - for customer_id in customer_ids: - stream = ga_service.search_stream(customer_id=customer_id, query=query) - for batch in stream: - for row in batch.results: - data = flatten(merge_lists(to_dict(row))) - if "segments_date" in data: - data["segments_date"] = datetime.strptime( - data["segments_date"], "%Y-%m-%d" - ).date() - row_data = {k: v for k, v in data.items() if k in allowed_keys} - for pk in report.primary_keys(): - if pk not in row_data or row_data[pk] is None or row_data[pk] == "": - row_data[pk] = "-" - row_data["customer_id"] = customer_id - yield row_data - - -def to_dict(item: Any) -> TDataItem: + query = "SELECT customer.id, customer.descriptive_name FROM customer" + stream = ga_service.search_stream(customer_id=customer_id, query=query) + for batch in stream: + for row in batch.results: + yield to_dict(row.customer) + + +@dlt.resource(write_disposition="replace") +def campaigns( + client: Resource, customer_id: str = dlt.secrets.value +) -> Iterator[TDataItem]: """ - Processes a batch result (page of results per dimension) accordingly - :param batch: + Dlt resource which loads dimensions. + :param client: :return: """ - return json.loads( - proto.Message.to_json( - item, - preserving_proto_field_name=True, - use_integers_for_enums=False, - including_default_value_fields=False, - ) - ) + # Issues a search request using streaming. + ga_service = client.get_service("GoogleAdsService") + query = "SELECT campaign.id, campaign.labels FROM campaign" + stream = ga_service.search_stream(customer_id=customer_id, query=query) + for batch in stream: + for row in batch.results: + yield to_dict(row.campaign) -def merge_lists(item: dict) -> dict: - replacements = {} - for k, v in item.get("metrics", {}).items(): - if isinstance(v, list): - replacements[k] = ",".join(v) - if len(replacements) == 0: - return item - item["metrics"].update(replacements) - return item - - -def extract_fields(data: dict, field_paths: list[str]) -> dict: - result = {} - for path in field_paths: - parts = path.split(".") - value: Any = data - for part in parts: - if isinstance(value, dict): - value = value.get(part) if part in value else value.get(part + "_") - else: - value = None - break - column_name = path.replace(".", "_") - result[column_name] = value - return result - - -def run_gaql_query( - client: GoogleAdsClient, - customer_ids: list[str], - query: str, - start_date: Optional[datetime] = None, - end_date: Optional[datetime] = None, +@dlt.resource(write_disposition="replace") +def change_events( + client: Resource, customer_id: str = dlt.secrets.value ) -> Iterator[TDataItem]: """ - Execute a raw Google Ads Query Language (GAQL) query. - Supports :interval_start and :interval_end placeholders for date filtering. + Dlt resource which loads dimensions. + :param client: + :return: """ + # Issues a search request using streaming. ga_service = client.get_service("GoogleAdsService") + query = "SELECT change_event.change_date_time FROM change_event WHERE change_event.change_date_time during LAST_14_DAYS LIMIT 1000" + stream = ga_service.search_stream(customer_id=customer_id, query=query) + for batch in stream: + for row in batch.results: + yield to_dict(row.change_event) + - if ":interval_start" in query: - start_str = start_date.strftime("%Y-%m-%d") if start_date else "1970-01-01" - query = query.replace(":interval_start", f"'{start_str}'") - - if ":interval_end" in query: - end_str = ( - end_date.strftime("%Y-%m-%d") - if end_date - else date.today().strftime("%Y-%m-%d") - ) - query = query.replace(":interval_end", f"'{end_str}'") - - field_paths = None - for customer_id in customer_ids: - stream = ga_service.search_stream(customer_id=customer_id, query=query) - for batch in stream: - if field_paths is None: - field_paths = list(batch.field_mask.paths) - for row in batch.results: - data = extract_fields(to_dict(row), field_paths) - data["customer_id"] = customer_id - yield data +@dlt.resource(write_disposition="replace") +def customer_clients( + client: Resource, customer_id: str = dlt.secrets.value +) -> Iterator[TDataItem]: + """ + Dlt resource which loads dimensions. + :param client: + :return: + """ + # Issues a search request using streaming. + ga_service = client.get_service("GoogleAdsService") + query = "SELECT customer_client.status FROM customer_client" + stream = ga_service.search_stream(customer_id=customer_id, query=query) + for batch in stream: + for row in batch.results: + yield to_dict(row.customer_client) diff --git a/omniload/source/google_analytics/adapter.py b/omniload/source/google_analytics/adapter.py index f4edb5eb3..d228fc6cb 100644 --- a/omniload/source/google_analytics/adapter.py +++ b/omniload/source/google_analytics/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,47 +19,67 @@ from typing import Iterator, List, Optional, Union import dlt -from dlt.common import pendulum from dlt.common.typing import DictStrAny, TDataItem from dlt.sources import DltResource from dlt.sources.credentials import GcpOAuthCredentials, GcpServiceAccountCredentials from google.analytics.data_v1beta import BetaAnalyticsDataClient -from google.analytics.data_v1beta.types import ( - Dimension, - Metric, - MinuteRange, -) +from google.analytics.data_v1beta.types import GetMetadataRequest, Metadata -from .helpers import get_realtime_report, get_report +from .helpers import basic_report +from .helpers.data_processing import to_dict +from .settings import START_DATE +TIME_DIMENSIONS = { + "date", + "isoYearIsoWeek", + "year", + "yearMonth", + "dateHour", + "dateHourMinute", +} -@dlt.source(max_table_nesting=0) + +@dlt.source(max_table_nesting=2) def google_analytics( - datetime_dimension: str, credentials: Union[ GcpOAuthCredentials, GcpServiceAccountCredentials ] = dlt.secrets.value, - property_ids: List[str] = dlt.config.value, + property_id: int = dlt.config.value, queries: List[DictStrAny] = dlt.config.value, - start_date: Optional[pendulum.DateTime] = pendulum.datetime(2024, 1, 1), - end_date: Optional[pendulum.DateTime] = None, - rows_per_page: int = 10000, - minute_range_objects: List[MinuteRange] | None = None, + start_date: Optional[str] = START_DATE, + rows_per_page: int = 1000, ) -> List[DltResource]: - validated_property_ids = [] - for pid in property_ids: - try: - int_pid = int(pid) - except ValueError: - raise ValueError( - f"{pid} is an invalid google property id. Please use a numeric id, and not your Measurement ID like G-7F1AE12JLR" - ) - if int_pid == 0: - raise ValueError( - "Google Analytics property id is 0. Did you forget to configure it?" - ) - validated_property_ids.append(int_pid) + """ + The DLT source for Google Analytics. Loads basic Analytics info to the pipeline. + Args: + credentials: Credentials to the Google Analytics Account. + property_id: A numeric Google Analytics property id. + More info: https://developers.google.com/analytics/devguides/reporting/data/v1/property-id. + queries: List containing info on all the reports being requested with all the dimensions and metrics per report. + Each query can specify its own time granularity by including the appropriate time dimension + (e.g. "date", "isoYearIsoWeek", "yearMonth", etc.). If no time dimension is specified, + "date" will be used for incremental loading. + start_date: The string version of the date in the format yyyy-mm-dd and some other values. + More info: https://developers.google.com/analytics/devguides/reporting/data/v1/rest/v1beta/DateRange. + Can be left empty for default incremental load behavior. + rows_per_page: Controls how many rows are retrieved per page in the reports. + Default is 10000, maximum possible is 100000. + + Returns: + resource_list: List containing all the resources in the Google Analytics Pipeline. + """ + # validate input params for the most common mistakes + try: + property_id = int(property_id) + except ValueError: + raise ValueError( + f"{property_id} is an invalid google property id. Please use a numeric id, and not your Measurement ID like G-7F1AE12JLR" + ) + if property_id == 0: + raise ValueError( + "Google Analytics property id is 0. Did you forget to configure it?" + ) if not rows_per_page: raise ValueError("Rows per page cannot be 0") # generate access token for credentials if we are using OAuth2.0 @@ -68,77 +88,81 @@ def google_analytics( # Build the service object for Google Analytics api. client = BetaAnalyticsDataClient(credentials=credentials.to_native_credentials()) - if len(queries) > 1: - raise ValueError( - "Google Analytics supports a single query ingestion at a time, please give only one query" + # get metadata needed for some resources + metadata = get_metadata(client=client, property_id=property_id) + resource_list = [metadata | metrics_table, metadata | dimensions_table] + for query in queries: + # always add "date" to dimensions so we are able to track the last day of a report + dimensions = query["dimensions"] + time_dimension = next( + (dim for dim in dimensions if dim in TIME_DIMENSIONS), + "date", # Default to "date" if no time dimension found ) - query = queries[0] - - # always add "date" to dimensions so we are able to track the last day of a report - dimensions = query["dimensions"] - - @dlt.resource( - name="custom", - merge_key=datetime_dimension, - write_disposition="merge", - ) - def basic_report( - incremental=dlt.sources.incremental( - datetime_dimension, - initial_value=start_date, - end_value=end_date, - range_end="closed", - range_start="closed", - ), - ) -> Iterator[TDataItem]: - start_date = incremental.last_value - end_date = incremental.end_value - if start_date is None: - start_date = pendulum.datetime(2024, 1, 1) - if end_date is None: - end_date = pendulum.yesterday() - for property_id in validated_property_ids: - yield from get_report( + if time_dimension not in dimensions: + dimensions += [time_dimension] + resource_name = query["resource_name"] + resource_list.append( + dlt.resource(basic_report, name=resource_name, write_disposition="append")( client=client, + rows_per_page=rows_per_page, property_id=property_id, - dimension_list=[Dimension(name=dimension) for dimension in dimensions], - metric_list=[Metric(name=metric) for metric in query["metrics"]], - per_page=rows_per_page, + dimensions=dimensions, + metrics=query["metrics"], + resource_name=resource_name, start_date=start_date, - end_date=end_date, + last_date=dlt.sources.incremental( + time_dimension, primary_key=() + ), # pass empty primary key to avoid unique checks, a primary key defined by the resource will be used ) + ) + return resource_list - # real time report - @dlt.resource( - name="realtime", - merge_key="ingested_at", - write_disposition="merge", - ) - def real_time_report() -> Iterator[TDataItem]: - for property_id in validated_property_ids: - yield from get_realtime_report( - client=client, - property_id=property_id, - dimension_list=[Dimension(name=dimension) for dimension in dimensions], - metric_list=[Metric(name=metric) for metric in query["metrics"]], - per_page=rows_per_page, - minute_range_objects=minute_range_objects, - ) - # res = dlt.resource( - # basic_report, name="basic_report", merge_key=datetime_dimension, write_disposition="merge" - # )( - # client=client, - # rows_per_page=rows_per_page, - # property_id=property_id, - # dimensions=dimensions, - # metrics=query["metrics"], - # resource_name=resource_name, - # last_date=dlt.sources.incremental( - # datetime_dimension, - # initial_value=start_date, - # end_value=end_date, - # ), - # ) - - return [basic_report, real_time_report] +@dlt.resource(selected=False) +def get_metadata( + client: BetaAnalyticsDataClient, property_id: int +) -> Iterator[Metadata]: + """ + Get all the metrics and dimensions for a report. + + Args: + client: The Google Analytics client used to make requests. + property_id: A reference to the Google Analytics project. + More info: https://developers.google.com/analytics/devguides/reporting/data/v1/property-id + + Yields: + Metadata objects. Only 1 is expected but yield is used as dlt resources require yield to be used. + """ + request = GetMetadataRequest(name=f"properties/{property_id}/metadata") + metadata: Metadata = client.get_metadata(request) + yield metadata + + +@dlt.transformer(data_from=get_metadata, write_disposition="replace", name="metrics") +def metrics_table(metadata: Metadata) -> Iterator[TDataItem]: + """ + Loads data for metrics. + + Args: + metadata: Metadata class object which contains all the information stored in the GA4 metadata. + + Yields: + Generator of dicts, 1 metric at a time. + """ + for metric in metadata.metrics: + yield to_dict(metric) + + +@dlt.transformer(data_from=get_metadata, write_disposition="replace", name="dimensions") +def dimensions_table(metadata: Metadata) -> Iterator[TDataItem]: + """ + Loads data for dimensions. + + Args: + metadata: Metadata class object which contains all the information stored in the GA4 metadata. + + Yields: + Generator of dicts, 1 dimension at a time. + """ + for dimension in metadata.dimensions: + yield to_dict(dimension) diff --git a/omniload/source/google_analytics/settings.py b/omniload/source/google_analytics/settings.py new file mode 100644 index 000000000..d6f8dc68c --- /dev/null +++ b/omniload/source/google_analytics/settings.py @@ -0,0 +1,17 @@ +# Copyright 2022-2026 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. + +"""Google analytics source settings and constants""" + +START_DATE = "2015-08-14" diff --git a/omniload/source/google_drive/adapter.py b/omniload/source/google_drive/adapter.py new file mode 100644 index 000000000..6a14a241c --- /dev/null +++ b/omniload/source/google_drive/adapter.py @@ -0,0 +1,130 @@ +# Copyright 2022-2026 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 resource downloads files and collects filepaths from Google Drive folder to destinations""" + +from copy import deepcopy +from pathlib import Path +from typing import Any, Sequence, Union + +import dlt +from dlt.common import logger +from dlt.sources import TDataItem, TDataItems +from dlt.sources.credentials import GcpOAuthCredentials, GcpServiceAccountCredentials +from googleapiclient.discovery import build # type: ignore + +from .helpers import download_file_from_google_drive +from .settings import FOLDER_IDS, STORAGE_FOLDER_PATH + +SCOPE = "https://www.googleapis.com/auth/drive.readonly" + + +@dlt.source(name="google_drive") +def google_drive_source( + credentials: Union[ + GcpOAuthCredentials, GcpServiceAccountCredentials + ] = dlt.secrets.value, + folder_ids: Sequence[str] = FOLDER_IDS, + storage_folder_path: str = STORAGE_FOLDER_PATH, + download: bool = False, + filter_by_mime_type: Sequence[str] = (), +) -> TDataItem: + """ + Retrieves files from specified Google Drive folders. + + Args: + credentials (Dict[str, Any]): The authorized user info in Google format for authenticating with Google Drive API. + Defaults to dlt.secrets.value. + folder_ids (Sequence[str]): A sequence of folder IDs from which to retrieve the files. + Defaults to FOLDER_IDS (see settings.py). + storage_folder_path (str): The path to the local folder where the downloaded files will be stored. + Defaults to STORAGE_FOLDER_PATH (see settings.py). + download (bool): Indicates whether to download the files or not. If True, the files will be downloaded and stored locally. + If False, only the file names and IDs will be yielded without downloading. + Defaults to False. + filter_by_mime_type (Sequence[str], optional): A sequence of MIME types used to filter attachments based on their content type. Default is an empty sequence. + + Yields: + TDataItem: A dictionary representing a file. + + """ + # check if credentials are provided + if isinstance(credentials, GcpOAuthCredentials): + credentials.auth(SCOPE) + + service = build("drive", "v3", credentials=credentials.to_native_credentials()) + + if download: + storage_folder_path_ = Path(storage_folder_path) + storage_folder_path_.mkdir(exist_ok=True, parents=True) + + yield get_files_uris(service, folder_ids) | download_files( + service, storage_folder_path_, filter_by_mime_type + ) + else: + yield get_files_uris(service, folder_ids) + + +@dlt.resource(name="files_ids") +def get_files_uris(service: Any, folder_ids: Sequence[str]) -> TDataItems: + # Query the Google Drive API to get the files with the specified folder ID and extension + for folder_id in folder_ids: + results = ( + service.files() + .list( + q=f"'{folder_id}' in parents", + fields="nextPageToken, files", + ) + .execute() + ) + + items = results.get("files", []) + + if not items: + logger.warning(f"No files found in directory {folder_id}!") + else: + yield [convert_response_to_standard(item) for item in items] + + +def convert_response_to_standard(response: TDataItem) -> TDataItem: + return { + "content_type": response["mimeType"], + "parents": response["parents"], + "file_id": response["id"], + "file_name": response["name"], + "modification_date": response["modifiedTime"], + "date": response["createdTime"], + } + + +@dlt.transformer(name="attachments") +def download_files( + items: TDataItems, + service: Any, + storage_folder_path: Path, + filter_by_mime_type: Sequence[str] = (), +) -> TDataItem: + for item in items: + if filter_by_mime_type and item["content_type"] not in filter_by_mime_type: + continue + + file_name, file_id = item["file_name"], item["file_id"] + result = deepcopy(item) + + file_path = storage_folder_path / file_name + download_file_from_google_drive(service, file_id, file_path.as_posix()) + if file_path.is_file(): + result["file_path"] = file_path.absolute().as_posix() + + yield result diff --git a/omniload/source/google_drive/helpers.py b/omniload/source/google_drive/helpers.py new file mode 100644 index 000000000..12527c2bf --- /dev/null +++ b/omniload/source/google_drive/helpers.py @@ -0,0 +1,40 @@ +# Copyright 2022-2026 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. + +import io +from typing import Any + +from googleapiclient.errors import HttpError # type: ignore +from googleapiclient.http import MediaIoBaseDownload # type: ignore + + +def download_file_from_google_drive(service: Any, file_id: str, file_path: str) -> None: + try: + # Create a request to download the file + request = service.files().get_media(fileId=file_id) + # Create a BytesIO object to store the file content + fh = io.BytesIO() + # Create a downloader object to handle the download + downloader = MediaIoBaseDownload(fh, request) + # Flag to track if the download is complete + done = False + # Download the next chunk of data and check if the download is complete + while not done: + _, done = downloader.next_chunk() + # Write the downloaded file content to the specified file path + with open(file_path, "wb") as f: + f.write(fh.getvalue()) + + except HttpError as error: + print(f"An error occurred: {error}") diff --git a/omniload/source/google_drive/settings.py b/omniload/source/google_drive/settings.py new file mode 100644 index 000000000..87c79e6f7 --- /dev/null +++ b/omniload/source/google_drive/settings.py @@ -0,0 +1,16 @@ +# Copyright 2022-2026 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. + +STORAGE_FOLDER_PATH = "google_drive/attachments" +FOLDER_IDS = ("google-drive-folder",) diff --git a/omniload/source/google_sheets/adapter.py b/omniload/source/google_sheets/adapter.py index 20b61182d..bbcf64dd2 100644 --- a/omniload/source/google_sheets/adapter.py +++ b/omniload/source/google_sheets/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -84,7 +84,7 @@ def google_spreadsheet( spreadsheet_id=spreadsheet_id, range_names=list(all_range_names), ) - assert len(all_range_names) == len(all_range_data), ( # noqa: S101 + assert len(all_range_names) == len(all_range_data), ( "Google Sheets API must return values for all requested ranges" ) @@ -97,17 +97,17 @@ def google_spreadsheet( range_data = [] metadata_table = [] for name, parsed_range, meta_range, values in all_range_data: - # # pass all ranges to spreadsheet info - including empty - # metadata_table.append( - # { - # "spreadsheet_id": spreadsheet_id, - # "title": spreadsheet_title, - # "range_name": name, - # "range": str(parsed_range), - # "range_parsed": parsed_range._asdict(), - # "skipped": True, - # } - # ) + # pass all ranges to spreadsheet info - including empty + metadata_table.append( + { + "spreadsheet_id": spreadsheet_id, + "title": spreadsheet_title, + "range_name": name, + "range": str(parsed_range), + "range_parsed": parsed_range._asdict(), + "skipped": True, + } + ) if values is None or len(values) == 0: logger.warning(f"Range {name} does not contain any data. Skipping.") continue @@ -119,7 +119,7 @@ def google_spreadsheet( f"First row of range {name} does not contain data. Skipping." ) continue - # metadata_table[-1]["skipped"] = False + metadata_table[-1]["skipped"] = False range_data.append((name, parsed_range, meta_range, values)) meta_values = api_calls.get_meta_for_ranges( diff --git a/omniload/source/hubspot/adapter.py b/omniload/source/hubspot/adapter.py index 1b66641f9..97f24d45f 100644 --- a/omniload/source/hubspot/adapter.py +++ b/omniload/source/hubspot/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,7 +13,8 @@ # limitations under the License. """ -This is a module that provides a DLT source to retrieve data from multiple endpoints of the HubSpot API using a specified API key. The retrieved data is returned as a tuple of Dlt resources, one for each endpoint. +This is a module that provides a dlt source to retrieve data from multiple endpoints of the HubSpot API +using a specified API key. The retrieved data is returned as a tuple of Dlt resources, one for each endpoint. The source retrieves data from the following endpoints: - CRM Companies @@ -22,6 +23,8 @@ - CRM Tickets - CRM Products - CRM Quotes +- CRM Owners +- CRM Pipelines - Web Analytics Events For each endpoint, a resource and transformer function are defined to retrieve data and transform it to a common format. @@ -32,112 +35,252 @@ Example: To retrieve data from all endpoints, use the following code: - -python - ->>> resources = hubspot(api_key="your_api_key") """ -from typing import Any, Dict, Iterator, List, Literal, Optional, Sequence, cast -from urllib.parse import parse_qs, quote, urlparse +from typing import ( + Any, + Dict, + Iterator, + List, + Literal, + Optional, + Sequence, +) +from urllib.parse import quote import dlt from dlt.common import pendulum from dlt.common.typing import TDataItems from dlt.sources import DltResource -from omniload.error import MissingValueError - from .helpers import ( - _get_property_names, + _get_property_names_types, + _to_dlt_columns_schema, fetch_data, - fetch_data_raw, - fetch_data_search, fetch_property_history, + get_properties_labels, ) from .settings import ( - ALL, + ALL_OBJECTS, + ARCHIVED_PARAM, CRM_OBJECT_ENDPOINTS, - CRM_OWNERS_ENDPOINT, - CRM_SCHEMAS_ENDPOINT, - DEFAULT_CALL_PROPS, - DEFAULT_CART_PROPS, - DEFAULT_COMMERCE_PAYMENT_PROPS, - DEFAULT_COMPANY_PROPS, - DEFAULT_CONTACT_PROPS, - DEFAULT_DEAL_PROPS, - DEFAULT_DISCOUNT_PROPS, - DEFAULT_EMAIL_PROPS, - DEFAULT_FEE_PROPS, - DEFAULT_FEEDBACK_SUBMISSION_PROPS, - DEFAULT_INVOICE_PROPS, - DEFAULT_LINE_ITEM_PROPS, - DEFAULT_MEETING_PROPS, - DEFAULT_NOTE_PROPS, - DEFAULT_PRODUCT_PROPS, - DEFAULT_QUOTE_PROPS, - DEFAULT_TASK_PROPS, - DEFAULT_TAX_PROPS, - DEFAULT_TICKET_PROPS, + CRM_PIPELINES_ENDPOINT, + ENTITY_PROPERTIES, + HS_TO_DLT_TYPE, + MAX_PROPS_LENGTH, OBJECT_TYPE_PLURAL, + OBJECT_TYPE_SINGULAR, + PIPELINES_OBJECTS, + PROPERTIES_WITH_CUSTOM_LABELS, + SOFT_DELETE_KEY, + STAGE_PROPERTY_PREFIX, STARTDATE, WEB_ANALYTICS_EVENTS_ENDPOINT, ) +from .utils import chunk_properties + +THubspotObjectType = Literal["company", "contact", "deal", "ticket", "product", "quote"] + + +def fetch_data_for_properties( + props: List[str], + api_key: str, + object_type: str, + soft_delete: bool, +) -> Iterator[TDataItems]: + """ + Fetch data for a given set of properties from the HubSpot API. + + Args: + props (List[str]): List of property names to fetch. + api_key (str): HubSpot API key for authentication. + object_type (str): The type of HubSpot object (e.g., 'company', 'contact'). + soft_delete (bool): Flag to fetch soft-deleted (archived) records. + + Yields: + Iterator[TDataItems]: Data retrieved from the HubSpot API. + """ + # The Hubspot API expects a comma separated string as properties + joined_props = ",".join(sorted(props)) + params: Dict[str, Any] = {"properties": joined_props, "limit": 100} + context: Optional[Dict[str, Any]] = ( + {SOFT_DELETE_KEY: False} if soft_delete else None + ) + + yield from fetch_data( + CRM_OBJECT_ENDPOINTS[object_type], api_key, params=params, context=context + ) + if soft_delete: + yield from fetch_data( + CRM_OBJECT_ENDPOINTS[object_type], + api_key, + params={**params, **ARCHIVED_PARAM}, + context={SOFT_DELETE_KEY: True}, + ) + + +def crm_objects( + object_type: str, + api_key: str, + props: List[str], + include_custom_props: bool = True, + archived: bool = False, +) -> Iterator[TDataItems]: + """ + Fetch CRM object data (e.g., companies, contacts) from the HubSpot API. + + Args: + object_type (str): Type of HubSpot object (e.g., 'company', 'contact'). + api_key (str): API key for HubSpot authentication. + props (List[str]): List of properties to retrieve. + include_custom_props (bool, optional): Include custom properties in the result. Defaults to True. + archived (bool, optional): Fetch archived (soft-deleted) objects. Defaults to False. + + Yields: + Iterator[TDataItems]: Data items retrieved from the API. + """ + props_to_type = fetch_props_with_types( + object_type, api_key, props, include_custom_props + ) + # We need column hints so that dlt can correctly set data types + # This is especially relevant for columns of type "number" in Hubspot + # that are returned as strings by the API + col_type_hints = { + prop: _to_dlt_columns_schema({prop: hb_type}) + for prop, hb_type in props_to_type.items() + } + for batch in fetch_data_for_properties( + list(props_to_type.keys()), api_key, object_type, archived + ): + yield dlt.mark.with_hints(batch, dlt.mark.make_hints(columns=col_type_hints)) + + +def crm_object_history( + object_type: str, + api_key: str, + props: List[str] = None, + include_custom_props: bool = True, +) -> Iterator[TDataItems]: + """ + Fetch the history of property changes for a given CRM object type. + + Args: + object_type (str): Type of HubSpot object (e.g., 'company', 'contact'). + api_key (str): API key for HubSpot authentication. + props (List[str]): List of properties to retrieve. Defaults to None. + include_custom_props (bool): Include custom properties in the result. Defaults to True. + + Yields: + Iterator[TDataItems]: Historical property data. + """ + + # Fetch the properties from ENTITY_PROPERTIES or default to "All" + props_entry: List[str] = props or ENTITY_PROPERTIES.get(object_type, []) + + # Fetch the properties with the option to include custom properties + props_to_type = fetch_props_with_types( + object_type, api_key, props_entry, include_custom_props + ) + col_type_hints = { + prop: _to_dlt_columns_schema({prop: hb_type}) + for prop, hb_type in props_to_type.items() + if hb_type in HS_TO_DLT_TYPE + } + # We need column hints so that dlt can correctly set data types + # This is especially relevant for columns of type "number" in Hubspot + # that are returned as strings by the API + for batch in fetch_property_history( + CRM_OBJECT_ENDPOINTS[object_type], + api_key, + ",".join(sorted(props_to_type.keys())), + ): + yield dlt.mark.with_hints(batch, dlt.mark.make_hints(columns=col_type_hints)) + + +def pivot_stages_properties( + data: List[Dict[str, Any]], + property_prefix: str = STAGE_PROPERTY_PREFIX, + id_prop: str = "id", +) -> List[Dict[str, Any]]: + """ + Transform the data by pivoting stage properties. -THubspotObjectType = Literal[ - "company", - "contact", - "deal", - "ticket", - "product", - "quote", - "call", - "email", - "feedback_submission", - "line_item", - "meeting", - "note", - "task", - "cart", - "discount", - "fee", - "invoice", - "commerce_payment", - "tax", -] - - -def _last_value_to_ms(last_value) -> Optional[str]: - """Convert dlt incremental last_value (ISO string or datetime) to ms timestamp string.""" - if last_value is None: - return None - dt: pendulum.DateTime = cast( - pendulum.DateTime, - ( - pendulum.parse(last_value) - if isinstance(last_value, str) - else pendulum.instance(last_value) - ), + Args: + data (List[Dict[str, Any]]): Data containing stage properties. + property_prefix (str, optional): Prefix for stage properties. Defaults to STAGE_PROPERTY_PREFIX. + id_prop (str, optional): Name of the ID property. Defaults to "id". + + Returns: + List[Dict[str, Any]]: Transformed data with pivoted stage properties. + """ + new_data: List[Dict[str, Any]] = [] + for record in data: + record_not_null: Dict[str, Any] = { + k: v for k, v in record.items() if v is not None + } + if id_prop not in record_not_null: + continue + id_val = record_not_null.pop(id_prop) + new_data += [ + { + id_prop: id_val, + property_prefix: v, + "stage": k.split(property_prefix)[1], + } + for k, v in record_not_null.items() + if k.startswith(property_prefix) + ] + return new_data + + +def stages_timing( + object_type: str, + api_key: str, + soft_delete: bool = False, +) -> Iterator[TDataItems]: + """ + Fetch stage timing data for a specific object type from the HubSpot API. Some entities, like, + deals and tickets have pipelines with multiple stages, which they can enter and exit. This function fetches + history of entering different stages for the given object. + + Args: + object_type (str): Type of HubSpot object (e.g., 'deal', 'ticket'). + api_key (str, optional): HubSpot API key for authentication. + soft_delete (bool, optional): Fetch soft-deleted (archived) records. Defaults to False. + + Yields: + Iterator[TDataItems]: Stage timing data. + """ + all_properties: List[str] = list( + _get_property_names_types(api_key, object_type).keys() ) - return str(int(dt.timestamp() * 1000)) + date_entered_properties: List[str] = [ + prop for prop in all_properties if prop.startswith(STAGE_PROPERTY_PREFIX) + ] + # Since the length of request should be less than MAX_PROPS_LENGTH, we cannot request + # data for the whole properties list. Therefore, in the following lines we request + # data iteratively for chunks of the properties list. + for chunk in chunk_properties(date_entered_properties, MAX_PROPS_LENGTH): + for data in fetch_data_for_properties(chunk, api_key, object_type, soft_delete): + yield pivot_stages_properties(data) -@dlt.source(name="hubspot", max_table_nesting=0) + +@dlt.source(name="hubspot") def hubspot( api_key: str = dlt.secrets.value, include_history: bool = False, + soft_delete: bool = False, include_custom_props: bool = True, - custom_object: Optional[str] = None, - start_date: Optional[Any] = None, - end_date: Optional[Any] = None, -) -> Sequence[DltResource]: + properties: Optional[Dict[str, List[str]]] = None, +) -> Iterator[DltResource]: """ - A DLT source that retrieves data from the HubSpot API using the + A dlt source that retrieves data from the HubSpot API using the specified API key. This function retrieves data for several HubSpot API endpoints, including companies, contacts, deals, tickets, products and web - analytics events. It returns a tuple of Dlt resources, one for + analytics events. It returns a tuple of dlt resources, one for each endpoint. Args: @@ -147,6 +290,14 @@ def hubspot( include_history (Optional[bool]): Whether to load history of property changes along with entities. The history entries are loaded to separate tables. + soft_delete (bool): + Whether to fetch deleted properties and mark them as `is_deleted`. + include_custom_props (bool): + Whether to include custom properties. + properties (Optional(dict)): + A dictionary containing lists of properties for all the resources. + Will override the default properties ENTITY_PROPERTIES from settings + For ex., {"contact": ["createdate", "email", "firstname", "hs_object_id", "lastmodifieddate", "lastname",]} Returns: Sequence[DltResource]: Dlt resources, one for each HubSpot API endpoint. @@ -156,591 +307,192 @@ def hubspot( HubSpot CRM API. The API key is passed to `fetch_data` as the `api_key` argument. """ + properties = properties or ENTITY_PROPERTIES - if start_date is None: - start_date = "1970-01-01T00:00:00Z" - elif not isinstance(start_date, str): - start_date = pendulum.instance(start_date).to_iso8601_string() - if end_date is not None and not isinstance(end_date, str): - end_date = pendulum.instance(end_date).to_iso8601_string() - - @dlt.resource( - name="companies", write_disposition="merge", primary_key=["hs_object_id"] - ) - def companies( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot companies resource""" - yield from crm_objects( - "company", - api_key, - include_history=include_history, - props=DEFAULT_COMPANY_PROPS, - include_custom_props=include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) - - @dlt.resource( - name="contacts", write_disposition="merge", primary_key=["hs_object_id"] - ) - def contacts( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot contacts resource""" - yield from crm_objects( - "contact", - api_key, - include_history, - DEFAULT_CONTACT_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) - - @dlt.resource(name="deals", write_disposition="merge", primary_key=["hs_object_id"]) - def deals( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot deals resource""" - yield from crm_objects( - "deal", - api_key, - include_history, - DEFAULT_DEAL_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) - - @dlt.resource( - name="tickets", write_disposition="merge", primary_key=["hs_object_id"] - ) - def tickets( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot tickets resource""" - yield from crm_objects( - "ticket", - api_key, - include_history, - DEFAULT_TICKET_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) - - @dlt.resource( - name="products", write_disposition="merge", primary_key=["hs_object_id"] - ) - def products( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot products resource""" - yield from crm_objects( - "product", - api_key, - include_history, - DEFAULT_PRODUCT_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) - - @dlt.resource(name="calls", write_disposition="merge", primary_key=["hs_object_id"]) - def calls( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), + @dlt.resource(name="owners", write_disposition="merge", primary_key="id") + def owners( + api_key: str = api_key, soft_delete: bool = soft_delete ) -> Iterator[TDataItems]: - """Hubspot calls resource""" - yield from crm_objects( - "call", - api_key, - include_history, - DEFAULT_CALL_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + """Fetch HubSpot owners data. The owners resource implemented separately, + because it doesn't have endpoint for properties requesting - @dlt.resource( - name="emails", write_disposition="merge", primary_key=["hs_object_id"] - ) - def emails( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot emails resource""" - yield from crm_objects( - "email", - api_key, - include_history, - DEFAULT_EMAIL_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + Args: + api_key (str): HubSpot API key for authentication. + soft_delete (bool, optional): Fetch soft-deleted (archived) owners. Defaults to False. - @dlt.resource( - name="feedback_submissions", - write_disposition="merge", - primary_key=["hs_object_id"], - ) - def feedback_submissions( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot feedback submissions resource""" - yield from crm_objects( - "feedback_submission", - api_key, - include_history, - DEFAULT_FEEDBACK_SUBMISSION_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + Yields: + Iterator[TDataItems]: Owner data. + """ - @dlt.resource( - name="line_items", write_disposition="merge", primary_key=["hs_object_id"] - ) - def line_items( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot line items resource""" - yield from crm_objects( - "line_item", - api_key, - include_history, - DEFAULT_LINE_ITEM_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + # Fetch data for owners + for page in fetch_data(endpoint=CRM_OBJECT_ENDPOINTS["owner"], api_key=api_key): + yield page - @dlt.resource( - name="meetings", write_disposition="merge", primary_key=["hs_object_id"] - ) - def meetings( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot meetings resource""" - yield from crm_objects( - "meeting", - api_key, - include_history, - DEFAULT_MEETING_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + # Fetch soft-deleted owners if requested + if soft_delete: + for page in fetch_data( + endpoint=CRM_OBJECT_ENDPOINTS["owner"], + params=ARCHIVED_PARAM, + api_key=api_key, + context={SOFT_DELETE_KEY: True}, + ): + yield page - @dlt.resource(name="notes", write_disposition="merge", primary_key=["hs_object_id"]) - def notes( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot notes resource""" - yield from crm_objects( - "note", - api_key, - include_history, - DEFAULT_NOTE_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + @dlt.resource(name="properties", write_disposition="replace") + def properties_custom_labels(api_key: str = api_key) -> Iterator[TDataItems]: + """ + A dlt resource that retrieves custom labels for given list of properties. - @dlt.resource(name="tasks", write_disposition="merge", primary_key=["hs_object_id"]) - def tasks( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot tasks resource""" - yield from crm_objects( - "task", - api_key, - include_history, - DEFAULT_TASK_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + Args: + api_key (str, optional): HubSpot API key for authentication. - @dlt.resource(name="carts", write_disposition="merge", primary_key=["hs_object_id"]) - def carts( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot carts resource""" - yield from crm_objects( - "cart", - api_key, - include_history, - DEFAULT_CART_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + Yields: + DltResource: A dlt resource containing properties for HubSpot objects. + """ - @dlt.resource( - name="discounts", write_disposition="merge", primary_key=["hs_object_id"] - ) - def discounts( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot discounts resource""" - yield from crm_objects( - "discount", - api_key, - include_history, - DEFAULT_DISCOUNT_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + def get_properties_description( + properties_list_inner: List[Dict[str, Any]], + ) -> Iterator[Dict[str, Any]]: + """Fetch properties.""" + for property_info in properties_list_inner: + yield from get_properties_labels( + api_key=api_key, + object_type=property_info["object_type"], + property_name=property_info["property_name"], + ) + + if PROPERTIES_WITH_CUSTOM_LABELS: + yield from get_properties_description(PROPERTIES_WITH_CUSTOM_LABELS) + else: + return + + def pipelines_for_objects( + pipelines_objects: List[str], + api_key_inner: str = api_key, + ) -> Iterator[DltResource]: + """ + Function that yields all resources for HubSpot objects, which have pipelines. + (could be deals or/and tickets, specified in PIPELINES_OBJECTS) - @dlt.resource(name="fees", write_disposition="merge", primary_key=["hs_object_id"]) - def fees( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot fees resource""" - yield from crm_objects( - "fee", - api_key, - include_history, - DEFAULT_FEE_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + Args: + pipelines_objects (list of strings): The list of objects, which have pipelines. + api_key_inner (str, optional): The API key used to authenticate with the HubSpot API. Defaults to dlt.secrets.value. - @dlt.resource( - name="invoices", write_disposition="merge", primary_key=["hs_object_id"] - ) - def invoices( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot invoices resource""" - yield from crm_objects( - "invoice", - api_key, - include_history, - DEFAULT_INVOICE_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + Yields: + Iterator[DltResource]: dlt resources for pipelines and stages. + """ - @dlt.resource( - name="commerce_payments", - write_disposition="merge", - primary_key=["hs_object_id"], - ) - def commerce_payments( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot commerce payments resource""" - yield from crm_objects( - "commerce_payment", - api_key, - include_history, - DEFAULT_COMMERCE_PAYMENT_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + def get_pipelines(object_type: str) -> Iterator[TDataItems]: + yield from fetch_data( + CRM_PIPELINES_ENDPOINT.format(objectType=object_type), + api_key=api_key_inner, + ) - @dlt.resource(name="taxes", write_disposition="merge", primary_key=["hs_object_id"]) - def taxes( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot taxes resource""" - yield from crm_objects( - "tax", - api_key, - include_history, - DEFAULT_TAX_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), + # get the pipelines data + for obj_type in pipelines_objects: + name = f"pipelines_{obj_type}" + yield dlt.resource( + get_pipelines, + name=name, + write_disposition="merge", + merge_key="id", + table_name=name, + primary_key="id", + )(obj_type) + + # get the history of entering for pipeline stages + name = f"stages_timing_{obj_type}" + if obj_type in OBJECT_TYPE_SINGULAR: + yield dlt.resource( + stages_timing, + name=name, + write_disposition="merge", + primary_key=["id", "stage"], + )( + OBJECT_TYPE_SINGULAR[obj_type], + api_key=api_key, + soft_delete=soft_delete, + ) + + # resources for all objects + for obj in ALL_OBJECTS: + yield dlt.resource( + crm_objects, + name=OBJECT_TYPE_PLURAL[obj], + write_disposition="merge", + primary_key="id", + )( + object_type=obj, + api_key=api_key, + props=properties.get(obj), + include_custom_props=include_custom_props, + archived=soft_delete, ) - @dlt.resource( - name="quotes", write_disposition="merge", primary_key=["hs_object_id"] - ) - def quotes( - api_key: str = api_key, - include_history: bool = include_history, - include_custom_props: bool = include_custom_props, - updated_at: dlt.sources.incremental[str] = dlt.sources.incremental( - "hs_lastmodifieddate", - initial_value=start_date, - end_value=end_date, - ), - ) -> Iterator[TDataItems]: - """Hubspot quotes resource""" - yield from crm_objects( - "quote", - api_key, - include_history, - DEFAULT_QUOTE_PROPS, - include_custom_props, - start_date_ms=_last_value_to_ms(updated_at.last_value), - end_date_ms=_last_value_to_ms(end_date), - ) + # corresponding history resources + if include_history: + for obj in ALL_OBJECTS: + yield dlt.resource( + crm_object_history, + name=f"{OBJECT_TYPE_PLURAL[obj]}_property_history", + write_disposition="append", + )( + object_type=obj, + api_key=api_key, + props=properties.get(obj), + include_custom_props=include_custom_props, + ) - @dlt.resource(name="owners", write_disposition="merge", primary_key="id") - def owners( - api_key: str = api_key, - ) -> Iterator[TDataItems]: - """Hubspot owners resource""" - yield from fetch_data(CRM_OWNERS_ENDPOINT, api_key, resource_name="owners") + # owners resource + yield owners - @dlt.resource(name="schemas", write_disposition="merge", primary_key="id") - def schemas( - api_key: str = api_key, - ) -> Iterator[TDataItems]: - """Hubspot schemas resource""" - yield from fetch_data(CRM_SCHEMAS_ENDPOINT, api_key, resource_name="schemas") + # pipelines resources + yield from pipelines_for_objects(PIPELINES_OBJECTS, api_key) - @dlt.resource(write_disposition="merge", primary_key="hs_object_id") - def custom( - api_key: str = api_key, - custom_object_name: str = custom_object, # ty: ignore[invalid-parameter-default] - ) -> Iterator[TDataItems]: - custom_objects = fetch_data_raw(CRM_SCHEMAS_ENDPOINT, api_key) - object_type_id = None - associations = None - if ":" in custom_object_name: - fields = custom_object_name.split(":") - if len(fields) == 2: - custom_object_name = fields[0] - associations = fields[1] - - custom_object_lowercase = custom_object_name.lower() - - for co in custom_objects["results"]: - if co["name"].lower() == custom_object_lowercase: - object_type_id = co["objectTypeId"] - break - - # sometimes people use the plural name of the object type by accident, we should try to match that if we can - if "labels" in co: - if custom_object_lowercase == co["labels"]["plural"].lower(): - object_type_id = co["objectTypeId"] - break - - if object_type_id is None: - raise ValueError(f"There is no such custom object as {custom_object_name}") - custom_object_properties = f"crm/v3/properties/{object_type_id}" - - props_pages = fetch_data(custom_object_properties, api_key) - props = [] - for page in props_pages: - props.extend([prop["name"] for prop in page]) - props = ",".join(sorted(list(set(props)))) - - custom_object_endpoint = f"crm/v3/objects/{object_type_id}/?properties={props}" - if associations: - custom_object_endpoint += f"&associations={associations}" - - """Hubspot custom object details resource""" - yield from fetch_data(custom_object_endpoint, api_key, resource_name="custom") - - return ( - companies, - contacts, - deals, - tickets, - products, - quotes, - calls, - emails, - feedback_submissions, - line_items, - meetings, - notes, - tasks, - carts, - discounts, - fees, - invoices, - commerce_payments, - taxes, - owners, - schemas, - custom, - ) + # custom properties labels resource + yield properties_custom_labels -def crm_objects( +def fetch_props_with_types( object_type: str, - api_key: str = dlt.secrets.value, - include_history: bool = False, - props: Optional[Sequence[str]] = None, + api_key: str, + props: List[str], include_custom_props: bool = True, - start_date_ms: Optional[str] = None, - end_date_ms: Optional[str] = None, -) -> Iterator[TDataItems]: - """Building blocks for CRM resources.""" - props_effective: List[str] = [] - if props == ALL: - props_effective = list(_get_property_names(api_key, object_type)) - - if include_custom_props: - props_effective.extend(_get_property_names(api_key, object_type)) +) -> Dict[str, str]: + """ + Fetch the mapping of properties to their types. - props_out = ",".join(sorted(list(set(props_effective)))) + Args: + object_type (str): Type of HubSpot object (e.g., 'company', 'contact'). + api_key (str): HubSpot API key for authentication. + props (List[str]): List of properties to fetch. + include_custom_props (bool, optional): Include custom properties in the result. Defaults to True. - if start_date_ms is not None: - _qs = parse_qs(urlparse(CRM_OBJECT_ENDPOINTS[object_type]).query) - assoc_types = [t for t in _qs.get("associations", [""])[0].split(",") if t] - yield from fetch_data_search( - object_type, - api_key, - props_out, - start_date_ms, - end_date_ms=end_date_ms, - association_types=assoc_types, + Returns: + Dict[str, str]: Mapping of property to type. + """ + unique_props = set(props) + props_to_type = _get_property_names_types(api_key, object_type) + all_props = set(props_to_type.keys()) + + all_custom = {prop for prop in all_props if not prop.startswith("hs_")} + + # Choose selected props + if unique_props == all_props: + selected = all_props if include_custom_props else all_props - all_custom + else: + non_existent = unique_props - all_props + if non_existent: + raise ValueError( + f"The requested props {non_existent} don't exist in the source!" + ) + selected = ( + unique_props.union(all_custom) if include_custom_props else unique_props ) - return - params = {"properties": props_out, "limit": 100} + props_to_type = {prop: props_to_type[prop] for prop in selected} - yield from fetch_data(CRM_OBJECT_ENDPOINTS[object_type], api_key, params=params) - if include_history: - # Get history separately, as requesting both all properties and history together - # is likely to hit hubspot's URL length limit - for history_entries in fetch_property_history( - CRM_OBJECT_ENDPOINTS[object_type], - api_key, - props_out, - ): - yield dlt.mark.with_table_name( - history_entries, - OBJECT_TYPE_PLURAL[object_type] + "_property_history", - ) + return props_to_type @dlt.resource @@ -751,20 +503,20 @@ def hubspot_events_for_objects( start_date: pendulum.DateTime = STARTDATE, ) -> DltResource: """ - A standalone DLT resources that retrieves web analytics events from the HubSpot API for a particular object type and list of object ids. + A standalone dlt resource that retrieves web analytics events from the HubSpot API for a particular object type and list of object ids. Args: - object_type(THubspotObjectType, required): One of the hubspot object types see definition of THubspotObjectType literal - object_ids: (List[THubspotObjectType], required): List of object ids to track events + object_type (THubspotObjectType): One of the hubspot object types see definition of THubspotObjectType literal. + object_ids (List[str]): List of object ids to track events. api_key (str, optional): The API key used to authenticate with the HubSpot API. Defaults to dlt.secrets.value. - start_date (datetime, optional): The initial date time from which start getting events, default to STARTDATE + start_date (pendulum.DateTime, optional): The initial date time from which start getting events, default to STARTDATE. Returns: - incremental dlt resource to track events for objects from the list + DltResource: Incremental dlt resource to track events for objects from the list. """ - end_date = pendulum.now().isoformat() - name = object_type + "_events" + end_date: str = pendulum.now().isoformat() + name: str = object_type + "_events" def get_web_analytics_events( occurred_at: dlt.sources.incremental[str], @@ -773,13 +525,11 @@ def get_web_analytics_events( A helper function that retrieves web analytics events for a given object type from the HubSpot API. Args: - object_type (str): The type of object for which to retrieve web analytics events. + occurred_at (dlt.sources.incremental[str]): Incremental source for event occurrence time. Yields: - dict: A dictionary representing a web analytics event. + Iterator[List[Dict[str, Any]]]: Web analytics event data. """ - if occurred_at.last_value is None: - raise MissingValueError("occurred_at.last_value", "HubSpot") for object_id in object_ids: yield from fetch_data( WEB_ANALYTICS_EVENTS_ENDPOINT.format( @@ -798,11 +548,4 @@ def get_web_analytics_events( write_disposition="append", selected=True, table_name=lambda e: name + "_" + str(e["eventType"]), - )( - dlt.sources.incremental( - "occurredAt", - initial_value=start_date.isoformat(), - range_end="closed", - range_start="closed", - ) - ) + )(dlt.sources.incremental("occurredAt", initial_value=start_date.isoformat())) diff --git a/omniload/source/hubspot/helpers.py b/omniload/source/hubspot/helpers.py index efa873d79..00e45e854 100644 --- a/omniload/source/hubspot/helpers.py +++ b/omniload/source/hubspot/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,15 +15,12 @@ """Hubspot source helpers""" import urllib.parse -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional, Union +from dlt.common.schema.typing import TColumnSchema from dlt.sources.helpers import requests -from .settings import ( - DEFAULT_LAST_MODIFIED_PROPERTY, - LAST_MODIFIED_PROPERTY, - OBJECT_TYPE_PLURAL, -) +from .settings import HS_TO_DLT_TYPE, OBJECT_TYPE_PLURAL BASE_URL = "https://api.hubapi.com/" @@ -49,11 +46,45 @@ def _get_headers(api_key: str) -> Dict[str, str]: return dict(authorization=f"Bearer {api_key}") +def pagination( + _data: Dict[str, Any], headers: Dict[str, Any] +) -> Optional[Dict[str, Any]]: + _next = _data.get("paging", {}).get("next", None) + # _next = False + if _next: + next_url = _next["link"] + # Get the next page response + r = requests.get(next_url, headers=headers) + return r.json() # type: ignore + else: + return None + + +def extract_association_data( + _obj: Dict[str, Any], + data: Dict[str, Any], + association: str, + headers: Dict[str, Any], +) -> List[Dict[str, Any]]: + values = [] + + while data is not None: + for r in data["results"]: + values.append( + { + "value": _obj["hs_object_id"], + f"{association}_id": r["id"], + } + ) + data = pagination(data, headers) + return values + + def extract_property_history(objects: List[Dict[str, Any]]) -> Iterator[Dict[str, Any]]: for item in objects: history = item.get("propertiesWithHistory") if not history: - continue + return # Yield a flat list of property history entries for key, changes in history.items(): if not changes: @@ -111,7 +142,7 @@ def fetch_data( endpoint: str, api_key: str, params: Optional[Dict[str, Any]] = None, - resource_name: Optional[str] = None, + context: Optional[Dict[str, Any]] = None, ) -> Iterator[List[Dict[str, Any]]]: """ Fetch data from HUBSPOT endpoint using a specified API key and yield the properties of each result. @@ -120,7 +151,8 @@ def fetch_data( Args: endpoint (str): The endpoint to fetch data from, as a string. api_key (str): The API key to use for authentication, as a string. - params: Optional dict of query params to include in the request + params: Optional dict of query params to include in the request. + context (Optional[Dict[str, Any]]): Additional data which need to be added in the resulting page. Yields: A List of CRM object dicts @@ -134,7 +166,7 @@ def fetch_data( 404 Not Found), a `requests.exceptions.HTTPError` exception will be raised. The `endpoint` argument should be a relative URL, which will be appended to the base URL for the - API. The `params` argument is used to pass additional query parameters to the request + API. The `params` argument is used to pass additional query parameters to the request. This function also includes a retry decorator that will automatically retry the API call up to 3 times with a 5-second delay between retries, using an exponential backoff strategy. @@ -148,270 +180,84 @@ def fetch_data( # Parse the API response and yield the properties of each result # Parse the response JSON data _data = r.json() - # Yield the properties of each result in the API response while _data is not None: if "results" in _data: _objects: List[Dict[str, Any]] = [] for _result in _data["results"]: - if resource_name == "schemas": - _objects.append( - { - "name": _result["labels"].get("singular", ""), - "objectTypeId": _result.get("objectTypeId", ""), - "id": _result.get("id", ""), - "fullyQualifiedName": _result.get("fullyQualifiedName", ""), - "properties": _result.get("properties", ""), - "createdAt": _result.get("createdAt", ""), - "updatedAt": _result.get("updatedAt", ""), - } - ) - elif resource_name == "owners": - _objects.append( - { - "id": _result.get("id", ""), - "email": _result.get("email", ""), - "type": _result.get("type", ""), - "firstName": _result.get("firstName", ""), - "lastName": _result.get("lastName", ""), - "createdAt": _result.get("createdAt", ""), - "updatedAt": _result.get("updatedAt", ""), - "userId": _result.get("userId", ""), - "teams": _result.get("teams", []), - } - ) - else: - _obj = _result.get("properties", _result) - if "id" not in _obj and "id" in _result: - # Move id from properties to top level - _obj["id"] = _result["id"] - - if "associations" in _result: - for association in _result["associations"]: - __values = [ - { - "value": _obj["hs_object_id"], - f"{association}_id": __r["id"], - } - for __r in _result["associations"][association][ - "results" - ] - ] - - # remove duplicates from list of dicts - __values = [ - dict(t) for t in {tuple(d.items()) for d in __values} - ] - - _obj[association] = __values - - _objects.append(_obj) + _obj = _result.get("properties", _result) + if "id" not in _obj and "id" in _result: + # Move id from properties to top level + _obj["id"] = _result["id"] + if "associations" in _result: + for association in _result["associations"]: + __data = _result["associations"][association] + + __values = extract_association_data( + _obj, __data, association, headers + ) + + # remove duplicates from list of dicts + __values = [ + dict(t) for t in {tuple(d.items()) for d in __values} + ] + + _obj[association] = __values + if context: + _obj.update(context) + _objects.append(_obj) yield _objects # Follow pagination links if they exist - _next = _data.get("paging", {}).get("next", None) - if _next: - next_url = _next["link"] - # Get the next page response - r = requests.get(next_url, headers=headers) - _data = r.json() - else: - _data = None + _data = pagination(_data, headers) -def _get_property_names(api_key: str, object_type: str) -> List[str]: +def _get_property_names_types( + api_key: str, object_type: str +) -> Dict[str, Union[str, None]]: """ - Retrieve property names for a given entity from the HubSpot API. + Retrieve property names and their types if present for a given entity from the HubSpot API. Args: entity: The entity name for which to retrieve property names. Returns: - A list of property names. + A dict of propery names and their types if present. Raises: Exception: If an error occurs during the API request. """ - properties = [] + props_to_type: Dict[str, str] = {} endpoint = f"/crm/v3/properties/{OBJECT_TYPE_PLURAL[object_type]}" for page in fetch_data(endpoint, api_key): - properties.extend([prop["name"] for prop in page]) - - return properties - - -def _fetch_associations_batch( - from_object_type: str, - to_object_type: str, - object_ids: List[str], - api_key: str, -) -> Dict[str, List[str]]: - """Fetch associations for a batch of objects via the HubSpot v4 batch associations API. - - Returns a dict mapping from_id -> list of to_ids. - Returns an empty dict if the association type is unsupported. - """ - if not object_ids: - return {} - - url = get_url( - f"/crm/v4/associations/{from_object_type}/{to_object_type}/batch/read" - ) - headers = _get_headers(api_key) - r = requests.post( - url, headers=headers, json={"inputs": [{"id": oid} for oid in object_ids]} - ) - - if r.status_code in (400, 404): - return {} - r.raise_for_status() - - result: Dict[str, List[str]] = {} - for item in r.json().get("results", []): - from_id = str(item.get("from", {}).get("id", "")) - to_ids = [ - str(a["toObjectId"]) for a in item.get("to", []) if a.get("toObjectId") - ] - if from_id and to_ids: - result[from_id] = to_ids - return result - - -def fetch_data_search( - object_type: str, - api_key: str, - properties: str, - start_date_ms: str, - end_date_ms: Optional[str] = None, - association_types: Optional[List[str]] = None, -) -> Iterator[List[Dict[str, Any]]]: - import logging - - logger = logging.getLogger("hubspot.search") - - url = get_url(f"/crm/v3/objects/{OBJECT_TYPE_PLURAL[object_type]}/search") - headers = _get_headers(api_key) - from_type = OBJECT_TYPE_PLURAL[object_type] - modified_prop = LAST_MODIFIED_PROPERTY.get( - object_type, DEFAULT_LAST_MODIFIED_PROPERTY - ) - - props_list = [p for p in properties.split(",") if p] - last_id: Optional[str] = None - - while True: - filters = [ - { - "propertyName": modified_prop, - "operator": "GTE", - "value": start_date_ms, - } - ] - if end_date_ms is not None: - filters.append( - { - "propertyName": modified_prop, - "operator": "LTE", - "value": end_date_ms, - } - ) - if last_id is not None: - filters.append( - { - "propertyName": "hs_object_id", - "operator": "GT", - "value": last_id, - } - ) - - logger.info( - f"[hubspot] search {object_type}: " - f"GTE={start_date_ms} LTE={end_date_ms} after_id={last_id}" - ) + for prop in page: + props_to_type[prop["name"]] = prop.get("type", None) - body: Dict[str, Any] = { - "filterGroups": [{"filters": filters}], - "properties": props_list, - "sorts": [{"propertyName": "hs_object_id", "direction": "ASCENDING"}], - "limit": 100, - } + return props_to_type - total_yielded = 0 - while True: - r = requests.post(url, headers=headers, json=body) - r.raise_for_status() - _data = r.json() - - if _data.get("status") == "error": - raise ValueError( - f"HubSpot search error: {_data.get('message')} (correlationId: {_data.get('correlationId')})" - ) - - if "results" in _data: - _objects: List[Dict[str, Any]] = [] - for _result in _data["results"]: - _obj = _result.get("properties", _result) - if "id" not in _obj and "id" in _result: - _obj["id"] = _result["id"] - _objects.append(_obj) - - obj_id = str(_obj.get("hs_object_id") or _obj.get("id") or "") - if last_id is None or int(obj_id) > int(last_id): - last_id = obj_id - - if association_types and _objects: - obj_ids = [ - str(obj.get("hs_object_id") or obj.get("id") or "") - for obj in _objects - ] - for assoc_type in association_types: - if not assoc_type: - continue - assoc_map = _fetch_associations_batch( - from_type, assoc_type, obj_ids, api_key - ) - for obj in _objects: - obj_id = str(obj.get("hs_object_id") or obj.get("id") or "") - values = [ - {"value": obj_id, f"{assoc_type}_id": aid} - for aid in assoc_map.get(obj_id, []) - ] - obj[assoc_type] = [ - dict(t) for t in {tuple(d.items()) for d in values} - ] - - total_yielded += len(_objects) - yield _objects - - # Break BEFORE trying to fetch beyond the 10k limit — HubSpot's - # search API hangs when paging past 10,000 results. - if total_yielded >= 10000: - break - - _next = _data.get("paging", {}).get("next", None) - if _next: - body["after"] = _next["after"] - else: - break - - logger.info( - f"[hubspot] search {object_type}: window done, " - f"yielded={total_yielded} last_id={last_id}" - ) - - # HubSpot search API has a 10,000 result hard limit. If we hit it, - # restart with the same date filters plus hs_object_id > last_id - # to continue from where we left off. - if total_yielded < 10000: - break - - -def fetch_data_raw( - endpoint: str, api_key: str, params: Optional[Dict[str, Any]] = None -) -> Dict[str, Any]: +def get_properties_labels( + api_key: str, object_type: str, property_name: str +) -> Iterator[Dict[str, Any]]: + endpoint = f"/crm/v3/properties/{object_type}/{property_name}" url = get_url(endpoint) headers = _get_headers(api_key) - r = requests.get(url, headers=headers, params=params) - return r.json() + r = requests.get(url, headers=headers) + _data: Optional[Dict[str, Any]] = r.json() + while _data is not None: + yield _data + _data = pagination(_data, headers) + + +def _to_dlt_columns_schema(col: Dict[str, str]) -> TColumnSchema: + """Converts hubspot column to dlt column schema that will be + used as a column hint.""" + col_name, col_type = next(iter(col.items())) + # NOTE: if col_type is not in HS_TO_DLT_TYPE, we return an empty dict. + # Downstream, this means no column hints are provided for this property. + return ( + {"name": col_name, "data_type": HS_TO_DLT_TYPE[col_type]} + if col_type in HS_TO_DLT_TYPE + else {} + ) diff --git a/omniload/source/hubspot/settings.py b/omniload/source/hubspot/settings.py index 5c8a0f584..b5abc46a4 100644 --- a/omniload/source/hubspot/settings.py +++ b/omniload/source/hubspot/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,57 +14,26 @@ """Hubspot source settings and constants""" +from typing import Dict + from dlt.common import pendulum +from dlt.common.data_types import TDataType -STARTDATE = pendulum.datetime(year=2000, month=1, day=1) +STARTDATE = pendulum.datetime(year=2024, month=2, day=10) CRM_CONTACTS_ENDPOINT = ( - "/crm/v3/objects/contacts?associations=companies,deals,products,tickets,quotes" -) -CRM_COMPANIES_ENDPOINT = "/crm/v3/objects/companies?associations=products" -CRM_DEALS_ENDPOINT = ( - "/crm/v3/objects/deals?associations=companies,contacts,products,tickets,quotes" -) -CRM_PRODUCTS_ENDPOINT = ( - "/crm/v3/objects/products?associations=companies,contacts,deals,tickets,quotes" -) -CRM_TICKETS_ENDPOINT = ( - "/crm/v3/objects/tickets?associations=companies,contacts,deals,products,quotes" -) -CRM_QUOTES_ENDPOINT = ( - "/crm/v3/objects/quotes?associations=companies,contacts,deals,products,tickets" -) -CRM_CALLS_ENDPOINT = ( - "/crm/v3/objects/calls?associations=contacts,companies,deals,products,quotes" -) -CRM_EMAILS_ENDPOINT = ( - "/crm/v3/objects/emails?associations=contacts,companies,deals,products,quotes" -) -CRM_FEEDBACK_SUBMISSIONS_ENDPOINT = "/crm/v3/objects/feedback_submissions?associations=contacts,companies,deals,products,quotes" -CRM_LINE_ITEMS_ENDPOINT = ( - "/crm/v3/objects/line_items?associations=contacts,companies,deals,products,quotes" -) -CRM_MEETINGS_ENDPOINT = ( - "/crm/v3/objects/meetings?associations=contacts,companies,deals,products,quotes" -) -CRM_NOTES_ENDPOINT = ( - "/crm/v3/objects/notes?associations=contacts,companies,deals,products,quotes" -) -CRM_TASKS_ENDPOINT = ( - "/crm/v3/objects/tasks?associations=contacts,companies,deals,products,quotes" + "/crm/v3/objects/contacts?associations=deals,products,tickets,quotes" ) -CRM_CARTS_ENDPOINT = ( - "/crm/v3/objects/carts?associations=contacts,companies,deals,products,quotes" +CRM_COMPANIES_ENDPOINT = ( + "/crm/v3/objects/companies?associations=contacts,deals,products,tickets,quotes" ) -CRM_DISCOUNTS_ENDPOINT = "/crm/v3/objects/discounts?associations=contacts,line_items,companies,deals,products,quotes" -CRM_FEES_ENDPOINT = "/crm/v3/objects/fees?associations=contacts,line_items,companies,deals,products,quotes" -CRM_INVOICES_ENDPOINT = "/crm/v3/objects/invoices?associations=contacts,line_items,companies,fees,products,quotes" -CRM_COMMERCE_PAYMENTS_ENDPOINT = "/crm/v3/objects/commerce_payments?associations=contacts,companies,deals,quotes,invoices,products,fees" -CRM_TAXES_ENDPOINT = ( - "/crm/v3/objects/taxes?associations=line_items,companies,deals,products,quotes,fees" -) -CRM_OWNERS_ENDPOINT = "/crm/v3/owners" -CRM_SCHEMAS_ENDPOINT = "/crm/v3/schemas" +CRM_DEALS_ENDPOINT = "/crm/v3/objects/deals" +CRM_PRODUCTS_ENDPOINT = "/crm/v3/objects/products" +CRM_TICKETS_ENDPOINT = "/crm/v3/objects/tickets" +CRM_QUOTES_ENDPOINT = "/crm/v3/objects/quotes" +CRM_OWNERS_ENDPOINT = "/crm/v3/owners/" +CRM_PROPERTIES_ENDPOINT = "/crm/v3/properties/{objectType}/{property_name}" +CRM_PIPELINES_ENDPOINT = "/crm/v3/pipelines/{objectType}" CRM_OBJECT_ENDPOINTS = { "contact": CRM_CONTACTS_ENDPOINT, @@ -73,19 +42,7 @@ "product": CRM_PRODUCTS_ENDPOINT, "ticket": CRM_TICKETS_ENDPOINT, "quote": CRM_QUOTES_ENDPOINT, - "call": CRM_CALLS_ENDPOINT, - "email": CRM_EMAILS_ENDPOINT, - "feedback_submission": CRM_FEEDBACK_SUBMISSIONS_ENDPOINT, - "line_item": CRM_LINE_ITEMS_ENDPOINT, - "meeting": CRM_MEETINGS_ENDPOINT, - "note": CRM_NOTES_ENDPOINT, - "task": CRM_TASKS_ENDPOINT, - "cart": CRM_CARTS_ENDPOINT, - "discount": CRM_DISCOUNTS_ENDPOINT, - "fee": CRM_FEES_ENDPOINT, - "invoice": CRM_INVOICES_ENDPOINT, - "commerce_payment": CRM_COMMERCE_PAYMENTS_ENDPOINT, - "tax": CRM_TAXES_ENDPOINT, + "owner": CRM_OWNERS_ENDPOINT, } WEB_ANALYTICS_EVENTS_ENDPOINT = "/events/v3/events?objectType={objectType}&objectId={objectId}&occurredAfter={occurredAfter}&occurredBefore={occurredBefore}&sort=-occurredAt" @@ -97,39 +54,11 @@ "tickets": "ticket", "products": "product", "quotes": "quote", - "calls": "call", - "emails": "email", - "feedback_submissions": "feedback_submission", - "line_items": "line_item", - "meetings": "meeting", - "notes": "note", - "tasks": "task", - "carts": "cart", - "discounts": "discount", - "fees": "fee", - "invoices": "invoice", - "commerce_payments": "commerce_payment", - "taxes": "tax", } OBJECT_TYPE_PLURAL = {v: k for k, v in OBJECT_TYPE_SINGULAR.items()} +ALL_OBJECTS = OBJECT_TYPE_PLURAL.keys() -# Contacts use "lastmodifieddate"; all other CRM objects use "hs_lastmodifieddate" -LAST_MODIFIED_PROPERTY = { - "contact": "lastmodifieddate", -} -DEFAULT_LAST_MODIFIED_PROPERTY = "hs_lastmodifieddate" - -DEFAULT_DEAL_PROPS = [ - "amount", - "closedate", - "createdate", - "dealname", - "dealstage", - "hs_lastmodifieddate", - "hs_object_id", - "pipeline", -] DEFAULT_COMPANY_PROPS = [ "createdate", @@ -148,6 +77,17 @@ "lastname", ] +DEFAULT_DEAL_PROPS = [ + "amount", + "closedate", + "createdate", + "dealname", + "dealstage", + "hs_lastmodifieddate", + "hs_object_id", + "pipeline", +] + DEFAULT_TICKET_PROPS = [ "createdate", "content", @@ -179,151 +119,28 @@ "hs_title", ] -DEFAULT_CALL_PROPS = [ - "hs_call_body", - "hs_call_direction", - "hs_call_disposition", - "hs_call_duration", - "hs_call_from_number", - "hs_call_status", - "hs_call_title", - "hs_call_to_number", - "hs_lastmodifieddate", - "hs_timestamp", -] - -DEFAULT_EMAIL_PROPS = [ - "hs_attachment_ids", - "hs_email_direction", - "hs_email_headers", - "hs_email_html", - "hs_email_status", - "hs_email_subject", - "hs_email_text", - "hs_timestamp", - "hs_lastmodifieddate", - "hubspot_owner_id", -] - -DEFAULT_FEEDBACK_SUBMISSION_PROPS = [ - "hs_createdate", - "hs_lastmodifieddate", - "hs_object_id", - "hs_sentiment", - "hs_survey_channel", -] - -DEFAULT_LINE_ITEM_PROPS = [ - "amount", - "description", - "hs_line_item_currency_code", - "hs_recurring_billing_end_date", - "hs_recurring_billing_start_date", - "hs_lastmodifieddate", - "hs_sku", - "name", - "price", - "quantity", - "recurringbillingfrequency", -] - -DEFAULT_MEETING_PROPS = [ - "hs_internal_meeting_notes", - "hs_meeting_body", - "hs_meeting_end_time", - "hs_meeting_external_url", - "hs_meeting_location", - "hs_meeting_outcome", - "hs_meeting_start_time", - "hs_meeting_title", - "hs_timestamp", - "hs_lastmodifieddate", - "hubspot_owner_id", -] - -DEFAULT_NOTE_PROPS = [ - "hs_attachment_ids", - "hs_note_body", - "hs_timestamp", - "hs_lastmodifieddate", - "hubspot_owner_id", -] - -DEFAULT_TASK_PROPS = [ - "hs_task_body", - "hs_task_priority", - "hs_task_status", - "hs_task_subject", - "hs_task_type", - "hs_timestamp", - "hs_lastmodifieddate", - "hubspot_owner_id", -] - -DEFAULT_CART_PROPS = [ - "hs_cart_discount", - "hs_cart_name", - "hs_cart_url", - "hs_createdate", - "hs_currency_code", - "hs_external_cart_id", - "hs_external_status", - "hs_lastmodifieddate", - "hs_object_id", - "hs_shipping_cost", - "hs_source_store", - "hs_tags", - "hs_tax", - "hs_total_price", -] - -DEFAULT_DISCOUNT_PROPS = [ - "hs_duration", - "hs_label", - "hs_lastmodifieddate", - "hs_sort_order", - "hs_type", - "hs_value", -] - -DEFAULT_FEE_PROPS = [ - "hs_label", - "hs_lastmodifieddate", - "hs_type", - "hs_value", -] - -DEFAULT_INVOICE_PROPS = [ - "hs_currency", - "hs_due_date", - "hs_invoice_date", - "hs_lastmodifieddate", - "hs_tax_id", -] - -DEFAULT_COMMERCE_PAYMENT_PROPS = [ - "hs_currency_code", - "hs_customer_email", - "hs_fees_amount", - "hs_initial_amount", - "hs_initiated_date", - "hs_internal_comment", - "hs_lastmodifieddate", - "hs_latest_status", - "hs_payment_method_type", - "hs_payout_date", - "hs_processor_type", - "hs_reference_number", - "hs_refunds_amount", - "hs_billing_address_city", - "hs_billing_address_country", -] +ENTITY_PROPERTIES = { + "company": DEFAULT_COMPANY_PROPS, + "contact": DEFAULT_CONTACT_PROPS, + "deal": DEFAULT_DEAL_PROPS, + "ticket": DEFAULT_TICKET_PROPS, + "product": DEFAULT_PRODUCT_PROPS, + "quote": DEFAULT_QUOTE_PROPS, +} -DEFAULT_TAX_PROPS = [ - "hs_label", - "hs_lastmodifieddate", - "hs_type", - "hs_value", -] -ALL = ("ALL",) +PIPELINES_OBJECTS = ["deals", "tickets"] +SOFT_DELETE_KEY = "is_deleted" +ARCHIVED_PARAM = {"archived": True} +PREPROCESSING = {"split": ["hs_merged_object_ids"]} +STAGE_PROPERTY_PREFIX = "hs_v2_date_entered_" +MAX_PROPS_LENGTH = 2000 +PROPERTIES_WITH_CUSTOM_LABELS = () + +HS_TO_DLT_TYPE: Dict[str, TDataType] = { + "bool": "bool", + "enumeration": "text", + "number": "double", + "datetime": "timestamp", + "string": "text", +} diff --git a/omniload/source/hubspot/utils.py b/omniload/source/hubspot/utils.py new file mode 100644 index 000000000..71a15c04c --- /dev/null +++ b/omniload/source/hubspot/utils.py @@ -0,0 +1,43 @@ +# Copyright 2022-2026 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, List + +from .settings import PREPROCESSING + + +def split_data(doc: Dict[str, Any]) -> Dict[str, Any]: + for prop in PREPROCESSING["split"]: + if prop in doc and doc[prop] is not None: + if isinstance(doc[prop], str): + doc[prop] = doc[prop].split(";") + return doc + + +def chunk_properties(properties: List[str], max_length: int) -> Iterator[List[str]]: + """Function which yields chunk of properties list, making sure that + for each chunk, len(",".join(chunk)) =< max_length. + """ + chunk: List[str] = [] + length = 0 + for prop in properties: + prop_len = len(prop) + (1 if chunk else 0) # include comma length if not first + if length + prop_len > max_length: + yield chunk + chunk, length = [prop], len(prop) + else: + chunk.append(prop) + length += prop_len + if chunk: + yield chunk diff --git a/omniload/source/imap/adapter.py b/omniload/source/imap/adapter.py new file mode 100644 index 000000000..b33b43367 --- /dev/null +++ b/omniload/source/imap/adapter.py @@ -0,0 +1,194 @@ +# Copyright 2022-2026 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. + +"""Reads messages and attachments from e-mail inbox via IMAP protocol""" + +import imaplib +from copy import deepcopy +from typing import Iterable, List, Optional, Sequence + +import dlt +from dlt.common import logger, pendulum +from dlt.sources import DltResource, TDataItem, TDataItems +from dlt.sources.filesystem import FileItem, FileItemDict + +from .helpers import ( + ImapFileItem, + extract_attachments, + extract_email_info, + get_message_uids, + get_message_with_internal_date, +) +from .settings import DEFAULT_CHUNK_SIZE, DEFAULT_START_DATE, GMAIL_GROUP + + +@dlt.source +def inbox_source( + host: str = dlt.secrets.value, + email_account: str = dlt.secrets.value, + password: str = dlt.secrets.value, + folder: str = "INBOX", + gmail_group: Optional[str] = GMAIL_GROUP, + start_date: pendulum.DateTime = DEFAULT_START_DATE, + filter_emails: Sequence[str] = None, + filter_by_mime_type: Sequence[str] = None, + chunksize: int = DEFAULT_CHUNK_SIZE, +) -> Sequence[DltResource]: + """This source collects inbox emails and downloads attachments to the local folder. + + Args: + host (str, optional): The hostname of the IMAP server. Default is 'dlt.secrets.value'. + email_account (str, optional): The email account used to log in to the IMAP server. Default is 'dlt.secrets.value'. + password (str, optional): The password for the email account. Default is 'dlt.secrets.value'. + folder (str, optional): The mailbox folder from which to collect emails. Default is 'INBOX'. + gmail_group (str, optional): The email address of the Google Group to filter emails sent to the group. Default is 'GMAIL_GROUP' from settings. + start_date (pendulum.Date, optional): The start date (with a day resolution) from which to collect emails. Default is 'DEFAULT_START_DATE' from settings. + filter_emails (Sequence[str], optional): A sequence of email addresses used to filter emails based on the 'FROM' field. Default is 'FILTER_EMAILS' from settings. + filter_by_mime_type (Sequence[str], optional): A sequence of MIME types used to filter attachments based on their content type. Default is an empty sequence. + chunksize (int, optional): The number of message UIDs to collect at a time. Default is 'DEFAULT_CHUNK_SIZE' from settings. + + Returns: + Sequence[DltResource]: Returns following resources: uids, messages, attachments + """ + + def _login(client: imaplib.IMAP4_SSL) -> None: + client.login(email_account, password) + r, dat = client.select(folder, readonly=True) + if r != "OK": + raise client.error(dat[-1]) + + @dlt.resource(name="uids") + def get_messages_uids( + initial_message_num: Optional[ + dlt.sources.incremental[int] + ] = dlt.sources.incremental("message_uid", initial_value=1), + ) -> TDataItem: + """Collects email message UIDs (Unique IDs) from the mailbox. + Args: + initial_message_num (int, optional): Controls incremental loading on UID + + Yields: + TDataItem: A dictionary containing the 'message_uid' of the collected email message. + """ + + last_message_num = initial_message_num.last_value + + with imaplib.IMAP4_SSL(host) as client: + _login(client) + + criteria = [ + f"(SINCE {start_date.strftime('%d-%b-%Y')})", + f"(UID {str(int(last_message_num))}:*)", + ] + + if gmail_group: + logger.info(f"Load all emails for Group: {gmail_group}") + criteria.extend([f"(TO {gmail_group})"]) + + if filter_emails: + logger.info(f"Load emails only from: {filter_emails}") + if len(filter_emails) == 1: + criteria.append(f"(FROM {filter_emails[0]})") + else: + email_filter = " ".join( + [f"FROM {email}" for email in filter_emails] + ) + criteria.append(f"(OR {email_filter})") + + uids = get_message_uids(client, criteria) + if uids: + for i in range(0, len(uids), chunksize): + yield uids[i : i + chunksize] + + @dlt.transformer(name="messages", primary_key="message_uid") + def get_messages( + items: TDataItems, + include_body: bool = True, + ) -> TDataItem: + """Reads email messages from the mailbox based on the provided message UIDs. + + Args: + items (TDataItems): An iterable containing dictionaries with 'message_uid' representing the email message UIDs. + include_body (bool, optional): If True, includes the email body in the result. Default is True. + + Yields: + TDataItem: A dictionary containing the extracted email information from the read email message. + """ + + with imaplib.IMAP4_SSL(host) as client: + _login(client) + + for item in items: + message_uid = str(item["message_uid"]) + msg, internal_date = get_message_with_internal_date(client, message_uid) + result = deepcopy(item) + result["modification_date"] = internal_date + result.update(extract_email_info(msg, include_body=include_body)) + + yield result + + @dlt.transformer( + name="attachments", + primary_key="file_hash", + ) + def get_attachments( + items: TDataItems, + ) -> Iterable[List[FileItem]]: + """Downloads attachments from email messages based on the provided message UIDs. + + Args: + items (TDataItems): An iterable containing dictionaries with 'message_uid' representing the email message UIDs. + + Yields: + Iterable[List[FileItem]]: A dictionary containing the attachment FileItem. + """ + + with imaplib.IMAP4_SSL(host) as client: + _login(client) + + files_dict: List[FileItemDict] = [] + + for item in items: + message_uid = str(item["message_uid"]) + msg, internal_date = get_message_with_internal_date(client, message_uid) + attachments = list(extract_attachments(msg, filter_by_mime_type)) + if len(attachments) == 0: + continue + + email_info = extract_email_info(msg) + + for attachment in attachments: + attachment["modification_date"] = internal_date + attachment["file_url"] = ( + f"imap://{email_account}/{message_uid}/{attachment['file_name']}" + ) + + file_dict = FileItemDict(attachment) + file_dict["message"] = dict(email_info) + file_dict["message"].update(item) + + files_dict.append(file_dict) + if len(files_dict) >= chunksize: + yield files_dict # type: ignore + files_dict = [] + + # yield remainder + if files_dict: + yield files_dict # type: ignore + + return ( + get_messages_uids, + get_messages_uids | get_attachments, + get_messages_uids | get_messages, + ) diff --git a/omniload/source/imap/helpers.py b/omniload/source/imap/helpers.py new file mode 100644 index 000000000..56aa2563a --- /dev/null +++ b/omniload/source/imap/helpers.py @@ -0,0 +1,200 @@ +# Copyright 2022-2026 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. + +import email +import hashlib +import imaplib +from email.header import decode_header, make_header +from email.message import Message +from email.utils import parsedate_to_datetime +from time import mktime +from typing import Any, Dict, Iterator, Optional, Sequence, Tuple + +from dlt.common import logger, pendulum +from dlt.sources import TDataItems +from dlt.sources.filesystem import FileItem + + +class ImapFileItem(FileItem): + """A Imap file item.""" + + file_hash: str + + +def decode_header_word(v: Any) -> Any: + if not isinstance(v, str): + return v + try: + v = str(make_header(decode_header(v))) + except Exception: + pass + + return v + + +def get_message_uids( + client: imaplib.IMAP4_SSL, criterias: Sequence[str] +) -> Optional[TDataItems]: + """Get the message uids from the imap server. + + Args: + client (imaplib.IMAP4_SSL): The imap client. + criterias (Sequence[str]): The search criterias. + + Returns: + Optional[TDataItems]: The list of message uids. + """ + status, messages = client.uid("search", *criterias) + + if status != "OK": + raise client.error(messages[-1]) + + message_uids = messages[0].split() + + if not message_uids: + logger.warning("No emails found.") + return None + + return [{"message_uid": int(message_uid)} for message_uid in message_uids] + + +def get_internal_date(client: imaplib.IMAP4_SSL, message_uid: str) -> Optional[Any]: + """Get the internal date of the email message. + + Parameters: + client (imaplib.IMAP4_SSL): The imap client. + message_uid (str): The uid of the message. + + Returns: + Optional[Any]: The internal date of the email message. + """ + # client.select() + status, data = client.uid("fetch", message_uid, "(INTERNALDATE)") + date = None + + if status != "OK": + raise client.error(data[-1]) + + timestruct = imaplib.Internaldate2tuple(data[0]) + date = pendulum.from_timestamp(mktime(timestruct)) + return date + + +def extract_email_info(msg: Message, include_body: bool = False) -> Dict[str, Any]: + """Extract the email information from the email message. + + Parameters: + msg (Message): The email message object. + include_body (bool, optional): If true, the body of the email will be included. + + Returns: + Dict[str, Any]: The email information. + """ + email_data = dict(msg) + dt = parsedate_to_datetime(msg["Date"]) + dt_pendulum = pendulum.instance(dt) + email_data["Date"] = dt_pendulum + email_data["content_type"] = msg.get_content_type() + if include_body: + email_data["body"] = get_email_body(msg) + + return { + k: decode_header_word(v) + for k, v in email_data.items() + if not k.startswith(("X-", "ARC-", "DKIM-")) + } + + +def get_message_with_internal_date( + client: imaplib.IMAP4_SSL, message_uid: str +) -> Tuple[Message, pendulum.DateTime]: + """Get the email message and internal date from the imap server. + + Parameters: + client (imaplib.IMAP4_SSL): The imap client. + message_uid (str): The uid of the message. + + Returns: + Tuple[Message, pendulum.DateTime]: The email message object and internal date as pendulum DateTime + """ + status, data = client.uid("fetch", message_uid, "(RFC822 INTERNALDATE)") + + if status == "OK": + try: + raw_email = data[0][1] + except (IndexError, TypeError): + raise Exception(f"Error getting content of email with uid {message_uid}.") + else: + raise client.error(data[-1]) + + msg = email.message_from_bytes(raw_email) + timestruct = imaplib.Internaldate2tuple(data[1]) + return msg, pendulum.from_timestamp(mktime(timestruct)) + + +def extract_attachments( + message: Message, filter_by_mime_type: Sequence[str] = None +) -> Iterator[ImapFileItem]: + """Extract the attachments from the email message. + + Parameters: + message (Message): The email message object. + filter_by_mime_type (str): The mime type to filter the attachments. + + Returns: + Iterable[ImapFileItem]: The attachments. + """ + for part in message.walk(): + content_type = part.get_content_type() + content_disposition = part.get_content_disposition() + + # Checks if the content is an attachment + if not content_disposition or content_disposition.lower() != "attachment": + continue + + # Checks if the mime type is in the filter list + if filter_by_mime_type and content_type not in filter_by_mime_type: + continue + + file_md = ImapFileItem( # type: ignore + file_name=part.get_filename(), + mime_type=content_type, + file_content=part.get_payload(decode=True), # type: ignore[typeddict-item] + ) + file_md["file_hash"] = hashlib.sha256(file_md["file_content"]).hexdigest() + file_md["size_in_bytes"] = len(file_md["file_content"]) + + yield file_md + + +def get_email_body(msg: Message) -> str: + """ + Get the body of the email message. + + Parameters: + msg (Message): The email message object. + + Returns: + str: The email body as a string. + """ + body = "" + if msg.is_multipart(): + for part in msg.walk(): + content_type = part.get_content_type() + if content_type == "text/plain": + body += part.get_payload(decode=True).decode(errors="ignore") # type: ignore[union-attr] + else: + body = msg.get_payload(decode=True).decode(errors="ignore") # type: ignore[union-attr] + + return body diff --git a/omniload/source/imap/settings.py b/omniload/source/imap/settings.py new file mode 100644 index 000000000..5d67a9da0 --- /dev/null +++ b/omniload/source/imap/settings.py @@ -0,0 +1,19 @@ +# Copyright 2022-2026 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 dlt.common import pendulum + +GMAIL_GROUP = None +DEFAULT_START_DATE = pendulum.datetime(1970, 1, 1) +DEFAULT_CHUNK_SIZE = 100 diff --git a/omniload/source/jira/adapter.py b/omniload/source/jira/adapter.py index 4b9520a45..5ec8ac1be 100644 --- a/omniload/source/jira/adapter.py +++ b/omniload/source/jira/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,366 +12,152 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -This source provides data extraction from Jira Cloud via the REST API v3. +"""This source uses Jira API and dlt to load data such as Issues, Users, Workflows and Projects to the database.""" -It defines several functions to fetch data from different parts of Jira including -projects, issues, users, boards, sprints, and various configuration objects like -issue types, statuses, and priorities. -""" - -from typing import Any, Iterable, Optional +from typing import Iterable, List, Optional import dlt -from dlt.common.typing import TDataItem - -from .helpers import get_client -from .settings import ( - DEFAULT_PAGE_SIZE, - DEFAULT_START_DATE, - ISSUE_FIELDS, -) - - -@dlt.source -def jira_source() -> Any: - """ - The main function that runs all the other functions to fetch data from Jira. - - Returns: - Sequence[DltResource]: A sequence of DltResource objects containing the fetched data. - """ - return [ - projects, - issues, - users, - issue_types, - statuses, - priorities, - resolutions, - project_versions, - project_components, - events, - issue_changelogs, - ] - - -@dlt.resource(write_disposition="replace") -def projects( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, - expand: Optional[str] = None, - recent: Optional[int] = None, -) -> Iterable[TDataItem]: - """ - Fetches and returns a list of projects from Jira. - - Args: - base_url (str): Jira instance URL (e.g., https://your-domain.atlassian.net) - email (str): User email for authentication - api_token (str): API token for authentication - expand (str): Comma-separated list of fields to expand - recent (int): Number of recent projects to return - - Yields: - dict: The project data. - """ - client = get_client(base_url, email, api_token) - yield from client.get_projects(expand=expand, recent=recent) - - -@dlt.resource( - write_disposition="merge", - primary_key="id", - max_table_nesting=2, -) -def issues( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, - jql: str = "order by updated DESC", - fields: Optional[str] = None, - expand: Optional[str] = None, - max_results: Optional[int] = None, - updated: dlt.sources.incremental[str] = dlt.sources.incremental( - "fields.updated", - initial_value=DEFAULT_START_DATE, - range_end="closed", - range_start="closed", - ), -) -> Iterable[TDataItem]: - """ - Fetches issues from Jira using JQL search. - - Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - jql (str): JQL query string - fields (str): Comma-separated list of fields to return - expand (str): Comma-separated list of fields to expand - max_results (int): Maximum number of results to return - updated (str): The date from which to fetch updated issues - - Yields: - dict: The issue data. - """ - client = get_client(base_url, email, api_token) - - # Build JQL with incremental filter - incremental_jql = jql - if updated.start_value: - date_filter = f"updated >= '{updated.start_value}'" - - # Check if JQL has ORDER BY clause and handle it properly - jql_upper = jql.upper() - if "ORDER BY" in jql_upper: - # Split at ORDER BY and add filter before it - order_by_index = jql_upper.find("ORDER BY") - main_query = jql[:order_by_index].strip() - order_clause = jql[order_by_index:].strip() - - if main_query and ( - "WHERE" in main_query.upper() - or "AND" in main_query.upper() - or "OR" in main_query.upper() - ): - incremental_jql = f"({main_query}) AND {date_filter} {order_clause}" - else: - if main_query: - incremental_jql = f"{main_query} AND {date_filter} {order_clause}" - else: - incremental_jql = f"{date_filter} {order_clause}" - else: - # No ORDER BY clause, use original logic - if "WHERE" in jql_upper or "AND" in jql_upper or "OR" in jql_upper: - incremental_jql = f"({jql}) AND {date_filter}" - else: - incremental_jql = f"{jql} AND {date_filter}" - - # Use default fields if not specified - if fields is None: - fields = ",".join(ISSUE_FIELDS) - - yield from client.search_issues( - jql=incremental_jql, fields=fields, expand=expand, max_results=max_results - ) - - -@dlt.resource(write_disposition="replace") -def users( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, - username: Optional[str] = None, - account_id: Optional[str] = None, - max_results: int = DEFAULT_PAGE_SIZE, -) -> Iterable[TDataItem]: - """ - Fetches users from Jira. - - Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - username (str): Username to search for - account_id (str): Account ID to search for - max_results (int): Maximum results per page - - Yields: - dict: The user data. - """ - client = get_client(base_url, email, api_token) - yield from client.get_users( - username=username, account_id=account_id, max_results=max_results - ) - - -@dlt.resource(write_disposition="replace") -def issue_types( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: - """ - Fetches all issue types from Jira. - - Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - - Yields: - dict: The issue type data. - """ - client = get_client(base_url, email, api_token) - yield from client.get_issue_types() - - -@dlt.resource(write_disposition="replace") -def statuses( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: - """ - Fetches all statuses from Jira. - - Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - - Yields: - dict: The status data. - """ - client = get_client(base_url, email, api_token) - yield from client.get_statuses() - - -@dlt.resource(write_disposition="replace") -def priorities( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: - """ - Fetches all priorities from Jira. - - Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - - Yields: - dict: The priority data. - """ - client = get_client(base_url, email, api_token) - yield from client.get_priorities() - - -@dlt.resource(write_disposition="replace") -def resolutions( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: - """ - Fetches all resolutions from Jira. +from dlt.common.typing import DictStrAny, TDataItem +from dlt.sources import DltResource +from dlt.sources.helpers import requests - Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - - Yields: - dict: The resolution data. - """ - client = get_client(base_url, email, api_token) - yield from client.get_resolutions() +from .settings import DEFAULT_ENDPOINTS, DEFAULT_PAGE_SIZE -@dlt.transformer( - data_from=projects, - write_disposition="replace", -) -@dlt.defer -def project_versions( - project: TDataItem, - base_url: str = dlt.secrets.value, +@dlt.source(max_table_nesting=2) +def jira( + subdomain: str = dlt.secrets.value, email: str = dlt.secrets.value, api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: + page_size: int = DEFAULT_PAGE_SIZE, +) -> Iterable[DltResource]: """ - Fetches versions for each project from Jira. + Jira source function that generates a list of resource functions based on endpoints. Args: - project (dict): The project data. - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - + subdomain: The subdomain for the Jira instance. + email: The email to authenticate with. + api_token: The API token to authenticate with. + page_size: Maximum number of results per page Returns: - list[dict]: The version data for the given project. - """ - client = get_client(base_url, email, api_token) - project_key = project.get("key") - if not project_key: - return [] - - return list(client.get_project_versions(project_key)) - - -@dlt.transformer( - data_from=projects, - write_disposition="replace", -) -@dlt.defer -def project_components( - project: TDataItem, - base_url: str = dlt.secrets.value, + Iterable[DltResource]: List of resource functions. + """ + resources = [] + for endpoint_name, endpoint_parameters in DEFAULT_ENDPOINTS.items(): + res_function = dlt.resource( + get_paginated_data, name=endpoint_name, write_disposition="replace" + )( + **endpoint_parameters, # type: ignore[arg-type] + subdomain=subdomain, + email=email, + api_token=api_token, + page_size=page_size, + ) + resources.append(res_function) + + return resources + + +@dlt.source(max_table_nesting=2) +def jira_search( + subdomain: str = dlt.secrets.value, email: str = dlt.secrets.value, api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: + page_size: int = DEFAULT_PAGE_SIZE, +) -> Iterable[DltResource]: """ - Fetches components for each project from Jira. + Jira search source function that generates a resource function for searching issues. Args: - project (dict): The project data. - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - + subdomain: The subdomain for the Jira instance. + email: The email to authenticate with. + api_token: The API token to authenticate with. + page_size: Maximum number of results per page Returns: - list[dict]: The component data for the given project. - """ - client = get_client(base_url, email, api_token) - project_key = project.get("key") - if not project_key: - return [] - - return list(client.get_project_components(project_key)) - - -@dlt.resource(write_disposition="replace") -def events( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, + Iterable[DltResource]: Resource function for searching issues. + """ + + @dlt.resource(write_disposition="replace") + def issues(jql_queries: List[str]) -> Iterable[TDataItem]: + api_path = "rest/api/3/search/jql" + + for jql in jql_queries: + params = { + "fields": "*all", + "expand": "fields,changelog,operations,transitions,names", + "validateQuery": "strict", + "jql": jql, + } + + yield from get_paginated_data( + api_path=api_path, + params=params, + subdomain=subdomain, + email=email, + api_token=api_token, + page_size=page_size, + data_path="issues", + use_cursor_pagination=True, + ) + + return issues + + +def get_paginated_data( + subdomain: str, + email: str, + api_token: str, + page_size: int, + api_path: str = "rest/api/3/search/jql", + data_path: Optional[str] = None, + params: Optional[DictStrAny] = None, + use_cursor_pagination: bool = False, ) -> Iterable[TDataItem]: """ - Fetches all event types from Jira (e.g., Issue Created, Issue Updated, etc.). + Function to fetch paginated data from a Jira API endpoint. Args: - base_url (str): Jira instance URL - email (str): User email for authentication - api_token (str): API token for authentication - + subdomain: The subdomain for the Jira instance. + email: The email to authenticate with. + api_token: The API token to authenticate with. + page_size: Maximum number of results per page + api_path: The API path for the Jira endpoint. + data_path: Optional data path to extract from the response. + params: Optional parameters for the API request. + use_cursor_pagination: If True, uses cursor-based pagination (nextPageToken) + for the /search/jql endpoint. If False, uses offset-based pagination (startAt). Yields: - dict: The event data. + Iterable[TDataItem]: Yields pages of data from the API. """ - client = get_client(base_url, email, api_token) - yield from client.get_events() + url = f"https://{subdomain}.atlassian.net/{api_path}" + headers = {"Accept": "application/json"} + auth = (email, api_token) + params = {} if params is None else dict(params) + params["maxResults"] = page_size + if not use_cursor_pagination: + params["startAt"] = start_at = 0 -@dlt.resource(write_disposition="replace", max_table_nesting=0) -def issue_changelogs( - base_url: str = dlt.secrets.value, - email: str = dlt.secrets.value, - api_token: str = dlt.secrets.value, -) -> Iterable[TDataItem]: - client = get_client(base_url, email, api_token) + while True: + response = requests.get(url, auth=auth, headers=headers, params=params) + response.raise_for_status() + result = response.json() - issue_keys = [] - for project in client.get_projects(): - project_key = project.get("key") - if project_key: - jql = f"project = {project_key} order by updated DESC" - for issue in client.search_issues(jql=jql, fields="key"): - issue_key = issue.get("key") - if issue_key: - issue_keys.append(issue_key) + if data_path: + results_page = result.pop(data_path) + else: + results_page = result - if issue_keys: - yield from client.get_changelogs_bulk(issue_keys) + if len(results_page) == 0: + break + + yield results_page + + if use_cursor_pagination: + next_page_token = result.get("nextPageToken") + if not next_page_token: + break + params["nextPageToken"] = next_page_token + else: + start_at += len(results_page) + params["startAt"] = start_at diff --git a/omniload/source/jira/settings.py b/omniload/source/jira/settings.py index 76debf7d3..cf527e803 100644 --- a/omniload/source/jira/settings.py +++ b/omniload/source/jira/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,173 +12,34 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Jira source settings and constants""" - -# Default start date for Jira API requests -DEFAULT_START_DATE = "2010-01-01" - -# Jira API request timeout in seconds -REQUEST_TIMEOUT = 300 - -# Default page size for paginated requests -DEFAULT_PAGE_SIZE = 100 - -# Maximum page size allowed by Jira API -MAX_PAGE_SIZE = 1000 - -# Base API path for Jira Cloud -API_BASE_PATH = "/rest/api/3" - -# Project fields to retrieve from Jira API -PROJECT_FIELDS = ( - "id", - "key", - "name", - "description", - "lead", - "projectCategory", - "projectTypeKey", - "simplified", - "style", - "favourite", - "isPrivate", - "properties", - "entityId", - "uuid", - "insight", -) - -# Issue fields to retrieve from Jira API -ISSUE_FIELDS = ( - "id", - "key", - "summary", - "description", - "issuetype", - "status", - "priority", - "resolution", - "assignee", - "reporter", - "creator", - "created", - "updated", - "resolutiondate", - "duedate", - "components", - "fixVersions", - "versions", - "labels", - "environment", - "project", - "parent", - "subtasks", - "issuelinks", - "votes", - "watches", - "worklog", - "attachments", - "comment", - "customfield_*", -) - -# User fields to retrieve from Jira API -USER_FIELDS = ( - "accountId", - "accountType", - "emailAddress", - "displayName", - "active", - "timeZone", - "groups", - "applicationRoles", - "expand", -) - -# Board fields to retrieve from Jira API (for Agile/Scrum boards) -BOARD_FIELDS = ( - "id", - "name", - "type", - "location", - "filter", - "subQuery", -) - -# Sprint fields to retrieve from Jira API -SPRINT_FIELDS = ( - "id", - "name", - "state", - "startDate", - "endDate", - "completeDate", - "originBoardId", - "goal", -) - -# Issue type fields to retrieve from Jira API -ISSUE_TYPE_FIELDS = ( - "id", - "name", - "description", - "iconUrl", - "subtask", - "avatarId", - "hierarchyLevel", -) - -# Status fields to retrieve from Jira API -STATUS_FIELDS = ( - "id", - "name", - "description", - "iconUrl", - "statusCategory", -) - -# Priority fields to retrieve from Jira API -PRIORITY_FIELDS = ( - "id", - "name", - "description", - "iconUrl", -) - -# Resolution fields to retrieve from Jira API -RESOLUTION_FIELDS = ( - "id", - "name", - "description", -) - -# Version fields to retrieve from Jira API -VERSION_FIELDS = ( - "id", - "name", - "description", - "archived", - "released", - "startDate", - "releaseDate", - "overdue", - "userStartDate", - "userReleaseDate", - "project", - "projectId", -) - -# Component fields to retrieve from Jira API -COMPONENT_FIELDS = ( - "id", - "name", - "description", - "lead", - "assigneeType", - "assignee", - "realAssigneeType", - "realAssignee", - "isAssigneeTypeValid", - "project", - "projectId", -) +# Define endpoints +DEFAULT_ENDPOINTS = { + "issues": { + "data_path": "issues", + "api_path": "rest/api/3/search/jql", + "use_cursor_pagination": True, + "params": { + "fields": "*all", + "expand": "fields,changelog,operations,transitions,names", + "validateQuery": "strict", + "jql": "created >= '2000-01-01' order by created DESC", + }, + }, + "users": { + "api_path": "rest/api/3/users", + "params": {"includeInactiveUsers": True}, + }, + "workflows": { + "data_path": "values", + "api_path": "/rest/api/3/workflow/search", + "params": {}, + }, + "projects": { + "data_path": "values", + "api_path": "rest/api/3/project/search", + "params": { + "expand": "description,lead,issueTypes,url,projectKeys,permissions,insight" + }, + }, +} +DEFAULT_PAGE_SIZE = 50 diff --git a/omniload/source/kafka/adapter.py b/omniload/source/kafka/adapter.py index d538fde44..399cd267b 100644 --- a/omniload/source/kafka/adapter.py +++ b/omniload/source/kafka/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,7 +14,7 @@ """A source to extract Kafka messages. -When extraction starts, partition length is checked - +When extraction starts, partitions length is checked - data is read only up to it, overriding the default Kafka's behavior of waiting for new messages in endless loop. """ @@ -23,15 +23,15 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Union import dlt -from confluent_kafka import Consumer, Message +from confluent_kafka import Consumer, Message # type: ignore from dlt.common import logger -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import TAnyDateTime, TDataItem from .helpers import ( KafkaCredentials, - KafkaEventProcessor, OffsetTracker, + default_msg_processor, ) @@ -43,10 +43,13 @@ def kafka_consumer( topics: Union[str, List[str]], credentials: Union[KafkaCredentials, Consumer] = dlt.secrets.value, - msg_processor: Optional[Callable[[Message], Dict[str, Any]]] = None, + msg_processor: Optional[ + Callable[[Message], Dict[str, Any]] + ] = default_msg_processor, batch_size: Optional[int] = 3000, - batch_timeout: Optional[float] = 3.0, + batch_timeout: Optional[int] = 3, start_from: Optional[TAnyDateTime] = None, + end_time: Optional[TAnyDateTime] = None, ) -> Iterable[TDataItem]: """Extract recent messages from the given Kafka topics. @@ -61,23 +64,21 @@ def kafka_consumer( Auth credentials or an initiated Kafka consumer. By default, is taken from secrets. msg_processor(Optional[Callable]): A function-converter, - which will process every Kafka message after it is read and - before it is transferred to the destination. + which'll process every Kafka message after it's read and + before it's transferred to the destination. batch_size (Optional[int]): Messages batch size to read at once. batch_timeout (Optional[int]): Maximum time to wait for a batch - to be consumed in seconds. + consume, in seconds. start_from (Optional[TAnyDateTime]): A timestamp, at which to start reading. Older messages are ignored. + end_time (Optional[TAnyDateTime]): A timestamp, at which to stop + reading. Newer messages are ignored. Yields: Iterable[TDataItem]: Kafka messages. """ - msg_processor = msg_processor or KafkaEventProcessor().process - if not isinstance(topics, list): topics = [topics] - batch_size = batch_size or 3000 - batch_timeout = batch_timeout or 3.0 if isinstance(credentials, Consumer): consumer = credentials @@ -92,33 +93,56 @@ def kafka_consumer( ) if start_from is not None: - start_from = ensure_pendulum_datetime_utc(start_from) + start_from = ensure_pendulum_datetime(start_from) + + if end_time is not None: + end_time = ensure_pendulum_datetime(end_time) + + if start_from is None: + raise ValueError("`start_from` must be provided if `end_time` is provided") + + if start_from > end_time: + raise ValueError("`start_from` must be before `end_time`") + + tracker = OffsetTracker( + consumer, topics, dlt.current.resource_state(), start_from, end_time + ) - tracker = OffsetTracker(consumer, topics, dlt.current.resource_state(), start_from) + else: + tracker = OffsetTracker( + consumer, topics, dlt.current.resource_state(), start_from + ) # read messages up to the maximum offsets, # not waiting for new messages with closing(consumer): - while True: + while tracker.has_unread: messages = consumer.consume(batch_size, timeout=batch_timeout) if not messages: break batch = [] for msg in messages: - err = msg.error() - if err is not None: + if msg.error(): + err = msg.error() if err.retriable() or not err.fatal(): logger.warning(f"ERROR: {err} - RETRYING") - elif isinstance(err, BaseException): - raise err else: - raise RuntimeError(f"Unknown error: {err}") + raise err else: - batch.append(msg_processor(msg)) - tracker.renew(msg) + topic = msg.topic() + partition = str(msg.partition()) + current_offset = msg.offset() + max_offset = tracker[topic][partition]["max"] + + # Only process the message if it's within the partition's max offset + if current_offset < max_offset: + batch.append(msg_processor(msg)) + tracker.renew(msg) + else: + logger.info( + f"Skipping message on {topic} partition {partition} at offset {current_offset} " + f"- beyond max offset {max_offset}" + ) yield batch - - if tracker.has_unread is False: - return diff --git a/omniload/source/kafka/helpers.py b/omniload/source/kafka/helpers.py index 99cd7487e..9c3255f57 100644 --- a/omniload/source/kafka/helpers.py +++ b/omniload/source/kafka/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,94 +14,61 @@ from typing import Any, Dict, List, Optional -import orjson -from confluent_kafka import Consumer, Message, TopicPartition -from confluent_kafka.admin import TopicMetadata -from dlt import config +from confluent_kafka import Consumer, Message, TopicPartition # type: ignore +from confluent_kafka.admin import TopicMetadata # type: ignore +from dlt import config, secrets from dlt.common import pendulum from dlt.common.configuration import configspec from dlt.common.configuration.specs import CredentialsConfiguration +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import DictStrAny, TSecretValue +from dlt.common.utils import digest128 -from omniload.source.kafka.model import KafkaDecodingOptions, KafkaEvent +def default_msg_processor(msg: Message) -> Dict[str, Any]: + """Basic Kafka message processor. -class KafkaEventProcessor: - """ - A processor for Kafka events with processing stages and configuration capabilities. - - It cycles through "decode", "deserialize" and "format". - """ - - def __init__(self, options: Optional[KafkaDecodingOptions] = None): - self.options = options or KafkaDecodingOptions() - - def process(self, msg: Message) -> Dict[str, Any]: - """ - Progress Kafka event. - - Returns the message value and metadata. Timestamp consists of two values: - (type of the timestamp, timestamp). Type represents one of the Python - Kafka constants: - TIMESTAMP_NOT_AVAILABLE - Timestamps not supported by broker. - TIMESTAMP_CREATE_TIME - Message creation time (or source / producer time). - TIMESTAMP_LOG_APPEND_TIME - Broker receive time. - - Args: - msg (confluent_kafka.Message): A single Kafka message. - - Returns: - dict: Processed Kafka message. - """ - - # Decode. - event = self.decode(msg) - - # Deserialize. - self.deserialize(event) - - # Format egress message based on input options. - return event.to_dict(self.options) - - def decode(self, msg: Message) -> KafkaEvent: - """ - Translate from Confluent library's `Message` instance to `Event` instance. - """ - return KafkaEvent( - ts=msg.timestamp(), - topic=msg.topic(), - partition=msg.partition(), - offset=msg.offset(), - key=msg.key(), - value=msg.value(), - ) - - def deserialize(self, event: KafkaEvent) -> KafkaEvent: - """ - Deserialize event key and value according to decoding options. - """ - if self.options.key_type is not None: - if self.options.key_type == "json": - if event.key is not None: - event.key = orjson.loads(event.key) - else: - raise NotImplementedError(f"Unknown key type: {self.options.key_type}") - if self.options.value_type is not None: - if self.options.value_type == "json": - if event.value is not None: - event.value = orjson.loads(event.value) - else: - raise NotImplementedError( - f"Unknown value type: {self.options.value_type}" - ) - return event + Returns the message value and metadata. Timestamp consists of two values: + (type of the timestamp, timestamp). Type represents one of the Python + Kafka constants: + TIMESTAMP_NOT_AVAILABLE - Timestamps not supported by broker. + TIMESTAMP_CREATE_TIME - Message creation time (or source / producer time). + TIMESTAMP_LOG_APPEND_TIME - Broker receive time. + Args: + msg (confluent_kafka.Message): A single Kafka message. -class OffsetTracker(dict): + Returns: + dict: Processed Kafka message. + """ + ts = msg.timestamp() + topic = msg.topic() + partition = msg.partition() + key = msg.key() + if key is not None: + key = key.decode("utf-8") + + return { + "_kafka": { + "partition": partition, + "topic": topic, + "key": key, + "offset": msg.offset(), + "ts": { + "type": ts[0], + "value": ensure_pendulum_datetime(ts[1] / 1e3), + }, + "data": msg.value().decode("utf-8"), + }, + "_kafka_msg_id": digest128(topic + str(partition) + str(key)), + } + + +class OffsetTracker(dict): # type: ignore """Object to control offsets of the given topics. - Tracks all the partitions of the given topics with two params: - current offset and maximum offset (partition length). + Tracks all the partitions of the given topics with three params: + current offset, maximum offset (partition length), and an end time. Args: consumer (confluent_kafka.Consumer): Kafka consumer. @@ -109,6 +76,8 @@ class OffsetTracker(dict): pl_state (DictStrAny): Pipeline current state. start_from (Optional[pendulum.DateTime]): A timestamp, after which messages are read. Older messages are ignored. + end_time (Optional[pendulum.DateTime]): A timestamp, before which messages + are read. Newer messages are ignored. """ def __init__( @@ -117,6 +86,7 @@ def __init__( topic_names: List[str], pl_state: DictStrAny, start_from: Optional[pendulum.DateTime] = None, + end_time: Optional[pendulum.DateTime] = None, ): super().__init__() @@ -128,7 +98,7 @@ def __init__( "offsets", {t_name: {} for t_name in topic_names} ) - self._init_partition_offsets(start_from) + self._init_partition_offsets(start_from, end_time) def _read_topics(self, topic_names: List[str]) -> Dict[str, TopicMetadata]: """Read the given topics metadata from Kafka. @@ -150,7 +120,11 @@ def _read_topics(self, topic_names: List[str]) -> Dict[str, TopicMetadata]: return tracked_topics - def _init_partition_offsets(self, start_from: Optional[pendulum.DateTime]) -> None: + def _init_partition_offsets( + self, + start_from: Optional[pendulum.DateTime] = None, + end_time: Optional[pendulum.DateTime] = None, + ) -> None: """Designate current and maximum offsets for every partition. Current offsets are read from the state, if present. Set equal @@ -159,6 +133,8 @@ def _init_partition_offsets(self, start_from: Optional[pendulum.DateTime]) -> No Args: start_from (pendulum.DateTime): A timestamp, at which to start reading. Older messages are ignored. + end_time (pendulum.DateTime): A timestamp, before which messages + are read. Newer messages are ignored. """ all_parts = [] for t_name, topic in self._topics.items(): @@ -174,27 +150,49 @@ def _init_partition_offsets(self, start_from: Optional[pendulum.DateTime]) -> No for part in topic.partitions ] - # get offsets for the timestamp, if given - if start_from is not None: + # get offsets for the timestamp ranges, if given + if start_from is not None and end_time is not None: + start_ts_offsets = self._consumer.offsets_for_times(parts) + end_ts_offsets = self._consumer.offsets_for_times( + [ + TopicPartition(t_name, part, end_time.int_timestamp * 1000) + for part in topic.partitions + ] + ) + elif start_from is not None: ts_offsets = self._consumer.offsets_for_times(parts) # designate current and maximum offsets for every partition for i, part in enumerate(parts): max_offset = self._consumer.get_watermark_offsets(part)[1] - if start_from is not None: + if start_from is not None and end_time is not None: + if start_ts_offsets[i].offset != -1: + cur_offset = start_ts_offsets[i].offset + else: + cur_offset = max_offset - 1 + if end_ts_offsets[i].offset != -1: + end_offset = end_ts_offsets[i].offset + else: + end_offset = max_offset + + elif start_from is not None: if ts_offsets[i].offset != -1: cur_offset = ts_offsets[i].offset else: cur_offset = max_offset - 1 + + end_offset = max_offset + else: cur_offset = ( self._cur_offsets[t_name].get(str(part.partition), -1) + 1 ) + end_offset = max_offset self[t_name][str(part.partition)] = { "cur": cur_offset, - "max": max_offset, + "max": end_offset, } parts[i].offset = cur_offset @@ -245,10 +243,12 @@ class KafkaCredentials(CredentialsConfiguration): bootstrap_servers: str = config.value group_id: str = config.value - security_protocol: Optional[str] = None - sasl_mechanisms: Optional[str] = None - sasl_username: Optional[str] = None - sasl_password: Optional[TSecretValue] = None + security_protocol: str = config.value + + # Optional SASL credentials + sasl_mechanisms: Optional[str] = config.value + sasl_username: Optional[str] = config.value + sasl_password: Optional[TSecretValue] = secrets.value def init_consumer(self) -> Consumer: """Init a Kafka consumer from this credentials. @@ -259,16 +259,17 @@ def init_consumer(self) -> Consumer: config = { "bootstrap.servers": self.bootstrap_servers, "group.id": self.group_id, + "security.protocol": self.security_protocol, "auto.offset.reset": "earliest", } - if self.security_protocol: - config["security.protocol"] = self.security_protocol - if self.sasl_mechanisms: - config["sasl.mechanisms"] = self.sasl_mechanisms - if self.sasl_username: - config["sasl.username"] = self.sasl_username - if self.sasl_password: - config["sasl.password"] = self.sasl_password + if self.sasl_mechanisms and self.sasl_username and self.sasl_password: + config.update( + { + "sasl.mechanisms": self.sasl_mechanisms, + "sasl.username": self.sasl_username, + "sasl.password": self.sasl_password, + } + ) return Consumer(config) diff --git a/omniload/source/kinesis/adapter.py b/omniload/source/kinesis/adapter.py index 664976af3..984bb0b1e 100644 --- a/omniload/source/kinesis/adapter.py +++ b/omniload/source/kinesis/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,32 +14,31 @@ """Reads messages from Kinesis queue.""" -from typing import Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional import dlt from dlt.common import json, pendulum from dlt.common.configuration.specs import AwsCredentials -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import StrStr, TAnyDateTime, TDataItem from dlt.common.utils import digest128 -from .helpers import get_shard_iterator, get_stream_address, max_sequence_by_shard +from .helpers import get_shard_iterator, max_sequence_by_shard @dlt.resource( name=lambda args: args["stream_name"], - primary_key="kinesis_msg_id", + primary_key="_kinesis_msg_id", standalone=True, - max_table_nesting=0, ) def kinesis_stream( - stream_name: str, - initial_at_timestamp: Optional[TAnyDateTime], - credentials: AwsCredentials, + stream_name: str = dlt.config.value, + credentials: AwsCredentials = dlt.secrets.value, last_msg: Optional[dlt.sources.incremental[StrStr]] = dlt.sources.incremental( - "kinesis", last_value_func=max_sequence_by_shard + "_kinesis", last_value_func=max_sequence_by_shard ), - max_number_of_messages: Optional[int] = None, + initial_at_timestamp: TAnyDateTime = 0.0, + max_number_of_messages: int = None, milliseconds_behind_latest: int = 1000, parse_json: bool = True, chunk_size: int = 1000, @@ -56,10 +55,10 @@ def kinesis_stream( initial_at_timestamp (TAnyDateTime): An initial timestamp used to generate AT_TIMESTAMP or LATEST iterator when timestamp value is 0 max_number_of_messages (int): Maximum number of messages to read in one run. Actual read may exceed that number by up to chunk_size. Defaults to None (no limit). milliseconds_behind_latest (int): The number of milliseconds behind the top of the shard to stop reading messages, defaults to 1000. - parse_json (bool): If True, assumes that messages are json strings, parses them and returns instead of `data` (otherwise). Defaults to True. + parse_json (bool): If True, assumes that messages are json strings, parses them and returns instead of `data` (otherwise). Defaults to False. chunk_size (int): The number of records to fetch at once. Defaults to 1000. Yields: - Iterable[TDataItem]: Messages. Contain Kinesis envelope in `kinesis` and bytes data in `data` (if `parse_json` disabled) + Iterable[TDataItem]: Messages. Contain Kinesis envelope in `_kinesis` and bytes data in `data` (if `parse_json` disabled) """ session = credentials._to_botocore_session() @@ -69,7 +68,7 @@ def kinesis_stream( initial_at_datetime = ( None if initial_at_timestamp is None - else ensure_pendulum_datetime_utc(initial_at_timestamp) + else ensure_pendulum_datetime(initial_at_timestamp) ) # set it in state resource_state = dlt.current.resource_state() @@ -77,9 +76,9 @@ def kinesis_stream( "initial_at_timestamp", initial_at_datetime ) # so next time we request shards at AT_TIMESTAMP that is now - resource_state["initial_at_timestamp"] = pendulum.now("UTC").subtract(seconds=1) + resource_state["initial_at_timestamp"] = pendulum.now().subtract(seconds=1) - shards_list = kinesis_client.list_shards(**get_stream_address(stream_name)) + shards_list = kinesis_client.list_shards(StreamName=stream_name) shards: List[StrStr] = shards_list["Shards"] while next_token := shards_list.get("NextToken"): shards_list = kinesis_client.list_shards(NextToken=next_token) @@ -90,38 +89,31 @@ def kinesis_stream( # get next shard to fetch messages from while shard_id := shard_ids.pop(0) if shard_ids else None: shard_iterator, _ = get_shard_iterator( - kinesis_client, - stream_name, - shard_id, - last_msg, # ty: ignore[invalid-argument-type] - initial_at_datetime, + kinesis_client, stream_name, shard_id, last_msg, initial_at_datetime ) - + # fetch messages from shard iterator or exit loop if not present while shard_iterator: records = [] records_response = kinesis_client.get_records( ShardIterator=shard_iterator, Limit=chunk_size, # The size of data can be up to 1 MB, it must be controlled by the user ) - for record in records_response["Records"]: sequence_number = record["SequenceNumber"] content = record["Data"] - arrival_time = record["ApproximateArrivalTimestamp"] - arrival_timestamp = arrival_time.astimezone(pendulum.UTC) - message = { - "kinesis": { + "_kinesis": { "shard_id": shard_id, "seq_no": sequence_number, - "ts": ensure_pendulum_datetime_utc(arrival_timestamp), + "ts": ensure_pendulum_datetime( + record["ApproximateArrivalTimestamp"] + ), "partition": record["PartitionKey"], "stream_name": stream_name, }, - "kinesis_msg_id": digest128(shard_id + sequence_number), + "_kinesis_msg_id": digest128(shard_id + sequence_number), } - if parse_json: message.update(json.loadb(content)) else: diff --git a/omniload/source/kinesis/helpers.py b/omniload/source/kinesis/helpers.py index 23d612c8c..78bf235bc 100644 --- a/omniload/source/kinesis/helpers.py +++ b/omniload/source/kinesis/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,7 +16,7 @@ import dlt from dlt.common import pendulum -from dlt.common.typing import DictStrAny, DictStrStr, StrAny, StrStr +from dlt.common.typing import DictStrAny, StrAny, StrStr def get_shard_iterator( @@ -24,7 +24,7 @@ def get_shard_iterator( stream_name: str, shard_id: str, last_msg: dlt.sources.incremental[StrStr], - initial_at_timestamp: pendulum.DateTime | None, + initial_at_timestamp: pendulum.DateTime, ) -> Tuple[str, StrAny]: """Gets shard `shard_id` of `stream_name` iterator. If `last_msg` incremental is present it may contain last message sequence for shard_id. in that case AFTER_SEQUENCE_NUMBER is created. @@ -44,9 +44,7 @@ def get_shard_iterator( elif initial_at_timestamp is None: # Fetch all records from the beginning iterator_params = dict(ShardIteratorType="TRIM_HORIZON") - elif initial_at_timestamp.timestamp() == 0.0: - # will sets to latest i.e only the messages at the tip of the stream are read iterator_params = dict(ShardIteratorType="LATEST") else: iterator_params = dict( @@ -54,7 +52,7 @@ def get_shard_iterator( ) shard_iterator: StrStr = kinesis_client.get_shard_iterator( - **get_stream_address(stream_name), ShardId=shard_id, **iterator_params + StreamName=stream_name, ShardId=shard_id, **iterator_params ) return shard_iterator["ShardIterator"], iterator_params @@ -77,20 +75,3 @@ def max_sequence_by_shard(values: Sequence[StrStr]) -> StrStr: # we compare message sequence at shard_id last_value[shard_id] = max(item["seq_no"], last_value.get(shard_id, "")) return last_value - - -def get_stream_address(stream_name: str) -> DictStrStr: - """ - Return address of stream, either as StreamName or StreamARN, when applicable. - - Examples: - - customer_events - - arn:aws:kinesis:eu-central-1:842404475894:stream/customer_events - - https://docs.aws.amazon.com/kinesis/latest/APIReference/API_StreamDescription.html#Streams-Type-StreamDescription-StreamName - https://docs.aws.amazon.com/kinesis/latest/APIReference/API_StreamDescription.html#Streams-Type-StreamDescription-StreamARN - """ - if stream_name.startswith("arn:"): - return {"StreamARN": stream_name} - else: - return {"StreamName": stream_name} diff --git a/omniload/source/matomo/adapter.py b/omniload/source/matomo/adapter.py new file mode 100644 index 000000000..5d0ced70a --- /dev/null +++ b/omniload/source/matomo/adapter.py @@ -0,0 +1,245 @@ +# Copyright 2022-2026 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. + +"""Loads reports and raw visits data from Matomo""" + +from typing import Iterable, Iterator, List + +import dlt +import pendulum +from dlt.common.typing import DictStrAny, TDataItem +from dlt.sources import DltResource + +from .helpers.data_processing import ( + get_matomo_date_range, + process_report, + remove_active_visits, +) +from .helpers.matomo_client import MatomoAPIClient + + +@dlt.source(max_table_nesting=2) +def matomo_reports( + api_token: str = dlt.secrets.value, + url: str = dlt.config.value, + queries: List[DictStrAny] = dlt.config.value, + site_id: int = dlt.config.value, +) -> Iterable[DltResource]: + """ + Executes and loads a set of reports defined in `queries` for a Matomo site site_id. + + Args: + api_token (str): The API access token for authentication. Defaults to the value in the `dlt.secrets` object. + url (str): The URL of the Matomo server. Defaults to the value in the `dlt.config` object. + queries (List[DictStrAny]): Dicts that contain information on the reports to retrieve. + site_id (int): Site ID of the website as per your Matomo account. + + Returns: + Iterable[DltResource]: List of dlt resources, each containing a Matomo report. + """ + + # Create an instance of the Matomo API client + client = MatomoAPIClient(api_token=api_token, url=url) + for query in queries: + name = query.get("resource_name") + batch_data = dlt.resource( + get_data_batch, write_disposition="append", name=name + )( + client=client, + query=query, + site_id=site_id, + last_date=dlt.sources.incremental( + "date", primary_key=() + ), # disables unique checks in incremental + ) + yield batch_data + + +def get_data_batch( + client: MatomoAPIClient, + query: DictStrAny, + last_date: dlt.sources.incremental[pendulum.DateTime], + site_id: int, +) -> Iterator[TDataItem]: + """ + Gets data for a query in a batch. + + Args: + client (MatomoAPIClient): API Client used to make calls to Matomo API. + query (DictStrAny): Dict specified in config.toml containing info on what data to retrieve from Matomo API. + last_date (dlt.sources.incremental[pendulum.DateTime]): Last date timestamp of the last data loaded for the batch data. + site_id (int): Specifies the Matomo site data is gathered from. + + Returns: + Iterator[TDataItem]: Iterator yielding data items for the query. + """ + + # unpack query data and process date range + name = query["resource_name"] + extra_params = query.get("extra_params", {}) + methods = query["methods"] + period = query.get("period", "day") + start_date = query.get("date", None) + date_range = get_matomo_date_range(start_date=start_date, last_date=last_date) + + # Get all method data returned for a single query and iterate through pages received + reports = client.get_query( + date=date_range, + extra_params=extra_params, + methods=methods, + period=period, + site_id=site_id, + ) + for method_data, method in zip(reports, methods): + # process data for every method and save it in the correct table (1 for each method data) + table_name = f"{name}_{method}" + processed_report_generator = process_report(method_data) + for data in processed_report_generator: + yield dlt.mark.with_table_name(data, table_name) + + +@dlt.source(max_table_nesting=2) +def matomo_visits( + api_token: str = dlt.secrets.value, + url: str = dlt.config.value, + live_events_site_id: int = dlt.config.value, + initial_load_past_days: int = 10, + visit_timeout_seconds: int = 1800, + visit_max_duration_seconds: int = 3600, + get_live_event_visitors: bool = False, +) -> List[DltResource]: + """ + Loads a list of visits. Initially loads the current day visits and all visits in `initial_past_days`. + On subsequent loads, continues from last load. Active visits are skipped until a visit is closed and does not get more actions. + + Args: + api_token (str): The API access token for authentication. Defaults to the value in the `dlt.secrets` object. + url (str): The URL of the Matomo server. Defaults to the value in the `dlt.config` object. + live_events_site_id (int): Site ID of the website for live events. + initial_load_past_days (int): Number of past days to load on initial load. Defaults to 10. + visit_timeout_seconds (int): Session timeout in seconds, after which a visit is considered closed. Defaults to 1800. + visit_max_duration_seconds (int): Maximum visit duration in seconds after which a visit is considered closed. Defaults to 3600. + get_live_event_visitors (bool): Option to retrieve data about unique visitors. Defaults to False. + + Returns: + List[DltResource]: A list of resources containing event data. + """ + + # Create an instance of the Matomo API client + client = MatomoAPIClient(api_token=api_token, url=url) + # 10 days into the past is a default + initial_server_timestamp = ( + pendulum.today().subtract(days=initial_load_past_days).timestamp() + ) + visits_data_generator = get_last_visits( + client=client, + site_id=live_events_site_id, + last_date=dlt.sources.incremental( + "serverTimestamp", initial_value=initial_server_timestamp + ), + visit_timeout_seconds=visit_timeout_seconds, + visit_max_duration_seconds=visit_max_duration_seconds, + ) + resource_list = [visits_data_generator] + if get_live_event_visitors: + resource_list.append( + visits_data_generator + | get_unique_visitors(client=client, site_id=live_events_site_id) + ) + return resource_list + + +@dlt.resource( + name="visits", write_disposition="append", primary_key="idVisit", selected=True +) +def get_last_visits( + client: MatomoAPIClient, + site_id: int, + last_date: dlt.sources.incremental[float], + visit_timeout_seconds: int = 1800, + visit_max_duration_seconds: int = 3600, + rows_per_page: int = 2000, +) -> Iterator[TDataItem]: + """ + Dlt resource which gets the visits in the given site id for the given timeframe. If there was a previous load, the chosen timeframe is always last_load -> now. + Otherwise, a start date can be provided to make the chosen timeframe start_date -> now. + The default behavior would be to get all visits available until now if neither last_load nor start_date are given. + + Args: + client (MatomoAPIClient): Used to make calls to Matomo API. + site_id (int): Every site in Matomo has a unique id. + last_date (dlt.sources.incremental[float]): Date of last load for this resource if it exists. + visit_timeout_seconds (int): A session timeout in seconds, after that inactivity time, visit is considered closed. Defaults to 1800. + visit_max_duration_seconds (int): Maximum visit length in seconds after which we consider the visit closed. Defaults to 3600. + rows_per_page (int): How many rows will be there per page. + + Returns: + Iterator[TDataItem]: Iterator of dicts containing information on last visits in the given timeframe. + """ + + extra_params = { + "minTimestamp": last_date.last_value, + } + now = pendulum.now().timestamp() + visits = client.get_method( + site_id=site_id, + method="Live.getLastVisitsDetails", + extra_params=extra_params, + rows_per_page=rows_per_page, + ) + for page_no, page in enumerate(visits): + # remove active visits only from the first page + if page_no == 0: + page = remove_active_visits( + page, visit_timeout_seconds, visit_max_duration_seconds, now + ) + yield page + + +@dlt.transformer( + data_from=get_last_visits, + write_disposition="merge", + name="visitors", + primary_key="visitorId", +) +def get_unique_visitors( + visits: List[DictStrAny], + client: MatomoAPIClient, + site_id: int, + chunk_size: int = 20, +) -> Iterator[TDataItem]: + """ + Dlt transformer. Receives information about visits from get_last_visits. + This version allows batch loading for visitors data, which is to avoid too-long-URL issue + + Args: + visits (List[DictStrAny]): List of dicts containing information on last visits in the given timeframe. + client (MatomoAPIClient): Used to make calls to Matomo API. + site_id (int): Every site in Matomo has a unique id. + chunk_size (int): Number of visitor IDs to process in each batch. Defaults to 20. + + Returns: + Iterator[TDataItem]: Dict containing information about the visitor. + """ + + visitor_ids = [visit["visitorId"] for visit in visits] + indexed_visitor_ids = [ + visitor_ids[i : i + chunk_size] for i in range(0, len(visitor_ids), chunk_size) + ] + for visitor_list in indexed_visitor_ids: + method_data = client.get_visitors_batch( + visitor_list=visitor_list, site_id=site_id + ) + for method_dict in method_data: + yield method_dict diff --git a/omniload/source/matomo/settings.py b/omniload/source/matomo/settings.py new file mode 100644 index 000000000..85b335032 --- /dev/null +++ b/omniload/source/matomo/settings.py @@ -0,0 +1,17 @@ +# Copyright 2022-2026 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. + +"""Matomo source settings and constants""" + +DEFAULT_START_DATE = "2000-01-01" diff --git a/omniload/source/mongodb/adapter.py b/omniload/source/mongodb/adapter.py index 3b3ed4871..0bf315fd0 100644 --- a/omniload/source/mongodb/adapter.py +++ b/omniload/source/mongodb/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,6 +17,7 @@ from typing import Any, Dict, Iterable, List, Mapping, Optional, Union import dlt +from dlt.common.configuration.specs.config_section_context import ConfigSectionContext from dlt.common.data_writers import TDataItemFormat from dlt.sources import DltResource @@ -25,16 +26,15 @@ MongoDbCollectionResourceConfiguration, client_from_credentials, collection_documents, - process_file_items, ) -@dlt.source(max_table_nesting=0) +@dlt.source def mongodb( connection_url: str = dlt.secrets.value, database: Optional[str] = dlt.config.value, collection_names: Optional[List[str]] = dlt.config.value, - incremental: Optional[dlt.sources.incremental] = None, + incremental: Optional[dlt.sources.incremental] = None, # type: ignore[type-arg] write_disposition: Optional[str] = dlt.config.value, parallel: Optional[bool] = dlt.config.value, limit: Optional[int] = None, @@ -90,7 +90,6 @@ def mongodb( primary_key="_id", write_disposition=write_disposition, spec=MongoDbCollectionConfiguration, - max_table_nesting=0, )( client, collection, @@ -107,23 +106,21 @@ def mongodb( name=lambda args: args["collection"], standalone=True, spec=MongoDbCollectionResourceConfiguration, - max_table_nesting=0, ) def mongodb_collection( connection_url: str = dlt.secrets.value, database: Optional[str] = dlt.config.value, collection: str = dlt.config.value, - incremental: Optional[dlt.sources.incremental] = None, + incremental: Optional[dlt.sources.incremental] = None, # type: ignore[type-arg] write_disposition: Optional[str] = dlt.config.value, parallel: Optional[bool] = False, limit: Optional[int] = None, - chunk_size: Optional[int] = 1000, + chunk_size: Optional[int] = 10000, data_item_format: Optional[TDataItemFormat] = "object", filter_: Optional[Dict[str, Any]] = None, projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = dlt.config.value, pymongoarrow_schema: Optional[Any] = None, - custom_query: Optional[List[Dict[str, Any]]] = None, -) -> DltResource: +) -> Any: """ A DLT source which loads a collection from a mongo database using PyMongo. @@ -149,7 +146,6 @@ def mongodb_collection( exclude (dict) - {"released": False, "runtime": False} Note: Can't mix include and exclude statements '{"title": True, "released": False}` pymongoarrow_schema (pymongoarrow.schema.Schema): Mapping of expected field types to convert BSON to Arrow - custom_query (Optional[List[Dict[str, Any]]]): Custom MongoDB aggregation pipeline to execute instead of find() Returns: Iterable[DltResource]: A list of DLT resources for each collection to be loaded. @@ -179,103 +175,4 @@ def mongodb_collection( filter_=filter_ or {}, projection=projection, pymongoarrow_schema=pymongoarrow_schema, - custom_query=custom_query, - ) - - -def mongodb_insert(uri: str): - """Creates a dlt.destination for inserting data into a MongoDB collection. - - Args: - uri (str): MongoDB connection URI including database. - - Returns: - dlt.destination: A DLT destination object configured for MongoDB. - """ - from urllib.parse import urlparse - - parsed_uri = urlparse(uri) - database = ( - parsed_uri.path.lstrip("/") if parsed_uri.path.lstrip("/") else "omniload_db" - ) - first_batch_per_table: dict[str, bool] = {} - BATCH_SIZE = 10000 - - def destination(items, table) -> None: - import pyarrow - from pymongo import MongoClient - - collection_name = table["name"] - - if collection_name not in first_batch_per_table: - first_batch_per_table[collection_name] = True - - with MongoClient(uri) as client: - db = client[database] - collection = db[collection_name] - - # Process documents - if isinstance(items, str): - documents = process_file_items(items) - elif isinstance(items, pyarrow.RecordBatch): - documents = items.to_pylist() - else: - documents = [item for item in items if isinstance(item, dict)] - - write_disposition = table.get("write_disposition") - - batches = [ - documents[i : i + BATCH_SIZE] - for i in range(0, len(documents), BATCH_SIZE) - ] - - if write_disposition == "merge": - from pymongo import ReplaceOne - - primary_keys = [ - col_name - for col_name, col_def in table.get("columns", {}).items() - if isinstance(col_def, dict) and col_def.get("primary_key") - ] - - if not primary_keys: - raise ValueError( - f"Merge operation requires primary keys for table '{collection_name}'. " - f"Please define primary keys in the table schema or use 'replace' write disposition." - ) - - for batch in batches: - operations = [ - ReplaceOne( - {key: doc[key] for key in primary_keys}, - doc, - upsert=True, - ) - for doc in batch - if all(key in doc for key in primary_keys) - ] - if operations: - collection.bulk_write(operations, ordered=False) - - elif write_disposition == "replace": - if first_batch_per_table[collection_name] and documents: - collection.delete_many({}) - first_batch_per_table[collection_name] = False - - for batch in batches: - if batch: - collection.insert_many(batch) - - else: - raise ValueError( - f"Unsupported write disposition '{write_disposition}' for MongoDB destination. " - ) - - return dlt.destination( - destination, - name="mongodb", - loader_file_format="typed-jsonl", - batch_size=1000, - naming_convention="snake_case", - loader_parallelism_strategy="sequential", ) diff --git a/omniload/source/mongodb/helpers.py b/omniload/source/mongodb/helpers.py index 55104e879..73bea27d9 100644 --- a/omniload/source/mongodb/helpers.py +++ b/omniload/source/mongodb/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Mongo database source helpers and destination utilities""" +"""Mongo database source helpers""" -import re from itertools import islice from typing import ( TYPE_CHECKING, @@ -37,7 +36,7 @@ from dlt.common import logger from dlt.common.configuration.specs import BaseConfiguration, configspec from dlt.common.data_writers import TDataItemFormat -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import TDataItem from dlt.common.utils import map_nested_in_place from pendulum import _datetime @@ -56,8 +55,7 @@ TCursor = Any try: - import pymongoarrow - import pymongoarrow.schema + import pymongoarrow # type: ignore PYMONGOARROW_AVAILABLE = True except ImportError: @@ -69,11 +67,9 @@ def __init__( self, client: TMongoClient, collection: TCollection, - chunk_size: Optional[int], + chunk_size: int, incremental: Optional[dlt.sources.incremental[Any]] = None, ) -> None: - chunk_size = chunk_size if chunk_size is not None else 10000 - self.client = client self.collection = collection self.incremental = incremental @@ -163,7 +159,7 @@ def _projection_op( else: try: # ensure primary_key isn't excluded - projection_dict.pop(self.incremental.primary_key) # ty: ignore[invalid-argument-type] + projection_dict.pop(self.incremental.primary_key) # type: ignore except KeyError: pass # primary_key was properly not included in exclusion projection else: @@ -173,7 +169,7 @@ def _projection_op( return projection_dict - def _limit(self, cursor: Cursor, limit: Optional[int] = None) -> TCursor: + def _limit(self, cursor: Cursor, limit: Optional[int] = None) -> TCursor: # type: ignore """Apply a limit to the cursor, if needed. Args: @@ -217,19 +213,12 @@ def load_documents( cursor = self.collection.find(filter=filter_op, projection=projection_op) if self._sort_op: - cursor = cursor.sort(self._sort_op) # ty: ignore[invalid-argument-type] + cursor = cursor.sort(self._sort_op) cursor = self._limit(cursor, limit) while docs_slice := list(islice(cursor, self.chunk_size)): - res = map_nested_in_place(convert_mongo_objs, docs_slice) - if len(res) > 0 and "_id" in res[0] and isinstance(res[0]["_id"], dict): - yield dlt.mark.with_hints( - res, - dlt.mark.make_hints(columns={"_id": {"data_type": "json"}}), - ) - else: - yield res + yield map_nested_in_place(convert_mongo_objs, docs_slice) class CollectionLoaderParallel(CollectionLoader): @@ -272,7 +261,7 @@ def _get_cursor( cursor = self.collection.find(filter=filter_op, projection=projection_op) if self._sort_op: - cursor = cursor.sort(self._sort_op) # ty: ignore[invalid-argument-type] + cursor = cursor.sort(self._sort_op) return cursor @@ -355,7 +344,8 @@ def load_documents( Yields: Iterator[Any]: An iterator of the loaded documents. """ - from pymongoarrow.context import PyMongoArrowContext + from pymongoarrow.context import PyMongoArrowContext # type: ignore + from pymongoarrow.lib import process_bson_stream # type: ignore filter_op = self._filter_op _raise_if_intersection(filter_op, filter_) @@ -368,15 +358,15 @@ def load_documents( filter_, batch_size=self.chunk_size, projection=projection_op ) if self._sort_op: - cursor = cursor.sort(self._sort_op) # ty: ignore[invalid-argument-type] + cursor = cursor.sort(self._sort_op) # type: ignore - cursor = self._limit(cursor, limit) + cursor = self._limit(cursor, limit) # type: ignore - context = PyMongoArrowContext( + context = PyMongoArrowContext.from_schema( schema=pymongoarrow_schema, codec_options=self.collection.codec_options ) for batch in cursor: - context.process_bson_stream(batch) + process_bson_stream(batch, context) table = context.finish() yield convert_arrow_columns(table) @@ -463,7 +453,7 @@ def _get_cursor( filter=filter_op, batch_size=self.chunk_size, projection=projection_op ) if self._sort_op: - cursor = cursor.sort(self._sort_op) # ty: ignore[invalid-argument-type] + cursor = cursor.sort(self._sort_op) # type: ignore return cursor @@ -475,183 +465,19 @@ def _run_batch( pymongoarrow_schema: Any = None, ) -> TDataItem: from pymongoarrow.context import PyMongoArrowContext + from pymongoarrow.lib import process_bson_stream cursor = cursor.clone() - context = PyMongoArrowContext( + context = PyMongoArrowContext.from_schema( schema=pymongoarrow_schema, codec_options=self.collection.codec_options ) for chunk in cursor.skip(batch["skip"]).limit(batch["limit"]): - context.process_bson_stream(chunk) + process_bson_stream(chunk, context) table = context.finish() yield convert_arrow_columns(table) -class CollectionAggregationLoader(CollectionLoader): - """ - MongoDB collection loader that uses aggregation pipelines instead of find queries. - """ - - def __init__( - self, - client: TMongoClient, - collection: TCollection, - chunk_size: Optional[int] = None, - incremental: Optional[dlt.sources.incremental[Any]] = None, - ) -> None: - chunk_size = chunk_size if chunk_size is not None else 10000 - super().__init__(client, collection, chunk_size, incremental) - self.custom_query: Optional[List[Dict[str, Any]]] = None - - def set_custom_query(self, query: List[Dict[str, Any]]): - """Set the custom aggregation pipeline query""" - self.custom_query = query - - def load_documents( - self, - filter_: Dict[str, Any], - limit: Optional[int] = None, - projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None, - ) -> Iterator[TDataItem]: - """Load documents using aggregation pipeline""" - if not self.custom_query: - # Fallback to parent method if no custom query - yield from super().load_documents(filter_, limit, projection) - return - - # Build aggregation pipeline - pipeline = list(self.custom_query) # Copy the query - - # For custom queries, we assume incremental filtering is already handled - # via interval placeholders (:interval_start, :interval_end) in the query itself. - # We don't add additional incremental filtering to avoid conflicts. - - # Add additional filter if provided - if filter_: - filter_match = {"$match": filter_} - pipeline.insert(0, filter_match) - - # Add limit if specified - if limit and limit > 0: - pipeline.append({"$limit": limit}) - - # Add maxTimeMS to prevent hanging - cursor = self.collection.aggregate( - pipeline, - allowDiskUse=True, - batchSize=min(self.chunk_size, 101), - maxTimeMS=30000, # 30 second timeout - ) - - docs_buffer = [] - try: - for doc in cursor: - docs_buffer.append(doc) - - if len(docs_buffer) >= self.chunk_size: - res = map_nested_in_place(convert_mongo_objs, docs_buffer) - if ( - len(res) > 0 - and "_id" in res[0] - and isinstance(res[0]["_id"], dict) - ): - yield dlt.mark.with_hints( - res, - dlt.mark.make_hints(columns={"_id": {"data_type": "json"}}), - ) - else: - yield res - docs_buffer = [] - - # Yield any remaining documents - if docs_buffer: - res = map_nested_in_place(convert_mongo_objs, docs_buffer) - if len(res) > 0 and "_id" in res[0] and isinstance(res[0]["_id"], dict): - yield dlt.mark.with_hints( - res, - dlt.mark.make_hints(columns={"_id": {"data_type": "json"}}), - ) - else: - yield res - finally: - cursor.close() - - -class CollectionAggregationLoaderParallel(CollectionAggregationLoader): - """ - MongoDB collection parallel loader that uses aggregation pipelines. - Note: Parallel loading is not supported for aggregation pipelines due to cursor limitations. - Falls back to sequential loading. - """ - - def load_documents( - self, - filter_: Dict[str, Any], - limit: Optional[int] = None, - projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None, - ) -> Iterator[TDataItem]: - """Load documents using aggregation pipeline (sequential only)""" - logger.warning( - "Parallel loading is not supported for MongoDB aggregation pipelines. Using sequential loading." - ) - yield from super().load_documents(filter_, limit, projection) - - -class CollectionAggregationArrowLoader(CollectionAggregationLoader): - """ - MongoDB collection aggregation loader that uses Apache Arrow for data processing. - """ - - def load_documents( - self, - filter_: Dict[str, Any], - limit: Optional[int] = None, - projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None, - pymongoarrow_schema: Any = None, - ) -> Iterator[Any]: - """Load documents using aggregation pipeline with Arrow format""" - logger.warning( - "Arrow format is not directly supported for MongoDB aggregation pipelines. Converting to Arrow after loading." - ) - - # Load documents normally and convert to arrow format - for batch in super().load_documents(filter_, limit, projection): - if batch: # Only process non-empty batches - try: - from dlt.common.libs.pyarrow import pyarrow - - # Convert dict batch to arrow table - table = pyarrow.Table.from_pylist(batch) - yield convert_arrow_columns(table) - except ImportError: - logger.warning( - "PyArrow not available, falling back to object format" - ) - yield batch - - -class CollectionAggregationArrowLoaderParallel(CollectionAggregationArrowLoader): - """ - MongoDB collection parallel aggregation loader with Arrow support. - Falls back to sequential loading. - """ - - def load_documents( - self, - filter_: Dict[str, Any], - limit: Optional[int] = None, - projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = None, - pymongoarrow_schema: Any = None, - ) -> Iterator[TDataItem]: - """Load documents using aggregation pipeline with Arrow format (sequential only)""" - logger.warning( - "Parallel loading is not supported for MongoDB aggregation pipelines. Using sequential loading." - ) - yield from super().load_documents( - filter_, limit, projection, pymongoarrow_schema - ) - - def collection_documents( client: TMongoClient, collection: TCollection, @@ -661,9 +487,8 @@ def collection_documents( incremental: Optional[dlt.sources.incremental[Any]] = None, parallel: bool = False, limit: Optional[int] = None, - chunk_size: Optional[int] = None, + chunk_size: Optional[int] = 10000, data_item_format: Optional[TDataItemFormat] = "object", - custom_query: Optional[List[Dict[str, Any]]] = None, ) -> Iterator[TDataItem]: """ A DLT source which loads data from a Mongo database using PyMongo. @@ -688,14 +513,10 @@ def collection_documents( Supported formats: object - Python objects (dicts, lists). arrow - Apache Arrow tables. - custom_query (Optional[List[Dict[str, Any]]]): Custom MongoDB aggregation pipeline to execute instead of find() Returns: Iterable[DltResource]: A list of DLT resources for each collection to be loaded. """ - - chunk_size = chunk_size if chunk_size is not None else 10000 - if data_item_format == "arrow" and not PYMONGOARROW_AVAILABLE: dlt.common.logger.warn( "'pymongoarrow' is not installed; falling back to standard MongoDB CollectionLoader." @@ -714,48 +535,21 @@ def collection_documents( "create a projection to select fields, `projection` will be ignored." ) - # If custom query is provided, use aggregation loaders - if custom_query: - if parallel: - if data_item_format == "arrow": - LoaderClass = CollectionAggregationArrowLoaderParallel - else: - LoaderClass = CollectionAggregationLoaderParallel + if parallel: + if data_item_format == "arrow": + LoaderClass = CollectionArrowLoaderParallel else: - if data_item_format == "arrow": - LoaderClass = CollectionAggregationArrowLoader - else: - LoaderClass = CollectionAggregationLoader + LoaderClass = CollectionLoaderParallel # type: ignore else: - if parallel: - if data_item_format == "arrow": - LoaderClass = CollectionArrowLoaderParallel - else: - LoaderClass = CollectionLoaderParallel + if data_item_format == "arrow": + LoaderClass = CollectionArrowLoader # type: ignore else: - if data_item_format == "arrow": - LoaderClass = CollectionArrowLoader - else: - LoaderClass = CollectionLoader + LoaderClass = CollectionLoader # type: ignore loader = LoaderClass( client, collection, incremental=incremental, chunk_size=chunk_size ) - - # Set custom query if provided - if custom_query and hasattr(loader, "set_custom_query"): - loader.set_custom_query(custom_query) # ty: ignore[call-non-callable] - - # Load documents based on loader type - if isinstance( - loader, - ( - CollectionArrowLoader, - CollectionArrowLoaderParallel, - CollectionAggregationArrowLoader, - CollectionAggregationArrowLoaderParallel, - ), - ): + if isinstance(loader, (CollectionArrowLoader, CollectionArrowLoaderParallel)): yield from loader.load_documents( limit=limit, filter_=filter_, @@ -778,12 +572,12 @@ def convert_mongo_objs(value: Any) -> Any: if isinstance(value, (ObjectId, Decimal128)): return str(value) if isinstance(value, _datetime.datetime): - return ensure_pendulum_datetime_utc(value) + return ensure_pendulum_datetime(value) if isinstance(value, Regex): return value.try_compile().pattern if isinstance(value, Timestamp): date = value.as_datetime() - return ensure_pendulum_datetime_utc(date) + return ensure_pendulum_datetime(date) return value @@ -810,7 +604,7 @@ def convert_arrow_columns(table: Any) -> Any: pyarrow.lib.Table: The table with the columns converted. """ from dlt.common.libs.pyarrow import pyarrow - from pymongoarrow.types import ( + from pymongoarrow.types import ( # type: ignore _is_binary, _is_code, _is_decimal128, @@ -872,7 +666,7 @@ def _raise_if_intersection(filter1: Dict[str, Any], filter2: Dict[str, Any]) -> @configspec class MongoDbCollectionConfiguration(BaseConfiguration): - incremental: Optional[dlt.sources.incremental] = None + incremental: Optional[dlt.sources.incremental] = None # type: ignore[type-arg] @configspec @@ -880,101 +674,10 @@ class MongoDbCollectionResourceConfiguration(BaseConfiguration): connection_url: dlt.TSecretValue = dlt.secrets.value database: Optional[str] = dlt.config.value collection: str = dlt.config.value - incremental: Optional[dlt.sources.incremental] = None + incremental: Optional[dlt.sources.incremental] = None # type: ignore[type-arg] write_disposition: Optional[str] = dlt.config.value parallel: Optional[bool] = False projection: Optional[Union[Mapping[str, Any], Iterable[str]]] = dlt.config.value -def convert_mongo_shell_to_extended_json(query_string: str) -> str: - """ - Convert MongoDB shell syntax to MongoDB Extended JSON v2 format. - - This function handles common MongoDB shell constructs like ISODate, ObjectId, - NumberLong, NumberDecimal, etc. and converts them to their Extended JSON equivalents - that can be parsed by bson.json_util. - - Args: - query_string: A string containing MongoDB shell syntax - - Returns: - A string with MongoDB Extended JSON v2 format - - Examples: - >>> convert_mongo_shell_to_extended_json('ISODate("2010-01-01T00:00:00.000Z")') - '{"$date": "2010-01-01T00:00:00.000Z"}' - - >>> convert_mongo_shell_to_extended_json('ObjectId("507f1f77bcf86cd799439011")') - '{"$oid": "507f1f77bcf86cd799439011"}' - """ - converted = query_string - - # Convert ISODate("...") to {"$date": "..."} - # Pattern matches ISODate("2010-01-01T00:00:00.000+0000") or similar - converted = re.sub(r'ISODate\("([^"]+)"\)', r'{"$date": "\1"}', converted) - - # Convert ObjectId("...") to {"$oid": "..."} - converted = re.sub(r'ObjectId\("([^"]+)"\)', r'{"$oid": "\1"}', converted) - - # Convert NumberLong(...) to {"$numberLong": "..."} - # Note: NumberLong can have quotes or not: NumberLong(123) or NumberLong("123") - converted = re.sub(r'NumberLong\("([^"]+)"\)', r'{"$numberLong": "\1"}', converted) - converted = re.sub(r"NumberLong\(([^)]+)\)", r'{"$numberLong": "\1"}', converted) - - # Convert NumberInt(...) to {"$numberInt": "..."} - converted = re.sub(r'NumberInt\("([^"]+)"\)', r'{"$numberInt": "\1"}', converted) - converted = re.sub(r"NumberInt\(([^)]+)\)", r'{"$numberInt": "\1"}', converted) - - # Convert NumberDecimal("...") to {"$numberDecimal": "..."} - converted = re.sub( - r'NumberDecimal\("([^"]+)"\)', r'{"$numberDecimal": "\1"}', converted - ) - - # Convert Timestamp(..., ...) to {"$timestamp": {"t": ..., "i": ...}} - # Timestamp(1234567890, 1) -> {"$timestamp": {"t": 1234567890, "i": 1}} - converted = re.sub( - r"Timestamp\((\d+),\s*(\d+)\)", r'{"$timestamp": {"t": \1, "i": \2}}', converted - ) - - # Convert BinData(..., "...") to {"$binary": {"base64": "...", "subType": "..."}} - converted = re.sub( - r'BinData\((\d+),\s*"([^"]+)"\)', - r'{"$binary": {"base64": "\2", "subType": "\1"}}', - converted, - ) - - # Convert MinKey() to {"$minKey": 1} - converted = re.sub(r"MinKey\(\)", r'{"$minKey": 1}', converted) - - # Convert MaxKey() to {"$maxKey": 1} - converted = re.sub(r"MaxKey\(\)", r'{"$maxKey": 1}', converted) - - # Convert UUID("...") to {"$uuid": "..."} - converted = re.sub(r'UUID\("([^"]+)"\)', r'{"$uuid": "\1"}', converted) - - # Convert DBRef("collection", "id") to {"$ref": "collection", "$id": "id"} - converted = re.sub( - r'DBRef\("([^"]+)",\s*"([^"]+)"\)', r'{"$ref": "\1", "$id": "\2"}', converted - ) - - # Convert Code("...") to {"$code": "..."} - converted = re.sub(r'Code\("([^"]+)"\)', r'{"$code": "\1"}', converted) - - return converted - - __source_name__ = "mongodb" - - -# MongoDB destination helper functions -def process_file_items(file_path: str) -> list[dict]: - """Process items from a file path (JSONL format).""" - from dlt.common import json - - documents = [] - with open(file_path, "r") as f: - for line in f: - if line.strip(): - doc = json.loads(line.strip()) - documents.append(doc) # Include all fields including DLT metadata - return documents diff --git a/omniload/source/mux/adapter.py b/omniload/source/mux/adapter.py new file mode 100644 index 000000000..6a51a32d9 --- /dev/null +++ b/omniload/source/mux/adapter.py @@ -0,0 +1,102 @@ +# Copyright 2022-2026 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. + +"""Loads Mux views data using https://docs.mux.com/api-reference""" + +from typing import Iterable + +import dlt +from dlt.common import pendulum +from dlt.common.typing import TDataItem +from dlt.sources import DltResource +from dlt.sources.helpers import requests +from requests.auth import HTTPBasicAuth + +from .settings import API_BASE_URL, DEFAULT_LIMIT + + +@dlt.source +def mux_source() -> Iterable[DltResource]: + """ + Source function that passes all video assets and every video view from yesterday to be loaded. + + Yields: + DltResource: Video assets and video views to be loaded. + """ + yield assets_resource + yield views_resource + + +@dlt.resource(write_disposition="merge") +def assets_resource( + mux_api_access_token: str = dlt.secrets.value, + mux_api_secret_key: str = dlt.secrets.value, + limit: int = DEFAULT_LIMIT, +) -> Iterable[TDataItem]: + """ + Resource function that yields metadata about every asset to be loaded. + + Args: + mux_api_access_token (str): API access token for Mux. + mux_api_secret_key (str): API secret key for Mux. + limit (int): Limit on the number of assets to retrieve. Defaults to DEFAULT_LIMIT. + + Yields: + TDataItem: Data of each asset. + """ + url = f"{API_BASE_URL}/video/v1/assets" + params = {"limit": limit} + + response = requests.get( + url, params=params, auth=HTTPBasicAuth(mux_api_access_token, mux_api_secret_key) + ) + response.raise_for_status() + yield response.json()["data"] + + +@dlt.resource(write_disposition="append") +def views_resource( + mux_api_access_token: str = dlt.secrets.value, + mux_api_secret_key: str = dlt.secrets.value, + limit: int = DEFAULT_LIMIT, +) -> Iterable[DltResource]: + """ + Resource function that yields metadata about every video view from yesterday to be loaded. + + Args: + mux_api_access_token (str): API access token for Mux. + mux_api_secret_key (str): API secret key for Mux. + limit (int): Limit on the number of video views to retrieve. Defaults to DEFAULT_LIMIT. + + Yields: + DltResource: Data for each video view from yesterday. + """ + url = f"{API_BASE_URL}/data/v1/video-views" + page = 1 + today = pendulum.today() + yest_start = today.subtract(days=1).int_timestamp + yest_end = today.int_timestamp + + while True: + params = {"limit": limit, "page": page, "timeframe[]": [yest_start, yest_end]} + response = requests.get( + url, + params=params, # type: ignore + auth=HTTPBasicAuth(mux_api_access_token, mux_api_secret_key), + ) + response.raise_for_status() + if response.json()["data"] == []: + break + yield response.json()["data"] + page += 1 diff --git a/omniload/source/mux/settings.py b/omniload/source/mux/settings.py new file mode 100644 index 000000000..8c7180fa1 --- /dev/null +++ b/omniload/source/mux/settings.py @@ -0,0 +1,18 @@ +# Copyright 2022-2026 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. + +"""Mux source settings and constants""" + +API_BASE_URL = "https://api.mux.com" +DEFAULT_LIMIT = 100 diff --git a/omniload/source/notion/adapter.py b/omniload/source/notion/adapter.py index 1dfde1844..a65bdc654 100644 --- a/omniload/source/notion/adapter.py +++ b/omniload/source/notion/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,13 +17,42 @@ from typing import Dict, Iterator, List, Optional import dlt +from dlt.common.typing import TDataItems from dlt.sources import DltResource from .helpers.client import NotionClient from .helpers.database import NotionDatabase -@dlt.source(max_table_nesting=1) +@dlt.resource +def notion_pages( + page_ids: Optional[List[str]] = None, api_key: str = dlt.secrets.value +) -> Iterator[TDataItems]: + """ + Retrieves pages from Notion. + + Args: + page_ids (Optional[List[str]]): A list of page ids. + Defaults to None. If None, the function will generate all pages + in the workspace that are accessible to the integration. + api_key (str): The Notion API secret key. + + Yields: + Iterator[TDataItems]: Pages from Notion. + """ + client = NotionClient(api_key) + pages = client.search(filter_criteria={"value": "page", "property": "object"}) + + for page in pages: + blocks = client.fetch_resource("blocks", page["id"], "children")["results"] + if page_ids and page["id"] not in page_ids: + continue + + if blocks: + yield blocks + + +@dlt.source def notion_databases( database_ids: Optional[List[Dict[str, str]]] = None, api_key: str = dlt.secrets.value, diff --git a/omniload/source/notion/settings.py b/omniload/source/notion/settings.py index 80837f96c..bc5c44044 100644 --- a/omniload/source/notion/settings.py +++ b/omniload/source/notion/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/personio/adapter.py b/omniload/source/personio/adapter.py index 72dcac002..4c4111862 100644 --- a/omniload/source/personio/adapter.py +++ b/omniload/source/personio/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,20 +18,19 @@ import dlt from dlt.common import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import TAnyDateTime, TDataItem from dlt.sources import DltResource from .helpers import PersonioAPI +from .settings import DEFAULT_ITEMS_PER_PAGE, FIRST_DAY_OF_MILLENNIUM -@dlt.source(name="personio", max_table_nesting=0) +@dlt.source(name="personio") def personio_source( - start_date: TAnyDateTime, - end_date: Optional[TAnyDateTime] = None, client_id: str = dlt.secrets.value, client_secret: str = dlt.secrets.value, - items_per_page: int = 200, + items_per_page: int = DEFAULT_ITEMS_PER_PAGE, ) -> Iterable[DltResource]: """ The source for the Personio pipeline. Available resources are employees, absences, and attendances. @@ -46,7 +45,7 @@ def personio_source( client = PersonioAPI(client_id, client_secret) - @dlt.resource(primary_key="id", write_disposition="merge", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="merge") def employees( updated_at: dlt.sources.incremental[ pendulum.DateTime @@ -77,7 +76,7 @@ def convert_item(item: TDataItem) -> TDataItem: name = label.lower() if value["type"] == "date" and value["value"]: - output[name] = ensure_pendulum_datetime_utc(value["value"]) + output[name] = ensure_pendulum_datetime(value["value"]) else: output[name] = value["value"] return output @@ -93,7 +92,7 @@ def convert_item(item: TDataItem) -> TDataItem: for page in pages: yield [convert_item(item) for item in page] - @dlt.resource(primary_key="id", write_disposition="replace", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="replace") def absence_types(items_per_page: int = items_per_page) -> Iterable[TDataItem]: """ The resource for absence types (time-off-types), supports pagination. @@ -112,7 +111,7 @@ def absence_types(items_per_page: int = items_per_page) -> Iterable[TDataItem]: for page in pages: yield [item.get("attributes", {}) for item in page] - @dlt.resource(primary_key="id", write_disposition="merge", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="merge") def absences( updated_at: dlt.sources.incremental[ pendulum.DateTime @@ -143,8 +142,8 @@ def absences( def convert_item(item: TDataItem) -> TDataItem: output = item.get("attributes", {}) - output["created_at"] = ensure_pendulum_datetime_utc(output["created_at"]) - output["updated_at"] = ensure_pendulum_datetime_utc(output["updated_at"]) + output["created_at"] = ensure_pendulum_datetime(output["created_at"]) + output["updated_at"] = ensure_pendulum_datetime(output["updated_at"]) return output pages = client.get_pages( @@ -156,10 +155,10 @@ def convert_item(item: TDataItem) -> TDataItem: for page in pages: yield [convert_item(item) for item in page] - @dlt.resource(primary_key="id", write_disposition="merge", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="merge") def attendances( - start_date: TAnyDateTime = start_date, - end_date: Optional[TAnyDateTime] = end_date, + start_date: TAnyDateTime = FIRST_DAY_OF_MILLENNIUM, + end_date: Optional[TAnyDateTime] = None, updated_at: dlt.sources.incremental[ pendulum.DateTime ] = dlt.sources.incremental( @@ -188,8 +187,8 @@ def attendances( params = { "limit": items_per_page, - "start_date": ensure_pendulum_datetime_utc(start_date).to_date_string(), - "end_date": ensure_pendulum_datetime_utc(end_date).to_date_string(), + "start_date": ensure_pendulum_datetime(start_date).to_date_string(), + "end_date": ensure_pendulum_datetime(end_date).to_date_string(), "updated_from": updated_iso, "includePending": True, } @@ -201,14 +200,14 @@ def attendances( def convert_item(item: TDataItem) -> TDataItem: """Converts an attendance item.""" output = dict(id=item["id"], **item.get("attributes")) - output["date"] = ensure_pendulum_datetime_utc(output["date"]).date() - output["updated_at"] = ensure_pendulum_datetime_utc(output["updated_at"]) + output["date"] = ensure_pendulum_datetime(output["date"]).date() + output["updated_at"] = ensure_pendulum_datetime(output["updated_at"]) return output for page in pages: yield [convert_item(item) for item in page] - @dlt.resource(primary_key="id", write_disposition="replace", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="replace") def projects() -> Iterable[TDataItem]: """ The resource for projects. @@ -222,14 +221,14 @@ def projects() -> Iterable[TDataItem]: def convert_item(item: TDataItem) -> TDataItem: """Converts an attendance item.""" output = dict(id=item["id"], **item.get("attributes")) - output["created_at"] = ensure_pendulum_datetime_utc(output["created_at"]) - output["updated_at"] = ensure_pendulum_datetime_utc(output["updated_at"]) + output["created_at"] = ensure_pendulum_datetime(output["created_at"]) + output["updated_at"] = ensure_pendulum_datetime(output["updated_at"]) return output for page in pages: yield [convert_item(item) for item in page] - @dlt.resource(primary_key="id", write_disposition="replace", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="replace") def document_categories() -> Iterable[TDataItem]: """ The resource for document_categories. @@ -248,7 +247,7 @@ def convert_item(item: TDataItem) -> TDataItem: for page in pages: yield [convert_item(item) for item in page] - @dlt.resource(primary_key="id", write_disposition="replace", max_table_nesting=0) + @dlt.resource(primary_key="id", write_disposition="replace") def custom_reports_list() -> Iterable[TDataItem]: """ The resource for custom_reports. @@ -314,7 +313,7 @@ def convert_item(item: TDataItem, report_id: str) -> TDataItem: for value in attributes: name = value["attribute_id"] if value["data_type"] == "date" and value["value"]: - output[name] = ensure_pendulum_datetime_utc(value["value"]) + output[name] = ensure_pendulum_datetime(value["value"]) else: output[name] = value["value"] return output diff --git a/omniload/source/personio/helpers.py b/omniload/source/personio/helpers.py index f5384902d..99bf8e632 100644 --- a/omniload/source/personio/helpers.py +++ b/omniload/source/personio/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/personio/settings.py b/omniload/source/personio/settings.py new file mode 100644 index 000000000..42cc2f2c4 --- /dev/null +++ b/omniload/source/personio/settings.py @@ -0,0 +1,16 @@ +# Copyright 2022-2026 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. + +DEFAULT_ITEMS_PER_PAGE = 200 +FIRST_DAY_OF_MILLENNIUM = "2000-01-01" diff --git a/omniload/source/pg_replication/adapter.py b/omniload/source/pg_replication/adapter.py new file mode 100644 index 000000000..bafd89266 --- /dev/null +++ b/omniload/source/pg_replication/adapter.py @@ -0,0 +1,119 @@ +# Copyright 2022-2026 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. + +"""Replicates postgres tables in batch using logical decoding.""" + +from typing import Dict, Iterable, Optional, Sequence, Union + +import dlt +from dlt.common import logger +from dlt.common.schema.typing import TTableSchemaColumns +from dlt.common.typing import TDataItem +from dlt.extract.items import DataItemWithMeta +from dlt.sources.credentials import ConnectionStringCredentials + +from .helpers import ItemGenerator, advance_slot, get_max_lsn + + +@dlt.resource( + name=lambda args: args["slot_name"], + standalone=True, +) +def replication_resource( + slot_name: str, + pub_name: str, + credentials: ConnectionStringCredentials = dlt.secrets.value, + include_columns: Optional[Dict[str, Sequence[str]]] = None, + columns: Optional[Dict[str, TTableSchemaColumns]] = None, + target_batch_size: int = 1000, + flush_slot: bool = True, +) -> Iterable[Union[TDataItem, DataItemWithMeta]]: + """Resource yielding data items for changes in one or more postgres tables. + + - Relies on a replication slot and publication that publishes DML operations + (i.e. `insert`, `update`, and/or `delete`). Helper `init_replication` can be + used to set this up. + - Maintains LSN of last consumed message in state to track progress. + - At start of the run, advances the slot upto last consumed message in previous run. + - Processes in batches to limit memory usage. + + Args: + slot_name (str): Name of the replication slot to consume replication messages from. + pub_name (str): Name of the publication that publishes DML operations for the table(s). + credentials (ConnectionStringCredentials): Postgres database credentials. + include_columns (Optional[Dict[str, Sequence[str]]]): Maps table name(s) to + sequence of names of columns to include in the generated data items. + Any column not in the sequence is excluded. If not provided, all columns + are included. For example: + ``` + include_columns={ + "table_x": ["col_a", "col_c"], + "table_y": ["col_x", "col_y", "col_z"], + } + ``` + columns (Optional[Dict[str, TTableHintTemplate[TAnySchemaColumns]]]): Maps + table name(s) to column hints to apply on the replicated table(s). For example: + ``` + columns={ + "table_x": {"col_a": {"data_type": "json"}}, + "table_y": {"col_y": {"precision": 32}}, + } + ``` + target_batch_size (int): Desired number of data items yielded in a batch. + Can be used to limit the data items in memory. Note that the number of + data items yielded can be (far) greater than `target_batch_size`, because + all messages belonging to the same transaction are always processed in + the same batch, regardless of the number of messages in the transaction + and regardless of the value of `target_batch_size`. The number of data + items can also be smaller than `target_batch_size` when the replication + slot is exhausted before a batch is full. + flush_slot (bool): Whether processed messages are discarded from the replication + slot. Recommended value is True. Be careful when setting False—not flushing + can eventually lead to a “disk full” condition on the server, because + the server retains all the WAL segments that might be needed to stream + the changes via all of the currently open replication slots. + + Yields: + Data items for changes published in the publication. + """ + # start where we left off in previous run + start_lsn = dlt.current.resource_state().get("last_commit_lsn", 0) + if flush_slot and start_lsn: + advance_slot(start_lsn, slot_name, credentials) + + # continue until last message in replication slot + options = {"publication_names": pub_name, "proto_version": "1"} + upto_lsn = get_max_lsn(slot_name, options, credentials) + if upto_lsn is None: + return + logger.info( + f"Replicating slot {slot_name} publication {pub_name} from {start_lsn} to {upto_lsn}" + ) + # generate items in batches + while True: + gen = ItemGenerator( + credentials=credentials, + slot_name=slot_name, + options=options, + upto_lsn=upto_lsn, + start_lsn=start_lsn, + target_batch_size=target_batch_size, + include_columns=include_columns, + columns=columns, + ) + yield from gen + if gen.generated_all: + dlt.current.resource_state()["last_commit_lsn"] = gen.last_commit_lsn + break + start_lsn = gen.last_commit_lsn diff --git a/omniload/source/pg_replication/decoders.py b/omniload/source/pg_replication/decoders.py new file mode 100644 index 000000000..209fcf2b4 --- /dev/null +++ b/omniload/source/pg_replication/decoders.py @@ -0,0 +1,442 @@ +# Copyright 2022-2026 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. + +# flake8: noqa +# file copied from https://raw.githubusercontent.com/dgea005/pypgoutput/master/src/pypgoutput/decoders.py +# we do this instead of importing `pypgoutput` because it depends on `psycopg2`, which causes errors when installing on macOS + +import io +from abc import ABC, abstractmethod +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import List, NamedTuple, Optional, Union + +# integer byte lengths +INT8 = 1 +INT16 = 2 +INT32 = 4 +INT64 = 8 + + +def convert_pg_ts(_ts_in_microseconds: int) -> datetime: + ts = datetime(2000, 1, 1, 0, 0, 0, 0, tzinfo=timezone.utc) + return ts + timedelta(microseconds=_ts_in_microseconds) + + +def convert_bytes_to_int(_in_bytes: bytes) -> int: + return int.from_bytes(_in_bytes, byteorder="big", signed=True) + + +def convert_bytes_to_utf8(_in_bytes: Union[bytes, bytearray]) -> str: + return (_in_bytes).decode("utf-8") + + +class ColumnData(NamedTuple): + # col_data_category is NOT the type. it means null value/toasted(not sent)/text formatted + col_data_category: Optional[str] + col_data_length: Optional[int] = None + col_data: Optional[str] = None + + def __repr__(self) -> str: + return f"[col_data_category='{self.col_data_category}', col_data_length={self.col_data_length}, col_data='{self.col_data}']" + + +class ColumnType(NamedTuple): + """https://www.postgresql.org/docs/12/catalog-pg-attribute.html""" + + part_of_pkey: int + name: str + type_id: int + atttypmod: int + + +class TupleData(NamedTuple): + n_columns: int + column_data: List[ColumnData] + + def __repr__(self) -> str: + return f"n_columns: {self.n_columns}, data: {self.column_data}" + + +# TODO: you can make decoding way faster by +# - moving all the decoding core to PgoutputMessage +# - use struct unpack and increase offset manually to reduce calls +# - use tuples to represent data, separate data from decoding! +class PgoutputMessage(ABC): + def __init__(self, buffer: bytes): + self.buffer: io.BytesIO = io.BytesIO(buffer) + self.byte1: str = self.read_utf8(1) + self.decode_buffer() + + @abstractmethod + def decode_buffer(self) -> None: + """Decoding is implemented for each message type""" + + @abstractmethod + def __repr__(self) -> str: + """Implemented for each message type""" + + def read_int8(self) -> int: + return convert_bytes_to_int(self.buffer.read(INT8)) + + def read_int16(self) -> int: + return convert_bytes_to_int(self.buffer.read(INT16)) + + def read_int32(self) -> int: + return convert_bytes_to_int(self.buffer.read(INT32)) + + def read_int64(self) -> int: + return convert_bytes_to_int(self.buffer.read(INT64)) + + def read_utf8(self, n: int = 1) -> str: + return convert_bytes_to_utf8(self.buffer.read(n)) + + def read_timestamp(self) -> datetime: + # 8 chars -> int64 -> timestamp + return convert_pg_ts(_ts_in_microseconds=self.read_int64()) + + def read_string(self) -> str: + output = bytearray() + while (next_char := self.buffer.read(1)) != b"\x00": + output += next_char + return convert_bytes_to_utf8(output) + + def read_tuple_data(self) -> TupleData: + """ + TupleData + Int16 Number of columns. + Next, one of the following submessages appears for each column (except generated columns): + Byte1('n') Identifies the data as NULL value. + Or + Byte1('u') Identifies unchanged TOASTed value (the actual value is not sent). + Or + Byte1('t') Identifies the data as text formatted value. + Int32 Length of the column value. + Byten The value of the column, in text format. (A future release might support additional formats.) n is the above length. + """ + # TODO: investigate what happens with the generated columns + column_data = list() + n_columns = self.read_int16() + for column in range(n_columns): + col_data_category = self.read_utf8() + if col_data_category in ("n", "u"): + # "n"=NULL, "t"=TOASTed + column_data.append(ColumnData(col_data_category=col_data_category)) + elif col_data_category == "t": + # t = tuple + col_data_length = self.read_int32() + col_data = self.read_utf8(col_data_length) + column_data.append( + ColumnData( + col_data_category=col_data_category, + col_data_length=col_data_length, + col_data=col_data, + ) + ) + return TupleData(n_columns=n_columns, column_data=column_data) + + +class Begin(PgoutputMessage): + """ + https://pgpedia.info/x/xlogrecptr.html + https://www.postgresql.org/docs/14/datatype-pg-lsn.html + + byte1 Byte1('B') Identifies the message as a begin message. + lsn Int64 The final LSN of the transaction. + commit_tx_ts Int64 Commit timestamp of the transaction. The value is in number of microseconds since PostgreSQL epoch (2000-01-01). + tx_xid Int32 Xid of the transaction. + """ + + byte1: str + lsn: int + commit_ts: datetime + tx_xid: int + + def decode_buffer(self) -> None: + if self.byte1 != "B": + raise ValueError("first byte in buffer does not match Begin message") + self.lsn = self.read_int64() + self.commit_ts = self.read_timestamp() + self.tx_xid = self.read_int64() + + def __repr__(self) -> str: + return ( + f"BEGIN \n\tbyte1: '{self.byte1}', \n\tLSN: {self.lsn}, " + f"\n\tcommit_ts {self.commit_ts}, \n\ttx_xid: {self.tx_xid}" + ) + + +class Commit(PgoutputMessage): + """ + byte1: Byte1('C') Identifies the message as a commit message. + flags: Int8 Flags; currently unused (must be 0). + lsn_commit: Int64 The LSN of the commit. + lsn: Int64 The end LSN of the transaction. + Int64 Commit timestamp of the transaction. The value is in number of microseconds since PostgreSQL epoch (2000-01-01). + """ + + byte1: str + flags: int + lsn_commit: int + lsn: int + commit_ts: datetime + + def decode_buffer(self) -> None: + if self.byte1 != "C": + raise ValueError("first byte in buffer does not match Commit message") + self.flags = self.read_int8() + self.lsn_commit = self.read_int64() + self.lsn = self.read_int64() + self.commit_ts = self.read_timestamp() + + def __repr__(self) -> str: + return ( + f"COMMIT \n\tbyte1: {self.byte1}, \n\tflags {self.flags}, \n\tlsn_commit: {self.lsn_commit}" + f"\n\tLSN: {self.lsn}, \n\tcommit_ts {self.commit_ts}" + ) + + +class Origin: + """ + Byte1('O') Identifies the message as an origin message. + Int64 The LSN of the commit on the origin server. + String Name of the origin. + Note that there can be multiple Origin messages inside a single transaction. + This seems to be what origin means: https://www.postgresql.org/docs/12/replication-origins.html + """ + + pass + + +class Relation(PgoutputMessage): + """ + Byte1('R') Identifies the message as a relation message. + Int32 ID of the relation. + String Namespace (empty string for pg_catalog). + String Relation name. + Int8 Replica identity setting for the relation (same as relreplident in pg_class). + # select relreplident from pg_class where relname = 'test_table'; + # from reading the documentation and looking at the tables this is not int8 but a single character + # background: https://www.postgresql.org/docs/10/sql-altertable.html#SQL-CREATETABLE-REPLICA-IDENTITY + Int16 Number of columns. + Next, the following message part appears for each column (except generated columns): + Int8 Flags for the column. Currently can be either 0 for no flags or 1 which marks the column as part of the key. + String Name of the column. + Int32 ID of the column's data type. + Int32 Type modifier of the column (atttypmod). + """ + + byte1: str + relation_id: int + namespace: str + relation_name: str + replica_identity_setting: str + n_columns: int + columns: List[ColumnType] + + def decode_buffer(self) -> None: + if self.byte1 != "R": + raise ValueError("first byte in buffer does not match Relation message") + self.relation_id = self.read_int32() + self.namespace = self.read_string() + self.relation_name = self.read_string() + self.replica_identity_setting = self.read_utf8() + self.n_columns = self.read_int16() + self.columns = list() + + for column in range(self.n_columns): + part_of_pkey = self.read_int8() + col_name = self.read_string() + data_type_id = self.read_int32() + # TODO: check on use of signed / unsigned + # check with select oid from pg_type where typname = ; timestamp == 1184, int4 = 23 + col_modifier = self.read_int32() + self.columns.append( + ColumnType( + part_of_pkey=part_of_pkey, + name=col_name, + type_id=data_type_id, + atttypmod=col_modifier, + ) + ) + + def __repr__(self) -> str: + return ( + f"RELATION \n\tbyte1: '{self.byte1}', \n\trelation_id: {self.relation_id}" + f",\n\tnamespace/schema: '{self.namespace}',\n\trelation_name: '{self.relation_name}'" + f",\n\treplica_identity_setting: '{self.replica_identity_setting}',\n\tn_columns: {self.n_columns} " + f",\n\tcolumns: {self.columns}" + ) + + +class PgType: + """ + Renamed to PgType not to collide with "type" + + Byte1('Y') Identifies the message as a type message. + Int32 ID of the data type. + String Namespace (empty string for pg_catalog). + String Name of the data type. + """ + + pass + + +class Insert(PgoutputMessage): + """ + Byte1('I') Identifies the message as an insert message. + Int32 ID of the relation corresponding to the ID in the relation message. + Byte1('N') Identifies the following TupleData message as a new tuple. + TupleData TupleData message part representing the contents of new tuple. + """ + + byte1: str + relation_id: int + new_tuple_byte: str + new_tuple: TupleData + + def decode_buffer(self) -> None: + if self.byte1 != "I": + raise ValueError( + f"first byte in buffer does not match Insert message (expected 'I', got '{self.byte1}'" + ) + self.relation_id = self.read_int32() + self.new_tuple_byte = self.read_utf8() + self.new_tuple = self.read_tuple_data() + + def __repr__(self) -> str: + return ( + f"INSERT \n\tbyte1: '{self.byte1}', \n\trelation_id: {self.relation_id} " + f"\n\tnew tuple byte: '{self.new_tuple_byte}', \n\tnew_tuple: {self.new_tuple}" + ) + + +class Update(PgoutputMessage): + """ + Byte1('U') Identifies the message as an update message. + Int32 ID of the relation corresponding to the ID in the relation message. + Byte1('K') Identifies the following TupleData submessage as a key. This field is optional and is only present if the update changed data in any of the column(s) that are part of the REPLICA IDENTITY index. + Byte1('O') Identifies the following TupleData submessage as an old tuple. This field is optional and is only present if table in which the update happened has REPLICA IDENTITY set to FULL. + TupleData TupleData message part representing the contents of the old tuple or primary key. Only present if the previous 'O' or 'K' part is present. + Byte1('N') Identifies the following TupleData message as a new tuple. + TupleData TupleData message part representing the contents of a new tuple. + + The Update message may contain either a 'K' message part or an 'O' message part or neither of them, but never both of them. + """ + + byte1: str + relation_id: int + next_byte_identifier: Optional[str] + optional_tuple_identifier: Optional[str] + old_tuple: Optional[TupleData] + new_tuple_byte: str + new_tuple: TupleData + + def decode_buffer(self) -> None: + self.optional_tuple_identifier = None + self.old_tuple = None + if self.byte1 != "U": + raise ValueError( + f"first byte in buffer does not match Update message (expected 'U', got '{self.byte1}'" + ) + self.relation_id = self.read_int32() + # TODO test update to PK, test update with REPLICA IDENTITY = FULL + self.next_byte_identifier = self.read_utf8() # one of K, O or N + if self.next_byte_identifier == "K" or self.next_byte_identifier == "O": + self.optional_tuple_identifier = self.next_byte_identifier + self.old_tuple = self.read_tuple_data() + self.new_tuple_byte = self.read_utf8() + else: + self.new_tuple_byte = self.next_byte_identifier + if self.new_tuple_byte != "N": + # TODO: test exception handling + raise ValueError( + f"did not find new_tuple_byte ('N') at position: {self.buffer.tell()}, found: '{self.new_tuple_byte}'" + ) + self.new_tuple = self.read_tuple_data() + + def __repr__(self) -> str: + return ( + f"UPDATE \n\tbyte1: '{self.byte1}', \n\trelation_id: {self.relation_id}" + f"\n\toptional_tuple_identifier: '{self.optional_tuple_identifier}', \n\toptional_old_tuple_data: {self.old_tuple}" + f"\n\tnew_tuple_byte: '{self.new_tuple_byte}', \n\tnew_tuple: {self.new_tuple}" + ) + + +class Delete(PgoutputMessage): + """ + Byte1('D') Identifies the message as a delete message. + Int32 ID of the relation corresponding to the ID in the relation message. + Byte1('K') Identifies the following TupleData submessage as a key. This field is present if the table in which the delete has happened uses an index as REPLICA IDENTITY. + Byte1('O') Identifies the following TupleData message as a old tuple. This field is present if the table in which the delete has happened has REPLICA IDENTITY set to FULL. + TupleData TupleData message part representing the contents of the old tuple or primary key, depending on the previous field. + + The Delete message may contain either a 'K' message part or an 'O' message part, but never both of them. + """ + + byte1: str + relation_id: int + message_type: str + old_tuple: TupleData + + def decode_buffer(self) -> None: + if self.byte1 != "D": + raise ValueError( + f"first byte in buffer does not match Delete message (expected 'D', got '{self.byte1}'" + ) + self.relation_id = self.read_int32() + self.message_type = self.read_utf8() + # TODO: test with replica identity full + if self.message_type not in ["K", "O"]: + raise ValueError( + f"message type byte is not 'K' or 'O', got: '{self.message_type}'" + ) + self.old_tuple = self.read_tuple_data() + + def __repr__(self) -> str: + return ( + f"DELETE \n\tbyte1: {self.byte1} \n\trelation_id: {self.relation_id} " + f"\n\tmessage_type: {self.message_type} \n\told_tuple: {self.old_tuple}" + ) + + +class Truncate(PgoutputMessage): + """ + Byte1('T') Identifies the message as a truncate message. + Int32 Number of relations + Int8 Option bits for TRUNCATE: 1 for CASCADE, 2 for RESTART IDENTITY + Int32 ID of the relation corresponding to the ID in the relation message. This field is repeated for each relation. + """ + + byte1: str + number_of_relations: int + option_bits: int + relation_ids: List[int] + + def decode_buffer(self) -> None: + if self.byte1 != "T": + raise ValueError( + f"first byte in buffer does not match Truncate message (expected 'T', got '{self.byte1}'" + ) + self.number_of_relations = self.read_int32() + self.option_bits = self.read_int8() + self.relation_ids = [] + for relation in range(self.number_of_relations): + self.relation_ids.append(self.read_int32()) + + def __repr__(self) -> str: + return ( + f"TRUNCATE \n\tbyte1: {self.byte1} \n\tn_relations: {self.number_of_relations} " + f"option_bits: {self.option_bits}, relation_ids: {self.relation_ids}" + ) diff --git a/omniload/source/pg_replication/exceptions.py b/omniload/source/pg_replication/exceptions.py new file mode 100644 index 000000000..11c221435 --- /dev/null +++ b/omniload/source/pg_replication/exceptions.py @@ -0,0 +1,21 @@ +# Copyright 2022-2026 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. + + +class NoPrimaryKeyException(Exception): + pass + + +class IncompatiblePostgresVersionException(Exception): + pass diff --git a/omniload/source/pg_replication/helpers.py b/omniload/source/pg_replication/helpers.py new file mode 100644 index 000000000..2bb1e7c4f --- /dev/null +++ b/omniload/source/pg_replication/helpers.py @@ -0,0 +1,855 @@ +# Copyright 2022-2026 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 dataclasses import dataclass, field +from typing import ( + Any, + Dict, + Iterator, + List, + Mapping, + Optional, + Sequence, + Union, +) + +import dlt +import psycopg2 +from dlt.common import logger +from dlt.common.configuration.specs.config_section_context import ConfigSectionContext +from dlt.common.data_writers.escape import escape_postgres_identifier +from dlt.common.normalizers.naming.snake_case import NamingConvention +from dlt.common.pendulum import pendulum +from dlt.common.schema.typing import ( + DLT_NAME_PREFIX, + TColumnNames, + TTableSchema, + TTableSchemaColumns, + TWriteDisposition, +) +from dlt.common.schema.utils import merge_column +from dlt.common.typing import TDataItem +from dlt.extract import DltResource +from dlt.extract.items import DataItemWithMeta +from dlt.sources.config import with_config +from dlt.sources.credentials import ConnectionStringCredentials +from dlt.sources.sql_database import ( + sql_database as core_sql_database, +) +from dlt.sources.sql_database import ( + sql_table as core_sql_table, +) +from psycopg2.extensions import connection as ConnectionExt +from psycopg2.extensions import cursor +from psycopg2.extras import ( + LogicalReplicationConnection, + ReplicationCursor, + ReplicationMessage, + StopReplication, +) + +from .decoders import ( + Begin, + ColumnData, + Delete, + Insert, + Relation, + Update, +) +from .exceptions import IncompatiblePostgresVersionException +from .schema_types import _to_dlt_column_schema, _to_dlt_val + +__source_name__ = "pg_replication" + + +@with_config( + sections=("sources", "pg_replication"), + sections_merge_style=ConfigSectionContext.resource_merge_style, + section_arg_name="slot_name", +) +def replication_connection( + slot_name: str, credentials: ConnectionStringCredentials = dlt.secrets.value +) -> LogicalReplicationConnection: + """Returns a psycopg2 LogicalReplicationConnection to interact with postgres replication functionality. + + Requires `slot_name` to be passed in order to generate compatible configuraiton section. + + Raises error if the user does not have the REPLICATION attribute assigned. + """ + return _get_rep_conn(credentials) + + +@with_config( + sections=("sources", "pg_replication"), + sections_merge_style=ConfigSectionContext.resource_merge_style, + section_arg_name="slot_name", +) +def init_replication( + slot_name: str = dlt.config.value, + pub_name: str = dlt.config.value, + schema_name: str = dlt.config.value, + table_names: Optional[Union[str, Sequence[str]]] = dlt.config.value, + credentials: ConnectionStringCredentials = dlt.secrets.value, + publish: str = "insert, update, delete", + persist_snapshots: bool = False, + include_columns: Optional[Mapping[str, Sequence[str]]] = None, + columns: Optional[Mapping[str, TTableSchemaColumns]] = None, + reset: bool = False, +) -> Optional[Union[DltResource, List[DltResource]]]: + """Initializes replication for one, several, or all tables within a schema. + + Can be called repeatedly with the same `slot_name` and `pub_name`: + - creates a replication slot and publication with provided names if they do not exist yet + - skips creation of slot and publication if they already exist (unless`reset` is set to `False`) + - supports addition of new tables by extending `table_names` + - removing tables is not supported, i.e. exluding a table from `table_names` + will not remove it from the publication + - switching from a table selection to an entire schema is possible by omitting + the `table_names` argument + - changing `publish` has no effect (altering the published DML operations is not supported) + - table snapshots can only be persisted on the first call (because the snapshot + is exported when the slot is created) + + Args: + slot_name (str): Name of the replication slot to create if it does not exist yet. + pub_name (str): Name of the publication to create if it does not exist yet. + schema_name (str): Name of the schema to replicate tables from. + table_names (Optional[Union[str, Sequence[str]]]): Name(s) of the table(s) + to include in the publication. If not provided, the whole schema with `schema_name` will be replicated + (also tables added to the schema after the publication was created). You need superuser privileges + for the whole schema replication. When specifying individual table names, the database role must + own the tables if it's not a superuser. + credentials (ConnectionStringCredentials): Postgres database credentials. + publish (str): Comma-separated string of DML operations. Can be used to + control which changes are included in the publication. Allowed operations + are `insert`, `update`, and `delete`. `truncate` is currently not + supported—messages of that type are ignored. + E.g. `publish="insert"` will create a publication that only publishes insert operations. + persist_snapshots (bool): Whether the table states in the snapshot exported + during replication slot creation are persisted to tables. If true, a + snapshot table is created in Postgres for all included tables, and corresponding + resources (`DltResource` objects) for these tables are created and returned. + The resources can be used to perform an initial load of all data present + in the tables at the moment the replication slot got created. + include_columns (Optional[Dict[str, Sequence[str]]]): Maps table name(s) to + sequence of names of columns to include in the snapshot table(s). + Any column not in the sequence is excluded. If not provided, all columns + are included. For example: + ``` + include_columns={ + "table_x": ["col_a", "col_c"], + "table_y": ["col_x", "col_y", "col_z"], + } + ``` + Argument is only used if `persist_snapshots` is `True`. + columns (Optional[Dict[str, TTableSchemaColumns]]): Maps + table name(s) to column hints to apply on the snapshot table resource(s). + For example: + ``` + columns={ + "table_x": {"col_a": {"data_type": "json"}}, + "table_y": {"col_y": {"precision": 32}}, + } + ``` + Argument is only used if `persist_snapshots` is `True`. + reset (bool): If set to True, the existing slot and publication are dropped + and recreated. Has no effect if a slot and publication with the provided + names do not yet exist. + + Returns: + - None if `persist_snapshots` is `False` + - a `DltResource` object or a list of `DltResource` objects for the snapshot + table(s) if `persist_snapshots` is `True` and the replication slot did not yet exist + """ + if isinstance(table_names, str): + table_names = [table_names] + + rep_conn = _get_rep_conn(credentials) + try: + with rep_conn.cursor() as cur: + if reset: + drop_replication_slot(slot_name, cur) + drop_publication(pub_name, cur) + create_publication(pub_name, cur, publish) + if table_names is None: + add_schema_to_publication(schema_name, pub_name, cur) + else: + add_tables_to_publication(table_names, schema_name, pub_name, cur) + + slot = create_replication_slot(slot_name, cur) + if persist_snapshots: + if slot is None: + raise RuntimeError( + f"Cannot create snapshots because slot {slot_name} is already created." + ) + + # get list of tables via sql_database if not provided + if table_names is None: + # do not include dlt tables + table_names = [ + table_name + for table_name in core_sql_database( + credentials, schema=schema_name, reflection_level="minimal" + ).resources.keys() + if not table_name.lower().startswith(DLT_NAME_PREFIX) + ] + + # need separate session to read the snapshot: https://stackoverflow.com/q/75852587 + with _get_conn(credentials) as conn: + with conn.cursor() as cur_snap: + # set the snapshot to the snaphost of the newly created replication slot + cur_snap.execute( + f""" + START TRANSACTION ISOLATION LEVEL REPEATABLE READ; + SET TRANSACTION SNAPSHOT '{slot["snapshot_name"]}'; + """ + ) + snapshot_tables = [ + ( + table_name, + persist_snapshot_table( + snapshot_name=slot["snapshot_name"], + table_name=table_name, + schema_name=schema_name, + cur=cur_snap, + include_columns=( + None + if include_columns is None + else include_columns.get(table_name) + ), + ), + _get_pk(cur_snap, table_name, schema_name), + ) + for table_name in table_names + ] + # commit all tables before creating resources that will reflect them + conn.commit() + snapshot_table_resources = [ + snapshot_table_resource( + snapshot_table_name=snapshot_table_name, + schema_name=schema_name, + table_name=table_name, + write_disposition="append" if publish == "insert" else "merge", + columns=None if columns is None else columns.get(table_name), + primary_key=primary_key, + credentials=credentials, + ) + for table_name, snapshot_table_name, primary_key in snapshot_tables + ] + if len(snapshot_table_resources) == 1: + return snapshot_table_resources[0] + return snapshot_table_resources + except Exception: + rep_conn.rollback() + raise + else: + rep_conn.commit() + finally: + rep_conn.close() + return None + + +@dlt.sources.config.with_config(sections=("sources", "pg_replication")) +def get_pg_version( + cur: cursor = None, + credentials: ConnectionStringCredentials = dlt.secrets.value, +) -> int: + """Returns Postgres server version as int.""" + if cur is not None: + return cur.connection.server_version + with _get_conn(credentials) as conn: + return conn.server_version + + +def create_publication( + name: str, + cur: cursor, + publish: str = "insert, update, delete", +) -> None: + """Creates a publication for logical replication if it doesn't exist yet. + + Does nothing if the publication already exists. + Raises error if the user does not have the CREATE privilege for the database. + """ + esc_name = escape_postgres_identifier(name) + try: + cur.execute(f"CREATE PUBLICATION {esc_name} WITH (publish = '{publish}');") + logger.info( + f"Successfully created publication {esc_name} with publish = '{publish}'." + ) + except psycopg2.errors.DuplicateObject: # the publication already exists + logger.info(f'Publication "{name}" already exists.') + + +def add_table_to_publication( + table_name: str, + schema_name: str, + pub_name: str, + cur: cursor, +) -> None: + """Adds a table to a publication for logical replication. + + Does nothing if the table is already a member of the publication. + Raises error if the user is not owner of the table. + """ + qual_name = _make_qualified_table_name(table_name, schema_name) + esc_pub_name = escape_postgres_identifier(pub_name) + try: + cur.execute(f"ALTER PUBLICATION {esc_pub_name} ADD TABLE {qual_name};") + logger.info( + f"Successfully added table {qual_name} to publication {esc_pub_name}." + ) + except psycopg2.errors.DuplicateObject: + logger.info( + f"Table {qual_name} is already a member of publication {esc_pub_name}." + ) + + +def add_tables_to_publication( + table_names: Union[str, Sequence[str]], + schema_name: str, + pub_name: str, + cur: cursor, +) -> None: + """Adds one or multiple tables to a publication for logical replication. + + Calls `add_table_to_publication` for each table in `table_names`. + """ + if isinstance(table_names, str): + table_names = table_names + for table_name in table_names: + add_table_to_publication(table_name, schema_name, pub_name, cur) + + +def add_schema_to_publication( + schema_name: str, + pub_name: str, + cur: cursor, +) -> None: + """Adds a schema to a publication for logical replication if the schema is not a member yet. + + Raises error if the user is not a superuser. + """ + if (version := get_pg_version(cur)) < 150000: + raise IncompatiblePostgresVersionException( + f"Cannot add schema to publication because the Postgres server version {version} is too low." + " Adding schemas to a publication is only supported for Postgres version 15 or higher." + " Upgrade your Postgres server version or set the `table_names` argument to explicitly specify table names." + ) + esc_schema_name = escape_postgres_identifier(schema_name) + esc_pub_name = escape_postgres_identifier(pub_name) + try: + cur.execute( + f"ALTER PUBLICATION {esc_pub_name} ADD TABLES IN SCHEMA {esc_schema_name};" + ) + logger.info( + f"Successfully added schema {esc_schema_name} to publication {esc_pub_name}." + ) + except psycopg2.errors.DuplicateObject: + logger.info( + f"Schema {esc_schema_name} is already a member of publication {esc_pub_name}." + ) + + +def create_replication_slot( # type: ignore[return] + name: str, cur: ReplicationCursor, output_plugin: str = "pgoutput" +) -> Optional[Dict[str, str]]: + """Creates a replication slot if it doesn't exist yet.""" + try: + cur.create_replication_slot(name, output_plugin=output_plugin) + logger.info(f'Successfully created replication slot "{name}".') + result = cur.fetchone() + return { + "slot_name": result[0], + "consistent_point": result[1], + "snapshot_name": result[2], + "output_plugin": result[3], + } + except psycopg2.errors.DuplicateObject: # the replication slot already exists + logger.info( + f'Replication slot "{name}" cannot be created because it already exists.' + ) + + +def drop_replication_slot(name: str, cur: ReplicationCursor) -> None: + """Drops a replication slot if it exists.""" + try: + cur.drop_replication_slot(name) + logger.info(f'Successfully dropped replication slot "{name}".') + except psycopg2.errors.UndefinedObject: # the replication slot does not exist + logger.info( + f'Replication slot "{name}" cannot be dropped because it does not exist.' + ) + + +def drop_publication(name: str, cur: ReplicationCursor) -> None: + """Drops a publication if it exists.""" + esc_name = escape_postgres_identifier(name) + try: + cur.execute(f"DROP PUBLICATION {esc_name};") + logger.info(f"Successfully dropped publication {esc_name}.") + except psycopg2.errors.UndefinedObject: # the publication does not exist + logger.info( + f"Publication {esc_name} cannot be dropped because it does not exist." + ) + + +def persist_snapshot_table( + snapshot_name: str, + table_name: str, + schema_name: str, + cur: cursor, + include_columns: Optional[Sequence[str]] = None, +) -> str: + """Persists exported snapshot table state. + + Reads snapshot table content and copies it into new table. + """ + col_str = "*" + if include_columns is not None: + col_str = ", ".join(map(escape_postgres_identifier, include_columns)) + # make sure to shorten identifier + naming = NamingConvention(63) + # name must start with _dlt so we skip this table when replicating + snapshot_table_name = naming.normalize_table_identifier( + f"_dlt_{table_name}_s_{snapshot_name}" + ) + snapshot_qual_name = _make_qualified_table_name(snapshot_table_name, schema_name) + qual_name = _make_qualified_table_name(table_name, schema_name) + cur.execute( + f""" + CREATE TABLE {snapshot_qual_name} AS SELECT {col_str} FROM {qual_name}; + """ + ) + logger.info(f"Successfully persisted snapshot table state in {snapshot_qual_name}.") + return snapshot_table_name + + +def snapshot_table_resource( + snapshot_table_name: str, + schema_name: str, + table_name: str, + primary_key: TColumnNames, + write_disposition: TWriteDisposition, + columns: TTableSchemaColumns = None, + credentials: ConnectionStringCredentials = dlt.secrets.value, +) -> DltResource: + """Returns a resource for a persisted snapshot table. + + Can be used to perform an initial load of the table, so all data that + existed in the table prior to initializing replication is also captured. + """ + resource: DltResource = core_sql_table( + credentials=credentials, + table=snapshot_table_name, + schema=schema_name, + reflection_level="full_with_precision", + ) + resource.apply_hints( + table_name=table_name, + write_disposition=write_disposition, + columns=columns, + primary_key=primary_key, + ) + return resource + + +def get_max_lsn( + slot_name: str, + options: Dict[str, str], + credentials: ConnectionStringCredentials, +) -> Optional[int]: + """Returns maximum Log Sequence Number (LSN) in replication slot. + + Returns None if the replication slot is empty. + Does not consume the slot, i.e. messages are not flushed. + Raises error if the replication slot or publication does not exist. + """ + # comma-separated value string + options_str = ", ".join( + f"'{x}'" for xs in list(map(list, options.items())) for x in xs + ) + with _get_conn(credentials) as conn: + with conn.cursor() as cur: + cur.execute( + "SELECT MAX(lsn) - '0/0' AS max_lsn " # subtract '0/0' to convert pg_lsn type to int (https://stackoverflow.com/a/73738472) + f"FROM pg_logical_slot_peek_binary_changes('{slot_name}', NULL, NULL, {options_str});" + ) + lsn: int = cur.fetchone()[0] + return lsn + + +def get_pub_ops( + pub_name: str, + cur: cursor, +) -> Dict[str, bool]: + """Returns dictionary of DML operations and their publish status.""" + + cur.execute( + f""" + SELECT pubinsert, pubupdate, pubdelete, pubtruncate + FROM pg_publication WHERE pubname = '{pub_name}' + """ + ) + result = cur.fetchone() + # credentials.connection.commit() + if result is None: + raise ValueError(f'Publication "{pub_name}" does not exist.') + return { + "insert": result[0], + "update": result[1], + "delete": result[2], + "truncate": result[3], + } + + +def lsn_int_to_hex(lsn: int) -> str: + """Convert integer LSN to postgres' hexadecimal representation.""" + # https://stackoverflow.com/questions/66797767/lsn-external-representation. + return f"{lsn >> 32 & 4294967295:X}/{lsn & 4294967295:08X}" + + +def advance_slot( + upto_lsn: int, + slot_name: str, + credentials: ConnectionStringCredentials, +) -> None: + """Advances position in the replication slot. + + Flushes all messages upto (and including) the message with LSN = `upto_lsn`. + This function is used as alternative to psycopg2's `send_feedback` method, because + the behavior of that method seems odd when used outside of `consume_stream`. + """ + if upto_lsn != 0: + with _get_conn(credentials) as conn: + with conn.cursor() as cur: + cur.execute( + f"SELECT * FROM pg_replication_slot_advance('{slot_name}', '{lsn_int_to_hex(upto_lsn)}');" + ) + + +def _get_conn( + credentials: ConnectionStringCredentials, + connection_factory: Optional[Any] = None, +) -> ConnectionExt: + """Returns a psycopg2 connection to interact with postgres.""" + conn = psycopg2.connect( + dsn=credentials.to_native_representation(), + connection_factory=connection_factory, + ) + if connection_factory is not LogicalReplicationConnection: + conn.autocommit = False + return conn # type: ignore[no-any-return] + + +def _get_rep_conn( + credentials: ConnectionStringCredentials, +) -> LogicalReplicationConnection: + """Returns a psycopg2 LogicalReplicationConnection to interact with postgres replication functionality. + + Raises error if the user does not have the REPLICATION attribute assigned. + """ + return _get_conn(credentials, LogicalReplicationConnection) # type: ignore[return-value] + + +def _make_qualified_table_name(table_name: str, schema_name: str) -> str: + """Escapes and combines a schema and table name.""" + return ( + escape_postgres_identifier(schema_name) + + "." + + escape_postgres_identifier(table_name) + ) + + +def _get_pk( + cur: cursor, + table_name: str, + schema_name: str, +) -> Optional[TColumnNames]: + """Returns primary key column(s) for postgres table. + + Returns None if no primary key columns exist. + """ + qual_name = _make_qualified_table_name(table_name, schema_name) + # https://wiki.postgresql.org/wiki/Retrieve_primary_key_columns + cur.execute( + f""" + SELECT a.attname + FROM pg_index i + JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) + WHERE i.indrelid = '{qual_name}'::regclass + AND i.indisprimary; + """ + ) + result = [tup[0] for tup in cur.fetchall()] + if len(result) == 0: + return None + elif len(result) == 1: + return result[0] # type: ignore[no-any-return] + return result + + +@dataclass +class ItemGenerator: + credentials: ConnectionStringCredentials + slot_name: str + options: Dict[str, str] + upto_lsn: int + start_lsn: int = 0 + target_batch_size: int = 1000 + include_columns: Optional[Dict[str, Sequence[str]]] = (None,) # type: ignore[assignment] + columns: Optional[Dict[str, TTableSchemaColumns]] = (None,) # type: ignore[assignment] + last_commit_lsn: Optional[int] = field(default=None, init=False) + generated_all: bool = False + + def __iter__(self) -> Iterator[Union[TDataItem, DataItemWithMeta]]: + """Yields replication messages from MessageConsumer. + + Starts replication of messages published by the publication from the replication slot. + Maintains LSN of last consumed Commit message in object state. + Does not advance the slot. + """ + try: + cur = _get_rep_conn(self.credentials).cursor() + pub_ops = get_pub_ops(self.options["publication_names"], cur) + cur.start_replication( + slot_name=self.slot_name, + start_lsn=self.start_lsn, + decode=False, + options=self.options, + ) + consumer = MessageConsumer( + upto_lsn=self.upto_lsn, + pub_ops=pub_ops, + target_batch_size=self.target_batch_size, + include_columns=self.include_columns, + columns=self.columns, + ) + cur.consume_stream(consumer) + except StopReplication: # completed batch or reached `upto_lsn` + pass + finally: + cur.connection.close() + self.last_commit_lsn = consumer.last_commit_lsn + for rel_id, data_items in consumer.data_items.items(): + table_name = consumer.last_table_schema[rel_id]["name"] + # skip dlt internal tables + # TODO: normalize the prefix. low prio. on postgres we should have lower case + if table_name.lower().startswith(DLT_NAME_PREFIX): + continue + yield data_items[0] # meta item with column hints only, no data + yield dlt.mark.with_table_name(data_items[1:], table_name) + self.generated_all = consumer.consumed_all + + +class MessageConsumer: + """Consumes messages from a ReplicationCursor sequentially. + + Generates data item for each `insert`, `update`, and `delete` message. + Processes in batches to limit memory usage. + Maintains message data needed by subsequent messages in internal state. + """ + + def __init__( + self, + upto_lsn: int, + pub_ops: Dict[str, bool], + target_batch_size: int = 1000, + include_columns: Optional[Dict[str, Sequence[str]]] = None, + columns: Optional[Dict[str, TTableSchemaColumns]] = None, + ) -> None: + self.upto_lsn = upto_lsn + self.pub_ops = pub_ops + self.target_batch_size = target_batch_size + self.include_columns = include_columns + self.columns = columns + + self.consumed_all: bool = False + # data_items attribute maintains all data items + self.data_items: Dict[int, List[Union[TDataItem, DataItemWithMeta]]] = ( + dict() + ) # maps relation_id to list of data items + # other attributes only maintain last-seen values + self.last_table_schema: Dict[int, TTableSchema] = ( + dict() + ) # maps relation_id to table schema + self.last_commit_ts: pendulum.DateTime + self.last_commit_lsn = None + + def __call__(self, msg: ReplicationMessage) -> None: + """Processes message received from stream.""" + self.process_msg(msg) + + def process_msg(self, msg: ReplicationMessage) -> None: + """Processes encoded replication message. + + Identifies message type and decodes accordingly. + Message treatment is different for various message types. + Breaks out of stream with StopReplication exception when + - `upto_lsn` is reached + - `target_batch_size` is reached + - a table's schema has changed + """ + op = msg.payload[0] + if op == 73: # ASCII for 'I' + self.process_change(Insert(msg.payload), msg.data_start) + elif op == 85: # ASCII for 'U' + self.process_change(Update(msg.payload), msg.data_start) + elif op == 68: # ASCII for 'D' + self.process_change(Delete(msg.payload), msg.data_start) + elif op == 66: # ASCII for 'B' + self.last_commit_ts = Begin(msg.payload).commit_ts # type: ignore[assignment] + elif op == 67: # ASCII for 'C' + self.process_commit(msg) + elif op == 82: # ASCII for 'R' + self.process_relation(Relation(msg.payload)) + elif op == 84: # ASCII for 'T' + logger.warning( + "The truncate operation is currently not supported. " + "Truncate replication messages are ignored." + ) + else: + raise ValueError(f"Unknown replication op {op}") + + def process_commit(self, msg: ReplicationMessage) -> None: + """Updates object state when Commit message is observed. + + Raises StopReplication when `upto_lsn` or `target_batch_size` is reached. + """ + self.last_commit_lsn = msg.data_start + if msg.data_start >= self.upto_lsn: + self.consumed_all = True + n_items = sum( + [len(items) for items in self.data_items.values()] + ) # combine items for all tables + if self.consumed_all or n_items >= self.target_batch_size: + raise StopReplication + + def process_relation(self, decoded_msg: Relation) -> None: + """Processes a replication message of type Relation. + + Stores table schema in object state. + Creates meta item to emit column hints while yielding data. + + Raises StopReplication when a table's schema changes. + """ + if ( + self.data_items.get(decoded_msg.relation_id) is not None + ): # table schema change + raise StopReplication + # get table schema information from source and store in object state + table_name = decoded_msg.relation_name + columns: TTableSchemaColumns = { + c.name: _to_dlt_column_schema(c) for c in decoded_msg.columns + } + + self.last_table_schema[decoded_msg.relation_id] = { + "name": table_name, + "columns": columns, + } + + # apply user input + # 1) exclude columns + include_columns = ( + None + if self.include_columns is None + else self.include_columns.get(table_name) + ) + if include_columns is not None: + columns = {k: v for k, v in columns.items() if k in include_columns} + # 2) override source hints + column_hints: TTableSchemaColumns = ( + dict() if self.columns is None else self.columns.get(table_name, dict()) + ) + for column_name, column_val in column_hints.items(): + columns[column_name] = merge_column(columns[column_name], column_val) + + # add hints for replication columns + columns["lsn"] = {"data_type": "bigint", "nullable": True} + if self.pub_ops["update"] or self.pub_ops["delete"]: + columns["lsn"]["dedup_sort"] = "desc" + if self.pub_ops["delete"]: + columns["deleted_ts"] = { + "hard_delete": True, + "data_type": "timestamp", + "nullable": True, + } + + # determine write disposition + write_disposition: TWriteDisposition = "append" + if self.pub_ops["update"] or self.pub_ops["delete"]: + write_disposition = "merge" + + # include meta item to emit hints while yielding data + meta_item = dlt.mark.with_hints( + [], + dlt.mark.make_hints( + table_name=table_name, + write_disposition=write_disposition, + columns=columns, + ), + create_table_variant=True, + ) + self.data_items[decoded_msg.relation_id] = [meta_item] + + def process_change( + self, decoded_msg: Union[Insert, Update, Delete], msg_start_lsn: int + ) -> None: + """Processes replication message of type Insert, Update, or Delete. + + Adds data item for inserted/updated/deleted record to instance attribute. + """ + if isinstance(decoded_msg, (Insert, Update)): + column_data = decoded_msg.new_tuple.column_data + elif isinstance(decoded_msg, Delete): + column_data = decoded_msg.old_tuple.column_data + table_name = self.last_table_schema[decoded_msg.relation_id]["name"] + data_item = self.gen_data_item( + data=column_data, + column_schema=self.last_table_schema[decoded_msg.relation_id]["columns"], + lsn=msg_start_lsn, + commit_ts=self.last_commit_ts, + for_delete=isinstance(decoded_msg, Delete), + include_columns=( + None + if self.include_columns is None + else self.include_columns.get(table_name) + ), + ) + self.data_items[decoded_msg.relation_id].append(data_item) + + @staticmethod + def gen_data_item( + data: List[ColumnData], + column_schema: TTableSchemaColumns, + lsn: int, + commit_ts: pendulum.DateTime, + for_delete: bool, + include_columns: Optional[Sequence[str]] = None, + ) -> TDataItem: + """Generates data item from replication message data and corresponding metadata.""" + data_item = { + schema["name"]: _to_dlt_val( + val=data.col_data, + data_type=schema["data_type"], + byte1=data.col_data_category, + for_delete=for_delete, + ) + for (schema, data) in zip(column_schema.values(), data) + if (True if include_columns is None else schema["name"] in include_columns) + } + data_item["lsn"] = lsn + if for_delete: + data_item["deleted_ts"] = commit_ts + return data_item diff --git a/omniload/source/pg_replication/schema_types.py b/omniload/source/pg_replication/schema_types.py new file mode 100644 index 000000000..518a5d3a5 --- /dev/null +++ b/omniload/source/pg_replication/schema_types.py @@ -0,0 +1,152 @@ +# Copyright 2022-2026 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. + +import json +from functools import lru_cache +from typing import Any, Dict, Optional + +from dlt.common import Decimal, logger +from dlt.common.data_types.type_helpers import coerce_value +from dlt.common.data_types.typing import TDataType +from dlt.common.schema.typing import TColumnSchema, TColumnType +from dlt.destinations.impl.postgres.factory import PostgresTypeMapper + +from .decoders import ColumnType + +_DUMMY_VALS: Dict[TDataType, Any] = { + "bigint": 0, + "binary": b" ", + "bool": True, + "json": [0], + "date": "2000-01-01", + "decimal": Decimal(0), + "double": 0.0, + "text": "", + "time": "00:00:00", + "timestamp": "2000-01-01T00:00:00", + "wei": 0, +} +"""Dummy values used to replace NULLs in NOT NULL columns in key-only delete records.""" + + +_PG_TYPES: Dict[int, str] = { + 16: "boolean", + 17: "bytea", + 20: "bigint", + 21: "smallint", + 23: "integer", + 701: "double precision", + 1043: "character varying", + 1082: "date", + 1083: "time without time zone", + 1114: "timestamp without time zone", + 1184: "timestamp with time zone", + 1700: "numeric", + 3802: "jsonb", + 114: "json", +} +"""Maps postgres type OID to type string. Only includes types present in PostgresTypeMapper.""" + + +def _get_precision(type_id: int, atttypmod: int) -> Optional[int]: + """Get precision from postgres type attributes.""" + # https://stackoverflow.com/a/3351120 + if type_id == 21: # smallint + return 16 + elif type_id == 23: # integer + return 32 + elif type_id == 20: # bigint + return 64 + if atttypmod != -1: + if type_id == 1700: # numeric + return ((atttypmod - 4) >> 16) & 65535 + elif type_id in ( + 1083, + 1184, + ): # time without time zone, timestamp with time zone + return atttypmod + elif type_id == 1043: # character varying + return atttypmod - 4 + return None + + +def _get_scale(type_id: int, atttypmod: int) -> Optional[int]: + """Get scale from postgres type attributes.""" + # https://stackoverflow.com/a/3351120 + if atttypmod != -1: + if type_id in (21, 23, 20): # smallint, integer, bigint + return 0 + if type_id == 1700: # numeric + return (atttypmod - 4) & 65535 + return None + + +@lru_cache(maxsize=None) +def _type_mapper() -> PostgresTypeMapper: + from dlt.destinations import postgres + + return PostgresTypeMapper(postgres().capabilities()) + + +def _to_dlt_column_type(type_id: int, atttypmod: int) -> TColumnType: + """Converts postgres type OID to dlt column type. + + Type OIDs not in _PG_TYPES mapping default to "text" type. + """ + pg_type = _PG_TYPES.get(type_id) + if pg_type is None: + logger.warning( + f"Type OID {type_id} is unknown and will be converted to `text` dlt type" + ) + pg_type = "character varying" + precision = _get_precision(type_id, atttypmod) + scale = _get_scale(type_id, atttypmod) + dlt_type = _type_mapper().from_destination_type(pg_type, precision, scale) + + # Set timezone flag for timestamp types + if type_id == 1114: # timestamp without time zone + dlt_type["timezone"] = False + elif type_id == 1184: # timestamp with time zone + dlt_type["timezone"] = True + return dlt_type + + +def _to_dlt_column_schema(col: ColumnType) -> TColumnSchema: + """Converts pypgoutput ColumnType to dlt column schema.""" + dlt_column_type = _to_dlt_column_type(col.type_id, col.atttypmod) + partial_column_schema = { + "name": col.name, + "primary_key": bool(col.part_of_pkey), + } + return {**dlt_column_type, **partial_column_schema} # type: ignore[typeddict-item] + + +def _to_dlt_val(val: str, data_type: TDataType, byte1: str, for_delete: bool) -> Any: + """Converts pgoutput's text-formatted value into dlt-compatible data value.""" + if byte1 == "n": + if for_delete: + # replace None with dummy value to prevent NOT NULL violations in staging table + return _DUMMY_VALS[data_type] + return None + elif byte1 == "t": + if data_type == "binary": + # https://www.postgresql.org/docs/current/datatype-binary.html#DATATYPE-BINARY-BYTEA-HEX-FORMAT + return bytes.fromhex(val.replace("\\x", "")) + elif data_type == "json": + return json.loads(val) + return coerce_value(data_type, "text", val) + else: + raise ValueError( + f"Byte1 in replication message must be 'n' or 't', not '{byte1}'." + ) diff --git a/omniload/source/pipedrive/adapter.py b/omniload/source/pipedrive/adapter.py index 11be9c82d..1b746bb5b 100644 --- a/omniload/source/pipedrive/adapter.py +++ b/omniload/source/pipedrive/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,32 +23,35 @@ To get an api key: https://pipedrive.readme.io/docs/how-to-find-the-api-token """ -from typing import Any, Dict, Iterator, List, Optional, Union # noqa: F401 +from typing import Any, Dict, Iterable, Iterator, List, Optional, Union import dlt from dlt.common import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc -from dlt.sources import DltResource, TDataItems +from dlt.common.time import ensure_pendulum_datetime +from dlt.common.typing import TDataItems +from dlt.sources import DltResource from .helpers import group_deal_flows from .helpers.custom_fields_munger import rename_fields, update_fields_mapping from .helpers.pages import get_pages, get_recent_items_incremental + +# Export v2 source for easy access +from .rest_v2 import pipedrive_v2_source from .settings import ENTITY_MAPPINGS, RECENTS_ENTITIES from .typing import TDataPage -@dlt.source(name="pipedrive", max_table_nesting=0) +@dlt.source(name="pipedrive") def pipedrive_source( pipedrive_api_key: str = dlt.secrets.value, - since_timestamp: Union[pendulum.DateTime, str] = "1970-01-01 00:00:00", -) -> Iterator[DltResource]: + since_timestamp: Optional[Union[pendulum.DateTime, str]] = "1970-01-01 00:00:00", +) -> Iterable[DltResource]: """ Get data from the Pipedrive API. Supports incremental loading and custom fields mapping. Args: pipedrive_api_key: https://pipedrive.readme.io/docs/how-to-find-the-api-token since_timestamp: Starting timestamp for incremental loading. By default complete history is loaded on first run. - incremental: Enable or disable incremental loading. Returns resources: custom_fields_mapping @@ -67,19 +70,23 @@ def pipedrive_source( stages users leads + projects + tasks For custom fields rename the `custom_fields_mapping` resource must be selected or loaded before other resources. Resources that depend on another resource are implemented as transformers so they can re-use the original resource data without re-downloading. Examples: deals_participants, deals_flow + + Note: For v2 API endpoints, use pipedrive_v2_source from pipedrive.rest_v2 """ # yield nice rename mapping yield create_state(pipedrive_api_key) | parsed_mapping # parse timestamp and build kwargs - since_timestamp = ensure_pendulum_datetime_utc(since_timestamp).strftime( + since_timestamp = ensure_pendulum_datetime(since_timestamp).strftime( "%Y-%m-%d %H:%M:%S" ) resource_kwargs: Any = ( @@ -107,10 +114,8 @@ def pipedrive_source( name="deals_flow", write_disposition="merge", primary_key="id" )(_get_deals_flow)(pipedrive_api_key) - yield leads( - pipedrive_api_key, - update_time=since_timestamp, # ty: ignore[invalid-argument-type] - ) + # if simple value is passed in place of incremental, it will be used as initial value + yield leads(pipedrive_api_key, update_time=since_timestamp) # type: ignore[arg-type] def _get_deals_flow( diff --git a/omniload/source/pipedrive/settings.py b/omniload/source/pipedrive/settings.py index bed27dc50..ffa0dd473 100644 --- a/omniload/source/pipedrive/settings.py +++ b/omniload/source/pipedrive/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -36,6 +36,97 @@ "organization": "organizations", "pipeline": "pipelines", "product": "products", + "project": "projects", "stage": "stages", + "task": "tasks", "user": "users", } + + +""" +Available Pipedrive API v2 endpoints for configuration. + +Note: Some endpoints (e.g., followers, deal_products) require nested configuration. +See NESTED_ENTITIES_V2 for examples. + +# For more details, see: https://developers.pipedrive.com/docs/api/v2 +""" +ENTITIES_V2 = { + "activities": {}, + "deals": { + "params": { + "include_fields": ( + "next_activity_id,last_activity_id,first_won_time,products_count," + "files_count,notes_count,followers_count,email_messages_count," + "activities_count,done_activities_count,undone_activities_count," + "participants_count,last_incoming_mail_time,last_outgoing_mail_time," + "smart_bcc_email" + ) + } + }, + "persons": { + "params": { + "include_fields": ( + "next_activity_id,last_activity_id,open_deals_count," + "related_open_deals_count,closed_deals_count,related_closed_deals_count," + "participant_open_deals_count,participant_closed_deals_count," + "email_messages_count,activities_count,done_activities_count," + "undone_activities_count,files_count,notes_count,followers_count," + "won_deals_count,related_won_deals_count,lost_deals_count," + "related_lost_deals_count,last_incoming_mail_time,last_outgoing_mail_time" + ) + } + }, + "organizations": { + "params": { + "include_fields": ( + "next_activity_id,last_activity_id,open_deals_count," + "related_open_deals_count,closed_deals_count,related_closed_deals_count," + "email_messages_count,activities_count,done_activities_count," + "undone_activities_count,files_count,notes_count,followers_count," + "won_deals_count,related_won_deals_count,lost_deals_count," + "related_lost_deals_count" + ) + } + }, + "products": {}, + "pipelines": {}, + "stages": {}, +} + +# Nested V2 API Endpoints Configuration +# Automatically loaded when their parent resource is included in use_v2_endpoints. +NESTED_ENTITIES_V2 = { + "deal_products": { + "parent": "deals", + "endpoint_path": "deals/{id}/products", + "params": { + "limit": 500, + }, + }, + "deal_followers": { + "parent": "deals", + "endpoint_path": "deals/{id}/followers", + "primary_key": [ + "user_id", + "_deals_id", + ], # Followers don't have 'id', use composite key + "include_from_parent": ["id"], # Include deal id from parent + "params": { + "limit": 500, + }, + }, +} + +# Default v2 resources to load when none are specified +# This curated set includes the most commonly used endpoints. +# Users can customize this list to match their needs. +# See ENTITIES_V2 above for all available v2 endpoints. +DEFAULT_V2_RESOURCES = [ + "deals", + "persons", + "organizations", + "products", + "pipelines", + "stages", +] diff --git a/omniload/source/pipedrive/typing.py b/omniload/source/pipedrive/typing.py index 05759fbf8..4269940c8 100644 --- a/omniload/source/pipedrive/typing.py +++ b/omniload/source/pipedrive/typing.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/pokemon/adapter.py b/omniload/source/pokemon/adapter.py new file mode 100644 index 000000000..7238ed907 --- /dev/null +++ b/omniload/source/pokemon/adapter.py @@ -0,0 +1,58 @@ +# Copyright 2022-2026 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 provides data extraction from an example source as a starting point for new pipelines. +Available resources: [berries, pokemon] +""" + +import typing as t +from typing import Any, Dict, Iterable, Sequence + +import dlt +from dlt.common.typing import TDataItem +from dlt.sources import DltResource +from dlt.sources.helpers import requests + +from .settings import BERRY_URL, POKEMON_URL + + +@dlt.resource(write_disposition="replace") +def berries() -> Iterable[TDataItem]: + """ + Returns a list of berries. + Yields: + dict: The berries data. + """ + yield requests.get(BERRY_URL).json()["results"] + + +@dlt.resource(write_disposition="replace") +def pokemon() -> Iterable[TDataItem]: + """ + Returns a list of pokemon. + Yields: + dict: The pokemon data. + """ + yield requests.get(POKEMON_URL).json()["results"] + + +@dlt.source +def source() -> Sequence[DltResource]: + """ + The source function that returns all availble resources. + Returns: + Sequence[DltResource]: A sequence of DltResource objects containing the fetched data. + """ + return [berries, pokemon] diff --git a/omniload/source/pokemon/helpers.py b/omniload/source/pokemon/helpers.py new file mode 100644 index 000000000..4ce71f7c7 --- /dev/null +++ b/omniload/source/pokemon/helpers.py @@ -0,0 +1,15 @@ +# Copyright 2022-2026 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. + +"""Pokemon pipeline helpers""" diff --git a/omniload/source/pokemon/settings.py b/omniload/source/pokemon/settings.py new file mode 100644 index 000000000..08ef7307e --- /dev/null +++ b/omniload/source/pokemon/settings.py @@ -0,0 +1,18 @@ +# Copyright 2022-2026 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. + +"""Pokemon Pipeline settings and constants""" + +BERRY_URL = "https://pokeapi.co/api/v2/berry" +POKEMON_URL = "https://pokeapi.co/api/v2/pokemon/" diff --git a/omniload/source/salesforce/adapter.py b/omniload/source/salesforce/adapter.py index 0a255000e..fd77b1bb4 100644 --- a/omniload/source/salesforce/adapter.py +++ b/omniload/source/salesforce/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,48 +12,43 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Iterable +"""Source for Salesforce depending on the simple_salesforce python package. + +Imported resources are: account, campaign, contact, lead, opportunity, pricebook_2, pricebook_entry, product_2, user and user_role + +Salesforce api docs: https://developer.salesforce.com/docs/apis + +To get the security token: https://onlinehelp.coveo.com/en/ces/7.0/administrator/getting_the_security_token_for_your_salesforce_account.htm +""" + +from typing import Iterable, Optional import dlt from dlt.common.typing import TDataItem from dlt.sources import DltResource, incremental -from simple_salesforce import Salesforce +from dlt.sources.helpers.requests import Session -from .helpers import get_records +from .helpers.client import SalesforceAuth, make_salesforce_client +from .helpers.records import get_records @dlt.source(name="salesforce") def salesforce_source( - username: str, - password: str, - token: str, - domain: str, - custom_object: str, + credentials: SalesforceAuth = dlt.secrets.value, + session: Optional[Session] = None, ) -> Iterable[DltResource]: - """ - Retrieves data from Salesforce using the Salesforce API. - - Args: - username (str): The username for authentication. - password (str): The password for authentication. - token (str): The security token for authentication. - - Yields: - DltResource: Data resources from Salesforce. - """ - - client = Salesforce(username, password, token, domain=domain) + client = make_salesforce_client(credentials, session) # define resources @dlt.resource(write_disposition="replace") - def user() -> Iterable[TDataItem]: + def sf_user() -> Iterable[TDataItem]: yield get_records(client, "User") @dlt.resource(write_disposition="replace") def user_role() -> Iterable[TDataItem]: yield get_records(client, "UserRole") - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def opportunity( last_timestamp: incremental[str] = dlt.sources.incremental( "SystemModstamp", initial_value=None @@ -63,7 +58,7 @@ def opportunity( client, "Opportunity", last_timestamp.last_value, "SystemModstamp" ) - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def opportunity_line_item( last_timestamp: incremental[str] = dlt.sources.incremental( "SystemModstamp", initial_value=None @@ -73,7 +68,7 @@ def opportunity_line_item( client, "OpportunityLineItem", last_timestamp.last_value, "SystemModstamp" ) - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def opportunity_contact_role( last_timestamp: incremental[str] = dlt.sources.incremental( "SystemModstamp", initial_value=None @@ -86,7 +81,7 @@ def opportunity_contact_role( "SystemModstamp", ) - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def account( last_timestamp: incremental[str] = dlt.sources.incremental( "LastModifiedDate", initial_value=None @@ -108,7 +103,7 @@ def lead() -> Iterable[TDataItem]: def campaign() -> Iterable[TDataItem]: yield get_records(client, "Campaign") - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def campaign_member( last_timestamp: incremental[str] = dlt.sources.incremental( "SystemModstamp", initial_value=None @@ -119,18 +114,18 @@ def campaign_member( ) @dlt.resource(write_disposition="replace") - def product() -> Iterable[TDataItem]: + def product_2() -> Iterable[TDataItem]: yield get_records(client, "Product2") @dlt.resource(write_disposition="replace") - def pricebook() -> Iterable[TDataItem]: + def pricebook_2() -> Iterable[TDataItem]: yield get_records(client, "Pricebook2") @dlt.resource(write_disposition="replace") def pricebook_entry() -> Iterable[TDataItem]: yield get_records(client, "PricebookEntry") - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def task( last_timestamp: incremental[str] = dlt.sources.incremental( "SystemModstamp", initial_value=None @@ -138,7 +133,7 @@ def task( ) -> Iterable[TDataItem]: yield get_records(client, "Task", last_timestamp.last_value, "SystemModstamp") - @dlt.resource(write_disposition="merge", primary_key="id") + @dlt.resource(write_disposition="merge") def event( last_timestamp: incremental[str] = dlt.sources.incremental( "SystemModstamp", initial_value=None @@ -146,12 +141,8 @@ def event( ) -> Iterable[TDataItem]: yield get_records(client, "Event", last_timestamp.last_value, "SystemModstamp") - @dlt.resource(write_disposition="replace") - def custom() -> Iterable[TDataItem]: - yield get_records(client, custom_object) - return ( - user, + sf_user, user_role, opportunity, opportunity_line_item, @@ -161,10 +152,9 @@ def custom() -> Iterable[TDataItem]: lead, campaign, campaign_member, - product, - pricebook, + product_2, + pricebook_2, pricebook_entry, task, event, - custom, ) diff --git a/omniload/source/salesforce/settings.py b/omniload/source/salesforce/settings.py new file mode 100644 index 000000000..e82de10f6 --- /dev/null +++ b/omniload/source/salesforce/settings.py @@ -0,0 +1,18 @@ +# Copyright 2022-2026 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. + +"""Salesforce source settings and constants""" + +# set to false to limit query to 100 records for testing +IS_PRODUCTION = True diff --git a/omniload/source/scrapy/adapter.py b/omniload/source/scrapy/adapter.py new file mode 100644 index 000000000..129008451 --- /dev/null +++ b/omniload/source/scrapy/adapter.py @@ -0,0 +1,83 @@ +# Copyright 2022-2026 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. + +"""Scraping source + +Integrates Dlt and Scrapy to facilitate scraping pipelines. +""" + +import typing as t + +import dlt +from dlt.sources import DltResource +from scrapy import Spider + +from .helpers import ScrapingConfig, create_pipeline_runner +from .types import AnyDict + + +def run_pipeline( + pipeline: dlt.Pipeline, + spider: t.Type[Spider], + *args: t.Any, + on_before_start: t.Callable[[DltResource], None] = None, + scrapy_settings: t.Optional[AnyDict] = None, + batch_size: t.Optional[int] = None, + queue_size: t.Optional[int] = None, + queue_result_timeout: t.Optional[float] = None, + **kwargs: t.Any, +) -> None: + """Simple runner for the scraping pipeline + + You can pass all parameters via kwargs to `dlt.pipeline.run(....)` + + ``` + destination: TDestinationReferenceArg = None, + staging: TDestinationReferenceArg = None, + dataset_name: str = None, + credentials: Any = None, + table_name: str = None, + write_disposition: TWriteDisposition = None, + columns: TAnySchemaColumns = None, + primary_key: TColumnNames = None, + schema: Schema = None, + loader_file_format: TLoaderFileFormat = None + ``` + """ + options: AnyDict = {} + if scrapy_settings: + options["scrapy_settings"] = scrapy_settings + + if batch_size: + options["batch_size"] = batch_size + + if queue_size: + options["queue_size"] = queue_size + + if queue_result_timeout: + options["queue_result_timeout"] = queue_result_timeout + + scraping_host = create_pipeline_runner(pipeline, spider, **options) + + if on_before_start: + on_before_start(scraping_host.pipeline_runner.scraping_resource) + + scraping_host.run(*args, **kwargs) + + +@dlt.source(spec=ScrapingConfig) +def _register() -> DltResource: + raise NotImplementedError( + "Due to internal architecture of Scrapy, we could not represent it as a generator. Please use `run_pipeline` function instead" + ) diff --git a/omniload/source/scrapy/helpers.py b/omniload/source/scrapy/helpers.py new file mode 100644 index 000000000..c50b224c2 --- /dev/null +++ b/omniload/source/scrapy/helpers.py @@ -0,0 +1,123 @@ +# Copyright 2022-2026 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. + +import os +import typing as t + +import dlt +from dlt.common import logger +from dlt.common.configuration.inject import with_config +from dlt.common.configuration.specs.base_configuration import ( + BaseConfiguration, + configspec, +) +from scrapy import Spider + +from .queue import ScrapingQueue +from .runner import PipelineRunner, ScrapingHost, ScrapyRunner, Signals +from .settings import SOURCE_SCRAPY_QUEUE_SIZE, SOURCE_SCRAPY_SETTINGS +from .types import AnyDict + + +@configspec +class ScrapingConfig(BaseConfiguration): + # Batch size for scraped items + batch_size: int = 100 + + # maxsize for queue + queue_size: t.Optional[int] = SOURCE_SCRAPY_QUEUE_SIZE + + # result wait timeout for our queue + queue_result_timeout: t.Optional[float] = 1.0 + + # List of start urls + start_urls: t.List[str] = None + start_urls_file: str = None + + +@with_config(sections=("sources", "scraping"), spec=ScrapingConfig) +def resolve_start_urls( + start_urls: t.Optional[t.List[str]] = dlt.config.value, + start_urls_file: t.Optional[str] = dlt.config.value, +) -> t.List[str]: + """Merges start urls + If both `start_urls` and `start_urls_file` given, we will merge them + and return deduplicated list of `start_urls` for scrapy spider. + """ + urls = set() + if os.path.exists(start_urls_file): + with open(start_urls_file, encoding="utf-8") as fp: + urls = {line for line in fp.readlines() if str(line).strip()} + + if start_urls: + for url in start_urls: + urls.add(url) + + return list(set(urls)) + + +@with_config(sections=("sources", "scraping"), spec=ScrapingConfig) +def create_pipeline_runner( + pipeline: dlt.Pipeline, + spider: t.Type[Spider], + batch_size: int = dlt.config.value, + queue_size: int = dlt.config.value, + queue_result_timeout: float = dlt.config.value, + scrapy_settings: t.Optional[AnyDict] = None, +) -> ScrapingHost: + """Creates scraping host instance + This helper only creates pipeline host, so running and controlling + scrapy runner and pipeline is completely delegated to advanced users + """ + queue = ScrapingQueue( # type: ignore + maxsize=queue_size, + batch_size=batch_size, + read_timeout=queue_result_timeout, + ) + + signals = Signals( + pipeline_name=pipeline.pipeline_name, + queue=queue, + ) + + settings = {**SOURCE_SCRAPY_SETTINGS} + if scrapy_settings: + settings = {**SOURCE_SCRAPY_SETTINGS, **scrapy_settings} + + # sync scrapy log level with dlt's logger unless explicitly overridden + if "LOG_LEVEL" not in settings: + if logger.is_logging(): + settings["LOG_LEVEL"] = logger.log_level() + else: + settings["LOG_LEVEL"] = "INFO" + + scrapy_runner = ScrapyRunner( + spider=spider, + start_urls=resolve_start_urls(), + signals=signals, + settings=settings, + ) + + pipeline_runner = PipelineRunner( + pipeline=pipeline, + queue=queue, + ) + + scraping_host = ScrapingHost( + queue, + scrapy_runner, + pipeline_runner, + ) + + return scraping_host diff --git a/omniload/source/scrapy/queue.py b/omniload/source/scrapy/queue.py new file mode 100644 index 000000000..6679d9ee0 --- /dev/null +++ b/omniload/source/scrapy/queue.py @@ -0,0 +1,104 @@ +# Copyright 2022-2026 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. + +import typing as t +from queue import Empty, Queue + +from dlt.common import logger + +# Please read more at https://mypy.readthedocs.io/en/stable/runtime_troubles.html#not-generic-runtime +T = t.TypeVar("T") + +if t.TYPE_CHECKING: + + class _Queue(Queue[T]): + pass + +else: + + class _Queue(Queue, t.Generic[T]): + pass + + +class QueueClosedError(Exception): + pass + + +class ScrapingQueue(_Queue[T]): + def __init__( + self, + maxsize: int = 0, + batch_size: int = 10, + read_timeout: float = 1.0, + ) -> None: + super().__init__(maxsize) + self.batch_size = batch_size + self.read_timeout = read_timeout + self._is_closed = False + + def get_batches(self) -> t.Iterator[t.Any]: + """Batching helper can be wrapped as a dlt.resource + + Returns: + Iterator[Any]: yields scraped items one by one + """ + batch: t.List[T] = [] + while True: + if len(batch) == self.batch_size: + yield batch + batch = [] + + try: + if self.is_closed: + try: + item = self.get_nowait() + except Empty: + raise QueueClosedError("Queue is closed") + else: + item = self.get(timeout=self.read_timeout) + batch.append(item) + + # Mark task as completed + self.task_done() + except Empty: + if batch: + yield batch + batch = [] + except QueueClosedError: + logger.info("Queue is closed and drained") + + if batch: + yield batch + + break + + def stream(self) -> t.Iterator[t.Any]: + """Streaming generator, wraps get_batches + and handles `GeneratorExit` if dlt closes it. + + Returns: + t.Iterator[t.Any]: returns batches of scraped content + """ + try: + yield from self.get_batches() + except GeneratorExit: + self.close() + + def close(self) -> None: + """Marks queue as closed""" + self._is_closed = True + + @property + def is_closed(self) -> bool: + return self._is_closed diff --git a/omniload/source/scrapy/runner.py b/omniload/source/scrapy/runner.py new file mode 100644 index 000000000..cd4372196 --- /dev/null +++ b/omniload/source/scrapy/runner.py @@ -0,0 +1,278 @@ +# Copyright 2022-2026 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 module contains abstractions to facilitate scraping and loading process""" + +import logging as _logging +import threading +import typing as t + +import dlt +from dlt.common import logger +from scrapy import Item, Spider, signals +from scrapy.crawler import Crawler, CrawlerRunner +from twisted.internet import reactor as _reactor_module + +from .queue import ScrapingQueue +from .types import AnyDict, Runnable + +T = t.TypeVar("T", bound=Item) + +_REACTOR: t.Any = _reactor_module +_REACTOR_THREAD: t.Optional[threading.Thread] = None +_REACTOR_LOCK = threading.Lock() + + +def _ensure_reactor() -> None: + """Start the Twisted reactor in a daemon thread if not already running.""" + global _REACTOR_THREAD + with _REACTOR_LOCK: + if _REACTOR_THREAD is not None and _REACTOR_THREAD.is_alive(): + return + + started = threading.Event() + + def _run_reactor() -> None: + _REACTOR.callWhenRunning(started.set) + _REACTOR.run(installSignalHandlers=False) + + _REACTOR_THREAD = threading.Thread(target=_run_reactor, daemon=True) + _REACTOR_THREAD.start() + started.wait() + + +class Signals(t.Generic[T]): + """Manages signal connections between a Scrapy Crawler and the scraping queue. + + Connects to spider_closed (not engine_stopped) to close the queue. This is + important because spider_closed fires AFTER all item_scraped callbacks — + guaranteed by the engine's deferred chain: + scraper.close_spider() waits for slot.is_idle() → spider_closed signal + Using engine_stopped would race with pending item_scraped callbacks. + """ + + def __init__(self, pipeline_name: str, queue: ScrapingQueue[T]) -> None: + self._stopped = False + self.queue = queue + self.pipeline_name = pipeline_name + self._crawler: t.Optional[Crawler] = None + self._runner: t.Optional[CrawlerRunner] = None + + def connect(self, crawler: Crawler, runner: CrawlerRunner) -> None: + """Connect signal handlers to a Crawler. Must be called from the reactor thread.""" + if self._crawler is not None: + raise RuntimeError("Signals already connected to a crawler") + self._crawler = crawler + self._runner = runner + crawler.signals.connect(self.on_item_scraped, signals.item_scraped) + crawler.signals.connect(self.on_spider_closed, signals.spider_closed) + + def disconnect(self) -> None: + """Disconnect signal handlers. Must be called from the reactor thread.""" + if self._crawler is not None: + self._crawler.signals.disconnect(self.on_item_scraped, signals.item_scraped) + self._crawler.signals.disconnect( + self.on_spider_closed, signals.spider_closed + ) + self._crawler = None + self._runner = None + + def on_item_scraped(self, item: T) -> None: + self.queue.put(item) + + def on_spider_closed(self, spider: t.Any, reason: str) -> None: + logger.info(f"Spider closed for pipeline={self.pipeline_name} reason={reason}") + self._initiate_stop() + + def _initiate_stop(self) -> None: + """Idempotent stop: close the queue and stop the crawler. + + Uses reactor.callFromThread to stop the runner, which is safe + from both the reactor thread and non-reactor threads. + """ + if self._stopped: + return + self._stopped = True + self.queue.close() + if self._runner is not None: + _REACTOR.callFromThread(self._runner.stop) + + +class ScrapyRunner(Runnable): + """Scrapy runner handles setup and teardown of scrapy crawling. + + Uses a persistent Twisted reactor running in a daemon thread, + allowing sequential pipeline runs in the same process. + """ + + def __init__( + self, + spider: t.Type[Spider], + start_urls: t.List[str], + settings: AnyDict, + signals: Signals[T], + ) -> None: + self.spider = spider + self.start_urls = start_urls + self.signals = signals + self.runner = CrawlerRunner(settings=settings) + # apply scrapy log level without touching root handlers (dlt owns those) + log_level = self.runner.settings.get("LOG_LEVEL", "WARNING") + _logging.getLogger("scrapy").setLevel(log_level) + _logging.getLogger("twisted").setLevel(_logging.ERROR) + + def run(self, *args: t.Any, **kwargs: t.Any) -> None: + """Runs scrapy crawler via the persistent reactor. + + Schedules the crawl in the reactor thread via callFromThread, + then blocks until the crawl completes. + """ + _ensure_reactor() + + done = threading.Event() + crawl_error: t.List[BaseException] = [] + + def _schedule() -> None: + crawler = self.runner.create_crawler(self.spider) + self.signals.connect(crawler, self.runner) + + d = self.runner.crawl( + crawler, + name="scraping_spider", + start_urls=self.start_urls, + **kwargs, + ) + + def _on_done(result: t.Any) -> None: + self.signals.disconnect() + done.set() + + def _on_error(failure: t.Any) -> None: + self.signals.disconnect() + crawl_error.append(failure.value) + done.set() + + d.addCallback(_on_done) + d.addErrback(_on_error) + + try: + logger.info("Starting the crawler") + _REACTOR.callFromThread(_schedule) + done.wait() + except Exception: + logger.error("Was unable to start crawling process") + raise + finally: + self.signals._initiate_stop() + logger.info("Scraping stopped") + + if crawl_error: + raise crawl_error[0] + + +class PipelineRunner(Runnable): + """Pipeline runner runs dlt pipeline in a separate thread + Since scrapy wants to run in the main thread it is the only available + option to host pipeline in a thread and communicate via the queue. + """ + + def __init__(self, pipeline: dlt.Pipeline, queue: ScrapingQueue[T]) -> None: + self.pipeline = pipeline + self.queue = queue + + if pipeline.dataset_name and not self.is_default_dataset_name(pipeline): + resource_name = pipeline.dataset_name + else: + resource_name = f"{pipeline.pipeline_name}_results" + + logger.info(f"Resource name: {resource_name}") + + self.scraping_resource = dlt.resource( + # Queue get_batches is a generator so we can + # pass it to pipeline.run and dlt will handle the rest. + self.queue.stream(), + name=resource_name, + ) + + def is_default_dataset_name(self, pipeline: dlt.Pipeline) -> bool: + default_name = pipeline.pipeline_name + pipeline.DEFAULT_DATASET_SUFFIX + return pipeline.dataset_name == default_name + + def run( + self, + *args: t.Any, + **kwargs: t.Any, + ) -> threading.Thread: + """You can use all regular dlt.pipeline.run() arguments + + ``` + destination: TDestinationReferenceArg = None, + staging: TDestinationReferenceArg = None, + dataset_name: str = None, + credentials: Any = None, + table_name: str = None, + write_disposition: TWriteDisposition = None, + columns: TAnySchemaColumns = None, + primary_key: TColumnNames = None, + schema: Schema = None, + loader_file_format: TLoaderFileFormat = None + ``` + """ + + def run() -> None: + try: + self.pipeline.run(self.scraping_resource, **kwargs) + except Exception: + logger.error("Error during pipeline.run call, closing the queue") + raise + finally: + self.queue.close() + + thread_runner = threading.Thread(target=run) + thread_runner.start() + return thread_runner + + +class ScrapingHost: + """Scraping host runs the pipeline and scrapy""" + + def __init__( + self, + queue: ScrapingQueue[T], + scrapy_runner: ScrapyRunner, + pipeline_runner: PipelineRunner, + ) -> None: + self.queue = queue + self.scrapy_runner = scrapy_runner + self.pipeline_runner = pipeline_runner + + def run( + self, + *args: t.Any, + **kwargs: t.Any, + ) -> None: + """You can pass kwargs which are passed to `pipeline.run`""" + logger.info("Starting pipeline") + pipeline_worker = self.pipeline_runner.run(*args, **kwargs) + + logger.info("Starting scrapy crawler") + self.scrapy_runner.run() + + # Wait for pipeline to finish its job + pipeline_worker.join(timeout=60) + if pipeline_worker.is_alive(): + logger.warning( + "Pipeline worker did not finish within timeout. Forcing queue close." + ) + self.queue.close() diff --git a/omniload/source/scrapy/settings.py b/omniload/source/scrapy/settings.py new file mode 100644 index 000000000..84ec5184a --- /dev/null +++ b/omniload/source/scrapy/settings.py @@ -0,0 +1,39 @@ +# Copyright 2022-2026 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 .types import AnyDict + +SOURCE_BATCH_SIZE: int = 10 +SOURCE_SCRAPY_QUEUE_SIZE: int = 3000 +SOURCE_SCRAPY_QUEUE_RESULT_TIMEOUT: int = 5 +SOURCE_SCRAPY_SETTINGS: AnyDict = { + "TELNETCONSOLE_ENABLED": False, + # How many sub pages to scrape + # https://docs.scrapy.org/en/latest/topics/settings.html#depth-limit + "DEPTH_LIMIT": 0, + "SPIDER_MIDDLEWARES": { + "scrapy.spidermiddlewares.depth.DepthMiddleware": 200, + "scrapy.spidermiddlewares.httperror.HttpErrorMiddleware": 300, + }, + "HTTPERROR_ALLOW_ALL": True, + "FAKEUSERAGENT_PROVIDERS": [ + # this is the first provider we'll try + "scrapy_fake_useragent.providers.FakeUserAgentProvider", + # if FakeUserAgentProvider fails, we'll use faker to generate a user-agent string for us + "scrapy_fake_useragent.providers.FakerProvider", + # fall back to USER_AGENT value + "scrapy_fake_useragent.providers.FixedUserAgentProvider", + ], + "USER_AGENT": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:123.0) Gecko/20100101 Firefox/123.0", +} diff --git a/omniload/source/scrapy/types.py b/omniload/source/scrapy/types.py new file mode 100644 index 000000000..19c6e1175 --- /dev/null +++ b/omniload/source/scrapy/types.py @@ -0,0 +1,22 @@ +# Copyright 2022-2026 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. + +import typing as t + +AnyDict = t.Dict[str, t.Any] + + +class Runnable(t.Protocol): + def run(self, *args: t.Any, **kwargs: t.Any) -> t.Any: + pass diff --git a/omniload/source/shopify/adapter.py b/omniload/source/shopify/adapter.py index 819c15ac1..1eafe610b 100644 --- a/omniload/source/shopify/adapter.py +++ b/omniload/source/shopify/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,27 +14,25 @@ """Fetches Shopify Orders and Products.""" -from typing import Any, Dict, Iterable, Optional # noqa: F401 +from typing import Any, Dict, Iterable, Iterator, Optional import dlt -from dlt.common import jsonpath as jp # noqa: F401 +from dlt.common import jsonpath as jp from dlt.common import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import TAnyDateTime, TDataItem from dlt.sources import DltResource -from omniload.error import MissingValueError - -from .helpers import ShopifyApi, ShopifyGraphQLApi, TOrderStatus +from .helpers import ShopifyApi, ShopifyPartnerApi, TOrderStatus from .settings import ( DEFAULT_API_VERSION, DEFAULT_ITEMS_PER_PAGE, - DEFAULT_PARTNER_API_VERSION, # noqa: F401 + DEFAULT_PARTNER_API_VERSION, FIRST_DAY_OF_MILLENNIUM, ) -@dlt.source(name="shopify", max_table_nesting=0) +@dlt.source(name="shopify") def shopify_source( private_app_password: str = dlt.secrets.value, api_version: str = DEFAULT_API_VERSION, @@ -74,99 +72,13 @@ def shopify_source( # build client client = ShopifyApi(shop_url, private_app_password, api_version) - start_date_obj = ensure_pendulum_datetime_utc(start_date) - end_date_obj = ensure_pendulum_datetime_utc(end_date) if end_date else None - created_at_min_obj = ensure_pendulum_datetime_utc(created_at_min) + start_date_obj = ensure_pendulum_datetime(start_date) + end_date_obj = ensure_pendulum_datetime(end_date) if end_date else None + created_at_min_obj = ensure_pendulum_datetime(created_at_min) # define resources - @dlt.resource( - primary_key="id", - write_disposition="merge", - columns={ - "body_html": { - "data_type": "text", - "nullable": True, - "description": "A description of the product. Supports HTML formatting.", - }, - "created_at": { - "data_type": "timestamp", - "nullable": False, - "description": "The date and time (ISO 8601 format) when the product was created.", - }, - "handle": { - "data_type": "text", - "nullable": False, - "description": "A unique human-friendly string for the product. Automatically generated from the product's title. Used by the Liquid templating language to refer to objects.", - }, - "id": { - "data_type": "bigint", - "nullable": False, - "primary_key": True, - "description": "An unsigned 64-bit integer that's used as a unique identifier for the product.", - }, - "images": { - "data_type": "json", - "nullable": True, - "description": "A list of product image objects, each one representing an image associated with the product.", - }, - "options": { - "data_type": "json", - "nullable": True, - "description": "The custom product properties. For example, Size, Color, and Material.", - }, - "product_type": { - "data_type": "text", - "nullable": True, - "description": "A categorization for the product used for filtering and searching products.", - }, - "published_at": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time (ISO 8601 format) when the product was published.", - }, - "published_scope": { - "data_type": "text", - "nullable": True, - "description": "Whether the product is published to the Point of Sale channel.", - }, - "status": { - "data_type": "text", - "nullable": False, - "description": "The status of the product.", - }, - "tags": { - "data_type": "text", - "nullable": True, - "description": "A string of comma-separated tags used for filtering and search.", - }, - "template_suffix": { - "data_type": "text", - "nullable": True, - "description": "The suffix of the Liquid template used for the product page.", - }, - "title": { - "data_type": "text", - "nullable": False, - "description": "The name of the product.", - }, - "updated_at": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time (ISO 8601 format) when the product was last modified.", - }, - "variants": { - "data_type": "json", - "nullable": True, - "description": "An array of product variants, each representing a different version of the product.", - }, - "vendor": { - "data_type": "text", - "nullable": True, - "description": "The name of the product's vendor.", - }, - }, - ) - def products_legacy( + @dlt.resource(primary_key="id", write_disposition="merge") + def products( updated_at: dlt.sources.incremental[ pendulum.DateTime ] = dlt.sources.incremental( @@ -174,8 +86,6 @@ def products_legacy( initial_value=start_date_obj, end_value=end_date_obj, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), created_at_min: pendulum.DateTime = created_at_min_obj, items_per_page: int = items_per_page, @@ -189,8 +99,6 @@ def products_legacy( Returns: Iterable[TDataItem]: A generator of products. """ - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Shopify") params = dict( updated_at_min=updated_at.last_value.isoformat(), limit=items_per_page, @@ -201,423 +109,7 @@ def products_legacy( params["updated_at_max"] = updated_at.end_value.isoformat() yield from client.get_pages("products", params) - @dlt.resource( - primary_key="id", - write_disposition="merge", - columns={ - "app_id": { - "data_type": "bigint", - "nullable": True, - "description": "The ID of the app that created the order.", - }, - "billing_address": { - "data_type": "json", - "nullable": True, - "description": "The mailing address associated with the payment method.", - }, - "browser_ip": { - "data_type": "text", - "nullable": True, - "description": "The IP address of the browser used by the customer when they placed the order.", - }, - "buyer_accepts_marketing": { - "data_type": "bool", - "nullable": True, - "description": "Whether the customer consented to receive email updates from the shop.", - }, - "cancel_reason": { - "data_type": "text", - "nullable": True, - "description": "The reason why the order was canceled.", - }, - "cancelled_at": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time when the order was canceled.", - }, - "cart_token": { - "data_type": "text", - "nullable": True, - "description": "A unique value referencing the cart associated with the order.", - }, - "checkout_token": { - "data_type": "text", - "nullable": True, - "description": "A unique value referencing the checkout associated with the order.", - }, - "client_details": { - "data_type": "json", - "nullable": True, - "description": "Information about the browser the customer used when placing the order.", - }, - "closed_at": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time when the order was closed.", - }, - "company": { - "data_type": "json", - "nullable": True, - "description": "Information about the purchasing company for the order.", - }, - "confirmation_number": { - "data_type": "text", - "nullable": True, - "description": "A randomly generated identifier for the order.", - }, - "confirmed": { - "data_type": "bool", - "nullable": True, - "description": "Whether inventory has been reserved for the order.", - }, - "created_at": { - "data_type": "timestamp", - "nullable": False, - "description": "The autogenerated date and time when the order was created.", - }, - "currency": { - "data_type": "text", - "nullable": False, - "description": "The three-letter code (ISO 4217 format) for the shop currency.", - }, - "current_total_additional_fees_set": { - "data_type": "json", - "nullable": True, - "description": "The current total additional fees on the order in shop and presentment currencies.", - }, - "current_total_discounts": { - "data_type": "decimal", - "nullable": True, - "description": "The current total discounts on the order in the shop currency.", - }, - "current_total_discounts_set": { - "data_type": "json", - "nullable": True, - "description": "The current total discounts on the order in shop and presentment currencies.", - }, - "current_total_duties_set": { - "data_type": "json", - "nullable": True, - "description": "The current total duties charged on the order in shop and presentment currencies.", - }, - "current_total_price": { - "data_type": "decimal", - "nullable": True, - "description": "The current total price of the order in the shop currency.", - }, - "current_total_price_set": { - "data_type": "json", - "nullable": True, - "description": "The current total price of the order in shop and presentment currencies.", - }, - "current_subtotal_price": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of prices for all line items after discounts and returns in the shop currency.", - }, - "current_subtotal_price_set": { - "data_type": "json", - "nullable": True, - "description": "The sum of the prices for all line items after discounts and returns in shop and presentment currencies.", - }, - "current_total_tax": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of the prices for all tax lines applied to the order in the shop currency.", - }, - "current_total_tax_set": { - "data_type": "json", - "nullable": True, - "description": "The sum of the prices for all tax lines applied to the order in shop and presentment currencies.", - }, - "customer": { - "data_type": "json", - "nullable": True, - "description": "Information about the customer.", - }, - "customer_locale": { - "data_type": "text", - "nullable": True, - "description": "The two or three-letter language code, optionally followed by a region modifier.", - }, - "discount_applications": { - "data_type": "json", - "nullable": True, - "description": "An ordered list of stacked discount applications.", - }, - "discount_codes": { - "data_type": "json", - "nullable": True, - "description": "A list of discounts applied to the order.", - }, - "email": { - "data_type": "text", - "nullable": True, - "description": "The customer's email address.", - }, - "estimated_taxes": { - "data_type": "bool", - "nullable": True, - "description": "Whether taxes on the order are estimated.", - }, - "financial_status": { - "data_type": "text", - "nullable": True, - "description": "The status of payments associated with the order.", - }, - "fulfillments": { - "data_type": "json", - "nullable": True, - "description": "An array of fulfillments associated with the order.", - }, - "fulfillment_status": { - "data_type": "text", - "nullable": True, - "description": "The order's status in terms of fulfilled line items.", - }, - "gateway": { - "data_type": "text", - "nullable": True, - "description": "The payment gateway used.", - }, - "id": { - "data_type": "bigint", - "nullable": False, - "primary_key": True, - "description": "The ID of the order, used for API purposes.", - }, - "landing_site": { - "data_type": "text", - "nullable": True, - "description": "The URL for the page where the buyer landed when they entered the shop.", - }, - "line_items": { - "data_type": "json", - "nullable": True, - "description": "A list of line item objects containing information about an item in the order.", - }, - "location_id": { - "data_type": "bigint", - "nullable": True, - "description": "The ID of one of the locations assigned to fulfill the order.", - }, - "merchant_of_record_app_id": { - "data_type": "bigint", - "nullable": True, - "description": "The application acting as Merchant of Record for the order.", - }, - "name": { - "data_type": "text", - "nullable": True, - "description": "The order name, generated by combining the order_number with the order prefix and suffix.", - }, - "note": { - "data_type": "text", - "nullable": True, - "description": "An optional note that a shop owner can attach to the order.", - }, - "note_attributes": { - "data_type": "json", - "nullable": True, - "description": "Extra information added to the order as key-value pairs.", - }, - "number": { - "data_type": "bigint", - "nullable": True, - "description": "The order's position in the shop's count of orders.", - }, - "order_number": { - "data_type": "bigint", - "nullable": True, - "description": "The order's position in the shop's count of orders, starting at 1001.", - }, - "original_total_additional_fees_set": { - "data_type": "json", - "nullable": True, - "description": "The original total additional fees on the order in shop and presentment currencies.", - }, - "original_total_duties_set": { - "data_type": "json", - "nullable": True, - "description": "The original total duties charged on the order in shop and presentment currencies.", - }, - "payment_terms": { - "data_type": "json", - "nullable": True, - "description": "The terms and conditions under which a payment should be processed.", - }, - "payment_gateway_names": { - "data_type": "json", - "nullable": True, - "description": "The list of payment gateways used for the order.", - }, - "phone": { - "data_type": "text", - "nullable": True, - "description": "The customer's phone number for receiving SMS notifications.", - }, - "po_number": { - "data_type": "text", - "nullable": True, - "description": "The purchase order number associated with the order.", - }, - "presentment_currency": { - "data_type": "text", - "nullable": True, - "description": "The presentment currency used to display prices to the customer.", - }, - "processed_at": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time when an order was processed.", - }, - "referring_site": { - "data_type": "text", - "nullable": True, - "description": "The website where the customer clicked a link to the shop.", - }, - "refunds": { - "data_type": "json", - "nullable": True, - "description": "A list of refunds applied to the order.", - }, - "shipping_address": { - "data_type": "json", - "nullable": True, - "description": "The mailing address where the order will be shipped.", - }, - "shipping_lines": { - "data_type": "json", - "nullable": True, - "description": "An array detailing the shipping methods used.", - }, - "source_name": { - "data_type": "text", - "nullable": True, - "description": "The source of the checkout.", - }, - "source_identifier": { - "data_type": "text", - "nullable": True, - "description": "The ID of the order placed on the originating platform.", - }, - "source_url": { - "data_type": "text", - "nullable": True, - "description": "A valid URL to the original order on the originating surface.", - }, - "subtotal_price": { - "data_type": "decimal", - "nullable": True, - "description": "The price of the order in the shop currency after discounts but before shipping, duties, taxes, and tips.", - }, - "subtotal_price_set": { - "data_type": "json", - "nullable": True, - "description": "The subtotal of the order in shop and presentment currencies after discounts but before shipping, duties, taxes, and tips.", - }, - "tags": { - "data_type": "text", - "nullable": True, - "description": "Tags attached to the order, formatted as a string of comma-separated values.", - }, - "tax_lines": { - "data_type": "json", - "nullable": True, - "description": "An array of tax line objects detailing taxes applied to the order.", - }, - "taxes_included": { - "data_type": "bool", - "nullable": True, - "description": "Whether taxes are included in the order subtotal.", - }, - "test": { - "data_type": "bool", - "nullable": True, - "description": "Whether this is a test order.", - }, - "token": { - "data_type": "text", - "nullable": True, - "description": "A unique value referencing the order.", - }, - "total_discounts": { - "data_type": "decimal", - "nullable": True, - "description": "The total discounts applied to the price of the order in the shop currency.", - }, - "total_discounts_set": { - "data_type": "json", - "nullable": True, - "description": "The total discounts applied to the price of the order in shop and presentment currencies.", - }, - "total_line_items_price": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of all line item prices in the shop currency.", - }, - "total_line_items_price_set": { - "data_type": "json", - "nullable": True, - "description": "The total of all line item prices in shop and presentment currencies.", - }, - "total_outstanding": { - "data_type": "decimal", - "nullable": True, - "description": "The total outstanding amount of the order in the shop currency.", - }, - "total_price": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of all line item prices, discounts, shipping, taxes, and tips in the shop currency.", - }, - "total_price_set": { - "data_type": "json", - "nullable": True, - "description": "The total price of the order in shop and presentment currencies.", - }, - "total_shipping_price_set": { - "data_type": "json", - "nullable": True, - "description": "The total shipping price of the order in shop and presentment currencies.", - }, - "total_tax": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of the prices for all tax lines applied to the order in the shop currency.", - }, - "total_tax_set": { - "data_type": "json", - "nullable": True, - "description": "The sum of the prices for all tax lines applied to the order in shop and presentment currencies.", - }, - "total_tip_received": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of all the tips in the order in the shop currency.", - }, - "total_weight": { - "data_type": "decimal", - "nullable": True, - "description": "The sum of all line item weights in grams.", - }, - "updated_at": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time when the order was last modified.", - }, - "user_id": { - "data_type": "bigint", - "nullable": True, - "description": "The ID of the user logged into Shopify POS who processed the order.", - }, - "order_status_url": { - "data_type": "text", - "nullable": True, - "description": "The URL pointing to the order status web page, if applicable.", - }, - }, - ) + @dlt.resource(primary_key="id", write_disposition="merge") def orders( updated_at: dlt.sources.incremental[ pendulum.DateTime @@ -626,8 +118,6 @@ def orders( initial_value=start_date_obj, end_value=end_date_obj, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), created_at_min: pendulum.DateTime = created_at_min_obj, items_per_page: int = items_per_page, @@ -642,8 +132,6 @@ def orders( Returns: Iterable[TDataItem]: A generator of orders. """ - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Shopify") params = dict( updated_at_min=updated_at.last_value.isoformat(), limit=items_per_page, @@ -664,8 +152,6 @@ def customers( initial_value=start_date_obj, end_value=end_date_obj, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), created_at_min: pendulum.DateTime = created_at_min_obj, items_per_page: int = items_per_page, @@ -679,8 +165,6 @@ def customers( Returns: Iterable[TDataItem]: A generator of customers. """ - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Shopify") params = dict( updated_at_min=updated_at.last_value.isoformat(), limit=items_per_page, @@ -691,1282 +175,67 @@ def customers( params["updated_at_max"] = updated_at.end_value.isoformat() yield from client.get_pages("customers", params) - @dlt.resource(primary_key="id", write_disposition="merge") - def events( - created_at: dlt.sources.incremental[ - pendulum.DateTime - ] = dlt.sources.incremental( - "created_at", - initial_value=start_date_obj, - end_value=end_date_obj, - range_end="closed", - range_start="closed", - ), - items_per_page: int = items_per_page, - ) -> Iterable[TDataItem]: - if created_at.last_value is None: - raise MissingValueError("created_at.last_value", "Shopify") - params = dict( - created_at_min=created_at.last_value.isoformat(), - limit=items_per_page, - order="created_at asc", - ) - yield from client.get_pages("events", params) - - @dlt.resource(primary_key="id", write_disposition="merge") - def price_rules( - updated_at: dlt.sources.incremental[ - pendulum.DateTime - ] = dlt.sources.incremental( - "updated_at", - initial_value=start_date_obj, - end_value=end_date_obj, - range_end="closed", - range_start="closed", - ), - items_per_page: int = items_per_page, - ) -> Iterable[TDataItem]: - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Shopify") - params = dict( - updated_at_min=updated_at.last_value.isoformat(), - limit=items_per_page, - order="updated_at asc", - ) - yield from client.get_pages("price_rules", params) + return (products, orders, customers) - @dlt.resource(primary_key="id", write_disposition="merge") - def transactions( - since_id: dlt.sources.incremental[int] = dlt.sources.incremental( - "id", - initial_value=None, - ), - items_per_page: int = items_per_page, - ) -> Iterable[TDataItem]: - params = dict( - limit=items_per_page, - ) - if since_id.start_value is not None: - params["since_id"] = since_id.start_value - yield from client.get_pages("shopify_payments/balance/transactions", params) - @dlt.resource( - primary_key="currency", - write_disposition={"disposition": "merge", "strategy": "scd2"}, - ) - def balance() -> Iterable[TDataItem]: - yield from client.get_pages("shopify_payments/balance", {}) +@dlt.resource +def shopify_partner_query( + query: str, + data_items_path: jp.TJsonPath, + pagination_cursor_path: jp.TJsonPath, + pagination_variable_name: str = "after", + variables: Optional[Dict[str, Any]] = None, + access_token: str = dlt.secrets.value, + organization_id: str = dlt.config.value, + api_version: str = DEFAULT_PARTNER_API_VERSION, +) -> Iterable[TDataItem]: + """ + Resource for getting paginated results from the Shopify Partner GraphQL API. - @dlt.resource(primary_key="id", write_disposition="merge") - def inventory_items( - updated_at: dlt.sources.incremental[ - pendulum.DateTime - ] = dlt.sources.incremental( - "updatedAt", - initial_value=start_date_obj, - end_value=end_date_obj, - allow_external_schedulers=True, - range_end="closed", - range_start="closed", - ), - items_per_page: int = items_per_page, - ) -> Iterable[TDataItem]: - client = ShopifyGraphQLApi( - base_url=shop_url, - access_token=private_app_password, - api_version="2024-07", - ) + This resource will run the given GraphQL query and extract a list of data items from the result. + It will then run the query again with a pagination cursor to get the next page of results. - query = """ - query inventoryItems($after: String, $query: String, $first: Int) { - inventoryItems(after: $after, first: $first, query: $query) { + Example: + query = '''query Transactions($after: String) { + transactions(after: $after, first: 100) { edges { - node { - id - countryCodeOfOrigin - createdAt - duplicateSkuCount - harmonizedSystemCode - inventoryHistoryUrl - legacyResourceId - measurement { - id - weight { - unit - value - } - } - - provinceCodeOfOrigin - requiresShipping - sku - tracked - trackedEditable { - locked - reason - } - unitCost { - amount - currencyCode - } - updatedAt - variant { - id - availableForSale - barcode - - compareAtPrice - createdAt - inventoryPolicy - inventoryQuantity - legacyResourceId - - position - price - product { + cursor + node { id } - requiresComponents - - selectedOptions { - name - value - } - sellableOnlineQuantity - - sellingPlanGroupsCount { - count - precision - } - sku - - taxCode - taxable - title - updatedAt - } - } - } - pageInfo { - endCursor - } - } - }""" - - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Shopify") - - yield from client.get_graphql_pages( - query, - data_items_path="data.inventoryItems.edges[*].node", - pagination_cursor_path="data.inventoryItems.pageInfo.endCursor", - pagination_variable_name="after", - variables={ - "query": f"updated_at:>'{updated_at.last_value.isoformat()}'", - "first": items_per_page, - }, - ) - - @dlt.resource(primary_key="id", write_disposition="merge") - def discounts(items_per_page: int = items_per_page) -> Iterable[TDataItem]: - client = ShopifyGraphQLApi( - base_url=shop_url, - access_token=private_app_password, - api_version="2024-07", - ) - - query = """ -query discountNodes($after: String, $query: String, $first: Int) { - discountNodes(after: $after, first: $first, query: $query) { - nodes { - id - discount { - __typename - ... on DiscountCodeApp { - appDiscountType { - app { - id - } - functionId - targetType - } - appliesOncePerCustomer - asyncUsageCount - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - codesCount { - count - precision - } - createdAt - customerSelection { - __typename - ... on DiscountCustomerAll { - allCustomers - } - ... on DiscountCustomerSegments { - segments { - creationDate - id - lastEditDate - name - query - } - } - ... on DiscountCustomers { - customers { - id - } - } - } - discountClass - discountId - endsAt - errorHistory { - errorsFirstOccurredAt - firstOccurredAt - hasBeenSharedSinceLastError - hasSharedRecentErrors - } - hasTimelineComment - recurringCycleLimit - shareableUrls { - targetItemImage { - id - url - } - targetType - title - url - } - startsAt - status - title - totalSales { - amount - currencyCode - } - updatedAt - usageLimit - } - ... on DiscountCodeBasic { - appliesOncePerCustomer - asyncUsageCount - codes: codes(first: 250) { - nodes { - id - code - } - } - codesCount { - count - precision - } - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - customerGets { - appliesOnOneTimePurchase - appliesOnSubscription - items { - __typename - ... on AllDiscountItems { - allItems - } - ... on DiscountCollections { - collectionsFirst250: collections(first: 250) { - nodes { - id - } - } - } - ... on DiscountProducts { - productsFirst250: products(first: 250) { - nodes { - id - } - } - productVariantsFirst250: productVariants(first: 250) { - nodes { - id - } - } - } - } - value { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountOnQuantity { - effect { - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountPercentage { - percentage - } - } - quantity { - quantity - } - } - ... on DiscountPercentage { - percentage - } - } - } - customerSelection { - __typename - ... on DiscountCustomerAll { - allCustomers - } - ... on DiscountCustomerSegments { - segments { - creationDate - id - lastEditDate - name - query - } - } - ... on DiscountCustomers { - customers { - id - } - } - } - discountClass - endsAt - hasTimelineComment - minimumRequirement { - __typename - ... on DiscountMinimumQuantity { - greaterThanOrEqualToQuantity - } - ... on DiscountMinimumSubtotal { - greaterThanOrEqualToSubtotal { - amount - currencyCode - } - } - } - recurringCycleLimit - shareableUrls { - url - title - } - shortSummary - startsAt - status - summary - title - totalSales { - amount - currencyCode - } - updatedAt - usageLimit - } - ... on DiscountCodeBxgy { - appliesOncePerCustomer - asyncUsageCount - codesFirst50: codes(first: 50) { - nodes { - id - code - } - } - codesCount { - count - precision - } - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - customerBuys { - items { - __typename - ... on AllDiscountItems { - allItems - } - ... on DiscountCollections { - collectionsFirst250: collections(first: 250) { - nodes { - id - } - } - } - ... on DiscountProducts { - productsFirst250: products(first: 250) { - nodes { - id - } - } - productVariantsFirst250: productVariants(first: 250) { - nodes { - id - } - } - } - } - value { - __typename - ... on DiscountPurchaseAmount { - amount - } - ... on DiscountQuantity { - quantity - } - } - } - customerGets { - appliesOnOneTimePurchase - appliesOnSubscription - items { - __typename - ... on AllDiscountItems { - allItems - } - ... on DiscountCollections { - collectionsFirst250: collections(first: 250) { - nodes { - id - } - } - } - ... on DiscountProducts { - productsFirst250: products(first: 250) { - nodes { - id - } - } - productVariantsFirst250: productVariants(first: 250) { - nodes { - id - } - } - } - } - value { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountOnQuantity { - effect { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountPercentage { - percentage - } - } - quantity { - quantity - } - } - ... on DiscountPercentage { - percentage - } - } - } - customerSelection { - __typename - ... on DiscountCustomerAll { - allCustomers - } - ... on DiscountCustomerSegments { - segments { - creationDate - id - lastEditDate - name - query - } - } - ... on DiscountCustomers { - customers { - id - } - } - } - discountClass - endsAt - hasTimelineComment - shareableUrls { - url - title - } - startsAt - status - summary - title - totalSales { - amount - currencyCode - } - updatedAt - usageLimit - usesPerOrderLimit - } - ... on DiscountCodeFreeShipping { - appliesOncePerCustomer - appliesOnOneTimePurchase - appliesOnSubscription - asyncUsageCount - codesFirst250: codes(first: 250) { - nodes { - id - code - } - } - codesCount { - count - precision - } - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - customerSelection { - __typename - ... on DiscountCustomerAll { - allCustomers - } - ... on DiscountCustomerSegments { - segments { - creationDate - id - lastEditDate - name - query - } - } - ... on DiscountCustomers { - customers { - id - } - } - } - destinationSelection { - __typename - ... on DiscountCountries { - countries - includeRestOfWorld - } - ... on DiscountCountryAll { - allCountries - } - } - discountClass - endsAt - hasTimelineComment - maximumShippingPrice { - amount - currencyCode - } - minimumRequirement { - __typename - ... on DiscountMinimumQuantity { - greaterThanOrEqualToQuantity - } - ... on DiscountMinimumSubtotal { - greaterThanOrEqualToSubtotal { - amount - currencyCode - } - } - } - recurringCycleLimit - shareableUrls { - targetItemImage { - id - url - } - targetType - title - url - } - shortSummary - startsAt - status - summary - title - totalSales { - amount - currencyCode - } - updatedAt - usageLimit - } - ... on DiscountAutomaticApp { - appDiscountType { - app { - apiKey - } - appBridge { - createPath - detailsPath - } - appKey - description - discountClass - functionId - targetType - title - } - asyncUsageCount - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - discountClass - discountId - startsAt - endsAt - errorHistory { - errorsFirstOccurredAt - firstOccurredAt - hasBeenSharedSinceLastError - hasSharedRecentErrors - } - status - title - updatedAt - } - ... on DiscountAutomaticBasic { - asyncUsageCount - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - customerGets { - appliesOnOneTimePurchase - appliesOnSubscription - items { - __typename - ... on AllDiscountItems { - allItems - } - ... on DiscountCollections { - collectionsFirst250: collections(first: 250) { - nodes { - id - } - } - } - ... on DiscountProducts { - productsFirst250: products(first: 250) { - nodes { - id - } - } - productVariantsFirst250: productVariants(first: 250) { - nodes { - id - } - } - } - } - value { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountOnQuantity { - effect { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountPercentage { - percentage - } - } - quantity { - quantity - } - } - ... on DiscountPercentage { - percentage - } - } - } - discountClass - endsAt - minimumRequirement { - __typename - ... on DiscountMinimumQuantity { - greaterThanOrEqualToQuantity - } - ... on DiscountMinimumSubtotal { - greaterThanOrEqualToSubtotal { - amount - currencyCode - } - } - } - recurringCycleLimit - shortSummary - startsAt - status - summary - title - updatedAt - } - ... on DiscountAutomaticBxgy { - asyncUsageCount - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - customerBuys { - items { - __typename - ... on AllDiscountItems { - allItems - } - ... on DiscountCollections { - collectionsFirst250: collections(first: 250) { - nodes { - id - } - } - } - ... on DiscountProducts { - productsFirst250: products(first: 250) { - nodes { - id - } - } - productVariantsFirst250: productVariants(first: 250) { - nodes { - id - } - } - } - } - value { - __typename - ... on DiscountPurchaseAmount { - amount - } - ... on DiscountQuantity { - quantity - } - } - } - customerGets { - appliesOnOneTimePurchase - appliesOnSubscription - items { - __typename - ... on AllDiscountItems { - allItems - } - ... on DiscountCollections { - collectionsFirst250: collections(first: 250) { - nodes { - id - } - } - } - ... on DiscountProducts { - productsFirst250: products(first: 250) { - nodes { - id - } - } - productVariantsFirst250: productVariants(first: 250) { - nodes { - id - } - } - } - } - value { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountOnQuantity { - effect { - __typename - ... on DiscountAmount { - amount { - amount - currencyCode - } - appliesOnEachItem - } - ... on DiscountPercentage { - percentage - } } - quantity { - quantity - } - } - ... on DiscountPercentage { - percentage - } - } - } - discountClass - endsAt - startsAt - status - summary - title - updatedAt - usesPerOrderLimit - } - ... on DiscountAutomaticFreeShipping { - appliesOnOneTimePurchase - appliesOnSubscription - asyncUsageCount - combinesWith { - orderDiscounts - productDiscounts - shippingDiscounts - } - createdAt - destinationSelection - discountClass - endsAt - hasTimelineComment - maximumShippingPrice { - amount - currencyCode - } - minimumRequirement { - __typename - ... on DiscountMinimumQuantity { - greaterThanOrEqualToQuantity - } - ... on DiscountMinimumSubtotal { - greaterThanOrEqualToSubtotal { - amount - currencyCode - } - } - } - recurringCycleLimit - shortSummary - startsAt - status - summary - title - totalSales { - amount - currencyCode - } - updatedAt - } - } - metafieldsFirst250: metafields(first: 250) { - nodes { - id - key - value - } - } - } - pageInfo { - endCursor - hasNextPage - hasPreviousPage - startCursor - } - } -} -""" - - yield from client.get_graphql_pages( - query, - data_items_path="data.discountNodes.nodes[*]", - pagination_cursor_path="data.discountNodes.pageInfo.endCursor", - pagination_variable_name="after", - variables={ - "first": items_per_page, - }, - ) - - @dlt.resource(primary_key="id", write_disposition="merge") - def taxonomy(items_per_page: int = items_per_page) -> Iterable[TDataItem]: - client = ShopifyGraphQLApi( - base_url=shop_url, - access_token=private_app_password, - api_version="2024-07", - ) - - query = """ -{ - taxonomy { - categories(first: 250) { - nodes { - id - isArchived - isLeaf - isRoot - level - name - parentId - fullName - ancestorIds - attributes(first: 250) { - nodes { - ... on TaxonomyAttribute { - id - } - ... on TaxonomyChoiceListAttribute { - id - name - } - ... on TaxonomyMeasurementAttribute { - id - name } - } - } - } - pageInfo { - endCursor - hasNextPage - hasPreviousPage - startCursor - } - } - } -} -""" + }''' - yield from client.get_graphql_pages( + partner_query_pages( query, - data_items_path="data.taxonomy.categories.nodes[*]", - pagination_cursor_path="data.taxonomy.categories.pageInfo.endCursor", - pagination_cursor_has_next_page_path="data.taxonomy.categories.pageInfo.hasNextPage", + data_items_path="data.transactions.edges[*].node", + pagination_cursor_path="data.transactions.edges[-1].cursor", pagination_variable_name="after", - variables={ - "first": items_per_page, - }, ) - @dlt.resource( - primary_key="id", - write_disposition="merge", - columns={ - "id": { - "data_type": "text", - "nullable": False, - "primary_key": True, - "description": "A globally unique ID for the product.", - }, - "availablePublicationsCount": { - "data_type": "json", - "nullable": False, - "description": "The number of publications that a resource is published to", - }, - "category": { - "data_type": "json", - "nullable": True, - "description": "The category of the product from Shopify's Standard Product Taxonomy.", - }, - "compareAtPriceRange": { - "data_type": "json", - "nullable": True, - "description": "The compare-at price range of the product in the shop's default currency.", - }, - "createdAt": { - "data_type": "timestamp", - "nullable": False, - "description": "The date and time when the product was created.", - }, - "defaultCursor": { - "data_type": "text", - "nullable": False, - "description": "A default cursor that returns the next record sorted by ID.", - }, - "description": { - "data_type": "text", - "nullable": False, - "description": "A single-line description of the product, with HTML tags removed.", - }, - "descriptionHtml": { - "data_type": "text", - "nullable": False, - "description": "The description of the product, with HTML tags.", - }, - "handle": { - "data_type": "text", - "nullable": False, - "description": "A unique, human-readable string of the product's title.", - }, - "metafields": { - "data_type": "json", - "nullable": True, - "description": "A list of custom fields associated with the product.", - }, - "options": { - "data_type": "json", - "nullable": True, - "description": "A list of product options, e.g., size, color.", - }, - "priceRangeV2": { - "data_type": "json", - "nullable": False, - "description": "The minimum and maximum prices of a product.", - }, - "productType": { - "data_type": "text", - "nullable": False, - "description": "The product type defined by the merchant.", - }, - "publishedAt": { - "data_type": "timestamp", - "nullable": True, - "description": "The date and time when the product was published.", - }, - "requiresSellingPlan": { - "data_type": "bool", - "nullable": True, - "description": "Whether the product can only be purchased with a selling plan.", - }, - "status": { - "data_type": "text", - "nullable": False, - "description": "The product status, which controls visibility across all sales channels.", - }, - "tags": { - "data_type": "text", - "nullable": True, - "description": "A comma-separated list of searchable keywords associated with the product.", - }, - "templateSuffix": { - "data_type": "text", - "nullable": True, - "description": "The theme template used when customers view the product in a store.", - }, - "title": { - "data_type": "text", - "nullable": False, - "description": "The name for the product that displays to customers.", - }, - "totalInventory": { - "data_type": "bigint", - "nullable": False, - "description": "The quantity of inventory that's in stock.", - }, - "tracksInventory": { - "data_type": "bool", - "nullable": False, - "description": "Whether inventory tracking is enabled for the product.", - }, - "updatedAt": { - "data_type": "timestamp", - "nullable": False, - "description": "The date and time when the product was last modified.", - }, - "variantsFirst250": { - "data_type": "json", - "nullable": False, - "description": "A list of variants associated with the product, first 250.", - }, - "variantsCount": { - "data_type": "json", - "nullable": False, - "description": "The number of variants associated with the product.", - }, - "vendor": { - "data_type": "text", - "nullable": False, - "description": "The name of the product's vendor.", - }, - }, + Args: + query: The GraphQL query to run. + data_items_path: The JSONPath to the data items in the query result. Should resolve to array items. + pagination_cursor_path: The JSONPath to the pagination cursor in the query result, will be piped to the next query via variables. + pagination_variable_name: The name of the variable to pass the pagination cursor to. + variables: Mapping of extra variables used in the query. + access_token: The Partner API Client access token, created in the Partner Dashboard. + organization_id: Your Organization ID, found in the Partner Dashboard. + api_version: The API version to use (e.g. 2024-01). Use `unstable` for the latest version. + Returns: + Iterable[TDataItem]: A generator of the query results. + """ + client = ShopifyPartnerApi( + access_token=access_token, + organization_id=organization_id, + api_version=api_version, ) - def products( - updated_at: dlt.sources.incremental[ - pendulum.DateTime - ] = dlt.sources.incremental( - "updatedAt", - initial_value=start_date_obj, - end_value=end_date_obj, - range_end="closed", - range_start="closed", - ), - items_per_page: int = items_per_page, - ) -> Iterable[TDataItem]: - client = ShopifyGraphQLApi( - base_url=shop_url, - access_token=private_app_password, - api_version="2024-07", - ) - - query = """ -query products($after: String, $query: String, $first: Int) { - products(after: $after, first: $first, query: $query) { - nodes { - availablePublicationsCount { - count - precision - } - category { - id - } - compareAtPriceRange { - maxVariantCompareAtPrice { - amount - currencyCode - } - minVariantCompareAtPrice { - amount - currencyCode - } - } - createdAt - defaultCursor - description - descriptionHtml - handle - id - metafields(first: 250) { - nodes { - id - key - value - } - } - options { - linkedMetafield { - key - namespace - } - name - optionValues { - hasVariants - id - linkedMetafieldValue - name - } - values - id - position - } - priceRangeV2 { - maxVariantPrice { - amount - currencyCode - } - minVariantPrice { - amount - currencyCode - } - } - productType - publishedAt - requiresSellingPlan - status - tags - templateSuffix - totalInventory - title - tracksInventory - updatedAt - vendor - variantsCount { - count - precision - } - variantsFirst250: variants(first: 250) { - nodes { - id - sku - } - } - } - pageInfo { - endCursor - hasNextPage - hasPreviousPage - } - } -} -""" - - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Shopify") - - variables = { - "first": items_per_page, - "query": f"updated_at:>'{updated_at.last_value.isoformat()}'", - } - - yield from client.get_graphql_pages( - query, - data_items_path="data.products.nodes[*]", - pagination_cursor_path="data.products.pageInfo.endCursor", - pagination_cursor_has_next_page_path="data.products.pageInfo.hasNextPage", - pagination_variable_name="after", - variables=variables, - ) - return ( - products, - products_legacy, - orders, - customers, - inventory_items, - transactions, - balance, - events, - price_rules, - discounts, - taxonomy, + yield from client.get_graphql_pages( + query, + data_items_path=data_items_path, + pagination_cursor_path=pagination_cursor_path, + pagination_variable_name=pagination_variable_name, + variables=variables, ) diff --git a/omniload/source/shopify/exceptions.py b/omniload/source/shopify/exceptions.py index 91cf7c1c0..3f1062398 100644 --- a/omniload/source/shopify/exceptions.py +++ b/omniload/source/shopify/exceptions.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/shopify/helpers.py b/omniload/source/shopify/helpers.py index bde9900c4..743aebef8 100644 --- a/omniload/source/shopify/helpers.py +++ b/omniload/source/shopify/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,8 +18,8 @@ from urllib.parse import urljoin from dlt.common import jsonpath -from dlt.common.time import ensure_pendulum_datetime_utc -from dlt.common.typing import Dict, DictStrAny, TDataItems +from dlt.common.time import ensure_pendulum_datetime +from dlt.common.typing import Dict, DictStrAny, TDataItem, TDataItems from dlt.sources.helpers import requests from .exceptions import ShopifyPartnerApiError @@ -28,53 +28,6 @@ TOrderStatus = Literal["open", "closed", "cancelled", "any"] -def convert_datetime_fields(item: Dict[str, Any]) -> Dict[str, Any]: - """Convert timestamp fields in the item to pendulum datetime objects - - The item is modified in place, including nested items. - - Args: - item: The item to convert - - Returns: - The same data item (for convenience) - """ - fields = ["created_at", "updated_at", "createdAt", "updatedAt"] - - def convert_nested(obj: Any) -> Any: - if isinstance(obj, dict): - for key, value in obj.items(): - if key in fields and isinstance(value, str): - obj[key] = ensure_pendulum_datetime_utc(value) - else: - obj[key] = convert_nested(value) - elif isinstance(obj, list): - return [convert_nested(elem) for elem in obj] - return obj - - return convert_nested(item) - - -def remove_nodes_key(item: Any) -> Any: - """ - Recursively remove the 'nodes' key from dictionaries if it's the only key and its value is an array. - - Args: - item: The item to process (can be a dict, list, or any other type) - - Returns: - The processed item - """ - if isinstance(item, dict): - if len(item) == 1 and "nodes" in item and isinstance(item["nodes"], list): - return [remove_nodes_key(node) for node in item["nodes"]] - return {k: remove_nodes_key(v) for k, v in item.items()} - elif isinstance(item, list): - return [remove_nodes_key(element) for element in item] - else: - return item - - class ShopifyApi: """ A Shopify API client that can be used to get pages of data from Shopify. @@ -111,38 +64,57 @@ def get_pages( """ url = urljoin(self.shop_url, f"/admin/api/{self.api_version}/{resource}.json") - resource_last = resource.split("/")[-1] - headers = {"X-Shopify-Access-Token": self.private_app_password} while url: response = requests.get(url, params=params, headers=headers) response.raise_for_status() json = response.json() - yield [convert_datetime_fields(item) for item in json[resource_last]] + # Get item list from the page + yield [self._convert_datetime_fields(item) for item in json[resource]] url = response.links.get("next", {}).get("url") # Query params are included in subsequent page URLs params = None + def _convert_datetime_fields(self, item: Dict[str, Any]) -> Dict[str, Any]: + """Convert timestamp fields in the item to pendulum datetime objects + + The item is modified in place. -class ShopifyGraphQLApi: - """Client for Shopify GraphQL API""" + Args: + item: The item to convert + + Returns: + The same data item (for convenience) + """ + fields = ["created_at", "updated_at"] + for field in fields: + if field in item: + item[field] = ensure_pendulum_datetime(item[field]) + return item + + +class ShopifyPartnerApi: + """Client for Shopify Partner grapql API""" def __init__( self, access_token: str, + organization_id: str, api_version: str = DEFAULT_PARTNER_API_VERSION, - base_url: str = "partners.shopify.com", ) -> None: + """ + Args: + access_token: The access token to use + organization_id: The organization id to query + api_version: The API version to use (e.g. 2023-01) + """ self.access_token = access_token + self.organization_id = organization_id self.api_version = api_version - self.base_url = base_url @property def graphql_url(self) -> str: - if self.base_url.startswith("https://"): - return f"{self.base_url}/admin/api/{self.api_version}/graphql.json" - - return f"https://{self.base_url}/admin/api/{self.api_version}/graphql.json" + return f"https://partners.shopify.com/{self.organization_id}/api/{self.api_version}/graphql.json" def run_graphql_query( self, query: str, variables: Optional[DictStrAny] = None @@ -173,30 +145,16 @@ def get_graphql_pages( data_items_path: jsonpath.TJsonPath, pagination_cursor_path: jsonpath.TJsonPath, pagination_variable_name: str, - pagination_cursor_has_next_page_path: Optional[jsonpath.TJsonPath] = None, variables: Optional[DictStrAny] = None, ) -> Iterable[TDataItems]: variables = dict(variables or {}) while True: data = self.run_graphql_query(query, variables) data_items = jsonpath.find_values(data_items_path, data) - if not data_items: break - - yield [ - remove_nodes_key(convert_datetime_fields(item)) for item in data_items - ] - + yield data_items cursors = jsonpath.find_values(pagination_cursor_path, data) if not cursors: break - - if pagination_cursor_has_next_page_path: - has_next_page = jsonpath.find_values( - pagination_cursor_has_next_page_path, data - ) - if not has_next_page or not has_next_page[0]: - break - variables[pagination_variable_name] = cursors[-1] diff --git a/omniload/source/shopify/settings.py b/omniload/source/shopify/settings.py index 6c7490525..04986cce1 100644 --- a/omniload/source/shopify/settings.py +++ b/omniload/source/shopify/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/slack/adapter.py b/omniload/source/slack/adapter.py index 7624f5c19..c624babf7 100644 --- a/omniload/source/slack/adapter.py +++ b/omniload/source/slack/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,7 +31,7 @@ ) -@dlt.source(name="slack", max_table_nesting=0) +@dlt.source(name="slack", max_table_nesting=2) def slack_source( page_size: int = MAX_PAGE_SIZE, access_token: str = dlt.secrets.value, @@ -40,6 +40,7 @@ def slack_source( selected_channels: Optional[List[str]] = dlt.config.value, table_per_channel: bool = True, replies: bool = False, + include_private_channels: bool = False, ) -> Iterable[DltResource]: """ The source for the Slack pipeline. Available resources are conversations, conversations_history @@ -54,6 +55,8 @@ def slack_source( table_per_channel: Boolean flag, True by default. If True - for each channel separate table with messages is created. Otherwise, all messages are put in one table. replies: Boolean flag indicating if you want a replies table to be present as well. False by default. + include_private_channels: Boolean flag indicating if you want to include private channels and group DMs. + Defaults to False. Requires appropriate OAuth scopes (groups:read, mpim:read). Returns: Iterable[DltResource]: A list of DltResource objects representing the data resources. @@ -71,7 +74,9 @@ def slack_source( ) def get_channels( - slack_api: SlackAPI, selected_channels: Optional[List[str]] + slack_api: SlackAPI, + selected_channels: Optional[List[str]], + include_private_channels: bool = False, ) -> Tuple[List[TDataItem], List[TDataItem]]: """ Returns channel fetched from slack and list of selected channels. @@ -79,15 +84,24 @@ def get_channels( Args: slack_api: Slack API instance. selected_channels: List of selected channels names or None. + include_private_channels: Whether to include private channels and group DMs. Returns: Tuple[List[TDataItem], List[TDataItem]]: fetched channels and selected fetched channels. """ channels: List[TDataItem] = [] + + # Define conversation types based on the include_private_channels parameter + if include_private_channels: + conversation_types = "public_channel,private_channel,mpim" + else: + conversation_types = "public_channel" + for page_data in slack_api.get_pages( resource="conversations.list", response_path="$.channels[*]", datetime_fields=DEFAULT_DATETIME_FIELDS, + params={"types": conversation_types}, ): channels.extend(page_data) @@ -101,7 +115,9 @@ def get_channels( fetch_channels = channels return channels, fetch_channels - channels, fetched_selected_channels = get_channels(api, selected_channels) + channels, fetched_selected_channels = get_channels( + api, selected_channels, include_private_channels + ) @dlt.resource(name="channels", primary_key="id", write_disposition="replace") def channels_resource() -> Iterable[TDataItem]: @@ -189,8 +205,6 @@ def messages_resource( initial_value=start_dt, end_value=end_dt, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), ) -> Iterable[TDataItem]: """ @@ -214,8 +228,6 @@ def per_table_messages_resource( initial_value=start_dt, end_value=end_dt, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), ) -> Iterable[TDataItem]: """Yield all messages for a given channel as a DLT resource. Keep blocks column without normalization. diff --git a/omniload/source/slack/helpers.py b/omniload/source/slack/helpers.py index 2be202545..90f837acf 100644 --- a/omniload/source/slack/helpers.py +++ b/omniload/source/slack/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,14 +14,15 @@ """Slack source helpers.""" +import logging from typing import Any, Generator, Iterable, List, Optional from urllib.parse import urljoin import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import Dict, TAnyDateTime, TDataItem from dlt.sources.helpers import requests -from jsonpath_ng.ext import parse +from jsonpath_ng.ext import parse # type: ignore from .settings import MAX_PAGE_SIZE, SLACK_API_URL @@ -55,9 +56,7 @@ def update_jsonpath(expression: str, json_data: TDataItem, value: Any) -> Any: return jsonpath.update_or_create(json_data, value) -def ensure_dt_type( - dt: Optional[TAnyDateTime] = None, to_ts: Optional[bool] = False -) -> Any: +def ensure_dt_type(dt: TAnyDateTime, to_ts: bool = False) -> Any: """Converts a datetime to a pendulum datetime or timestamp. Args: dt: The datetime to convert. @@ -67,7 +66,7 @@ def ensure_dt_type( """ if dt is None: return None - out_dt = ensure_pendulum_datetime_utc(dt) + out_dt = ensure_pendulum_datetime(dt) if to_ts: return out_dt.timestamp() return out_dt @@ -97,7 +96,7 @@ def headers(self) -> Dict[str, str]: return {"Authorization": f"Bearer {self.access_token}"} def parameters( - self, params: Optional[Dict[str, Any]] = None, next_cursor: Optional[str] = None + self, params: Optional[Dict[str, Any]] = None, next_cursor: str = None ) -> Dict[str, str]: """ Generate the query parameters to use for the request. @@ -135,9 +134,7 @@ def _get_next_cursor(self, response: Dict[str, Any]) -> Any: return next(extract_jsonpath(cursor_jsonpath, response), None) def _convert_datetime_fields( - self, - item: Dict[str, Any], - datetime_fields: Optional[List[str]] = None, + self, item: Dict[str, Any], datetime_fields: List[str] ) -> Dict[str, Any]: """Convert timestamp fields in the item to pendulum datetime objects. @@ -166,10 +163,10 @@ def _convert_datetime_fields( def get_pages( self, resource: str, - response_path: str, - params: Optional[Dict[str, Any]] = None, - datetime_fields: Optional[List[str]] = None, - context: Optional[Dict[str, Any]] = None, + response_path: str = None, + params: Dict[str, Any] = None, + datetime_fields: List[str] = None, + context: Dict[str, Any] = None, ) -> Iterable[TDataItem]: """Get all pages from slack using requests. Iterates through all pages and yield each page items.\ diff --git a/omniload/source/slack/settings.py b/omniload/source/slack/settings.py index 2b0c2c7f6..560801faa 100644 --- a/omniload/source/slack/settings.py +++ b/omniload/source/slack/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/omniload/source/strapi/adapter.py b/omniload/source/strapi/adapter.py new file mode 100644 index 000000000..275074689 --- /dev/null +++ b/omniload/source/strapi/adapter.py @@ -0,0 +1,49 @@ +# Copyright 2022-2026 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. + +""" +Basic strapi source +""" + +from typing import Iterable, List + +import dlt +from dlt.sources import DltResource + +from .helpers import get_endpoint + + +@dlt.source +def strapi_source( + endpoints: List[str], + api_secret_key: str = dlt.secrets.value, + domain: str = dlt.secrets.value, +) -> Iterable[DltResource]: + """ + Source function for retrieving data from Strapi. + + Args: + endpoints (List[str]): List of collections to retrieve data from. + api_secret_key (str): API secret key for authentication. Defaults to the value in the `dlt.secrets` object. + domain (str): Domain name for the Strapi API. Defaults to the value in the `dlt.secrets` object. + + Yields: + DltResource: Data resources from the specified collections. + """ + for endpoint in endpoints: + yield dlt.resource( # type: ignore + get_endpoint(api_secret_key, domain, endpoint), + name=endpoint, + write_disposition="replace", + ) diff --git a/omniload/source/strapi/helpers.py b/omniload/source/strapi/helpers.py new file mode 100644 index 000000000..f771aa12d --- /dev/null +++ b/omniload/source/strapi/helpers.py @@ -0,0 +1,56 @@ +# Copyright 2022-2026 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. + +"""Strapi source helpers""" + +import math +from typing import Iterable + +from dlt.common.typing import TDataItem +from dlt.sources.helpers import requests + + +def get_endpoint(token: str, domain: str, endpoint: str) -> Iterable[TDataItem]: + """ + A generator that yields data from a paginated API endpoint. + + Args: + token (str): The access token for the API. + domain (str): The domain name of the API. + endpoint (str): The API endpoint to query, defaults to ''. + + Yields: + TDataItem: A data item from the API endpoint. + """ + api_endpoint = f"https://{domain}/api/{endpoint}" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {token}"} + page_size = 25 + params = { + "pagination[start]": 0, + "pagination[limit]": page_size, + "pagination[withCount]": 1, + } + + # get the total number of pages + response = requests.get(api_endpoint, headers=headers, params=params) + total_results = response.json()["meta"]["pagination"]["total"] + pages_total = math.ceil(total_results / page_size) + + # yield page by page + for page_number in range(pages_total): + params["pagination[start]"] = page_number * page_size + response = requests.get(api_endpoint, headers=headers, params=params) + data = response.json().get("data") + if data: + yield from data diff --git a/omniload/source/strapi/settings.py b/omniload/source/strapi/settings.py new file mode 100644 index 000000000..d1965fde9 --- /dev/null +++ b/omniload/source/strapi/settings.py @@ -0,0 +1,15 @@ +# Copyright 2022-2026 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. + +"""Strapi source settings and constants""" diff --git a/omniload/source/stripe/adapter.py b/omniload/source/stripe/adapter.py index b5fbc4477..a2c052dd7 100644 --- a/omniload/source/stripe/adapter.py +++ b/omniload/source/stripe/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,7 @@ # 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.""" +"""This source uses Stripe API and dlt to load data such as Customer, Subscription, Event etc. to the database.""" from typing import Any, Dict, Generator, Iterable, Optional, Tuple @@ -21,17 +21,13 @@ from dlt.sources import DltResource from pendulum import DateTime -from .helpers import ( - async_parallel_pagination, - generate_date_ranges, - pagination, - transform_date, -) +from .helpers import pagination, transform_date +from .settings import ENDPOINTS, INCREMENTAL_ENDPOINTS -@dlt.source(max_table_nesting=0) +@dlt.source def stripe_source( - endpoints: Tuple[str, ...], + endpoints: Tuple[str, ...] = ENDPOINTS, stripe_secret_key: str = dlt.secrets.value, start_date: Optional[DateTime] = None, end_date: Optional[DateTime] = None, @@ -69,116 +65,49 @@ def stripe_resource( )(endpoint) -@dlt.source(max_table_nesting=0) -def async_stripe_source( - endpoints: Tuple[str, ...], +@dlt.source +def incremental_stripe_source( + endpoints: Tuple[str, ...] = INCREMENTAL_ENDPOINTS, stripe_secret_key: str = dlt.secrets.value, - start_date: Optional[DateTime] = None, + initial_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. + As Stripe API does not include the "updated" key in its responses, + we are only able to perform incremental downloads from endpoints where all objects are uneditable. + This source yields the resources with incremental loading based on "append" mode. + You will load only the newest data without duplicating and without downloading a huge amount of data each time. Args: - endpoints (Tuple[str, ...]): A tuple of endpoint names to retrieve data from. + endpoints (tuple): A tuple of endpoint names to retrieve data from. Defaults to Stripe API endpoints with uneditable data. 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. - + initial_start_date (Optional[DateTime]): An optional parameter that specifies the initial value for dlt.sources.incremental. + If parameter is not None, then load only data that were created after initial_start_date on the first run. + Defaults to None. Format: datetime(YYYY, MM, DD). + end_date (Optional[DateTime]): An optional end date to limit the data retrieved. + Defaults to None. Format: datetime(YYYY, MM, DD). Returns: - Iterable[DltResource]: Resources with data in RANDOM ORDER (optimized for speed). + Iterable[DltResource]: Resources with only that data has not yet been loaded. """ 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 + ) - async def async_stripe_resource(endpoint: str): - yield async_parallel_pagination(endpoint, max_workers, rate_limit_delay) + def incremental_resource( + endpoint: str, + created: Optional[Any] = dlt.sources.incremental( + "created", initial_value=start_date_unix + ), + ) -> Generator[Dict[Any, Any], Any, None]: + start_value = created.last_value + yield from pagination(endpoint, start_date=start_value, end_date=end_date) for endpoint in endpoints: yield dlt.resource( - async_stripe_resource, + incremental_resource, name=endpoint, - write_disposition="replace", + write_disposition="append", + primary_key="id", )(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 is not None and 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 is not None and created.end_value is not None - else int(pendulum.now().timestamp()) - ) - date_range: dict[str, Any] - 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, Any], - ) -> 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, - ) - ) diff --git a/omniload/source/stripe/helpers.py b/omniload/source/stripe/helpers.py index 3feae6478..a8b5d6912 100644 --- a/omniload/source/stripe/helpers.py +++ b/omniload/source/stripe/helpers.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,13 +14,10 @@ """Stripe analytics source helpers""" -import asyncio -import math -from datetime import datetime, timedelta -from typing import Any, Dict, Iterable, List, Optional, Union +from typing import Any, Dict, Iterable, Optional, Union import stripe -from dlt.common import pendulum +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import TDataItem from pendulum import DateTime @@ -56,264 +53,9 @@ def pagination( 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 _ 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]: # ty: ignore[invalid-return-type] - """ - 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]: # ty: ignore[invalid-return-type] - """ - 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): # ty: ignore[not-iterable] - chunk_data.append(batch) - - return chunk_data - - -def generate_date_ranges(start_ts: int, end_ts: int) -> Iterable[Dict[str, Any]]: - """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 + # convert ISO 8601 strings, datetimes and unix timestamps to a unix timestamp + return int(ensure_pendulum_datetime(date).timestamp()) def stripe_get_data( @@ -333,54 +75,6 @@ def stripe_get_data( 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) + # to_dict works across stripe-python versions, dict() conversion broke in v15 + # when StripeObject stopped inheriting from dict + return resource_dict.to_dict() # type: ignore[no-any-return] diff --git a/omniload/source/stripe/settings.py b/omniload/source/stripe/settings.py index 1258ea505..9d9576229 100644 --- a/omniload/source/stripe/settings.py +++ b/omniload/source/stripe/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,65 +16,14 @@ # the most popular endpoints # Full list of the Stripe API endpoints you can find here: https://stripe.com/docs/api. -ENDPOINTS = { - "account": "Account", - "applepaydomain": "ApplePayDomain", - "apple_pay_domain": "ApplePayDomain", - "applicationfee": "ApplicationFee", - "application_fee": "ApplicationFee", - "checkoutsession": "CheckoutSession", - "checkout_session": "CheckoutSession", - "coupon": "Coupon", - "charge": "Charge", - "customer": "Customer", - "dispute": "Dispute", - "paymentintent": "PaymentIntent", - "payment_intent": "PaymentIntent", - "paymentlink": "PaymentLink", - "payment_link": "PaymentLink", - "paymentmethod": "PaymentMethod", - "payment_method": "PaymentMethod", - "paymentmethoddomain": "PaymentMethodDomain", - "payment_method_domain": "PaymentMethodDomain", - "payout": "Payout", - "plan": "Plan", - "price": "Price", - "product": "Product", - "promotioncode": "PromotionCode", - "promotion_code": "PromotionCode", - "quote": "Quote", - "refund": "Refund", - "review": "Review", - "setupattempt": "SetupAttempt", - "setup_attempt": "SetupAttempt", - "setupintent": "SetupIntent", - "setup_intent": "SetupIntent", - "shippingrate": "ShippingRate", - "shipping_rate": "ShippingRate", - "subscription": "Subscription", - "subscriptionitem": "SubscriptionItem", - "subscription_item": "SubscriptionItem", - "subscriptionschedule": "SubscriptionSchedule", - "subscription_schedule": "SubscriptionSchedule", - "transfer": "Transfer", - "taxcode": "TaxCode", - "tax_code": "TaxCode", - "taxid": "TaxId", - "tax_id": "TaxId", - "taxrate": "TaxRate", - "tax_rate": "TaxRate", - "topup": "Topup", - "top_up": "Topup", - "webhookendpoint": "WebhookEndpoint", - "webhook_endpoint": "WebhookEndpoint", - "invoice": "Invoice", - "invoiceitem": "InvoiceItem", - "invoice_item": "InvoiceItem", - "invoicelineitem": "InvoiceLineItem", - "invoice_line_item": "InvoiceLineItem", - "balancetransaction": "BalanceTransaction", - "balance_transaction": "BalanceTransaction", - "creditnote": "CreditNote", - "credit_note": "CreditNote", - "event": "Event", -} +ENDPOINTS = ( + "Subscription", + "Account", + "Coupon", + "Customer", + "Invoice", + "Product", + "Price", +) +# possible incremental endpoints +INCREMENTAL_ENDPOINTS = ("Event", "BalanceTransaction") diff --git a/omniload/source/unstructured/adapter.py b/omniload/source/unstructured/adapter.py new file mode 100644 index 000000000..17b712d8f --- /dev/null +++ b/omniload/source/unstructured/adapter.py @@ -0,0 +1,116 @@ +# Copyright 2022-2026 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 converts unstructured data from a specified data resource to structured data using provided queries.""" + +import asyncio +import os +from typing import Dict, Optional + +import dlt +from dlt.common import logger +from dlt.sources import DltResource, TDataItem + +from .helpers import ( + aprocess_file_to_structured, + process_file_to_structured, + vectorstore_mapping, +) +from .settings import INVOICE_QUERIES + + +@dlt.resource +def unstructured_to_structured_resource( + queries: Optional[Dict[str, str]] = dlt.config.value, + openai_api_key: str = dlt.secrets.value, + vectorstore: str = "chroma", + table_name: str = "unstructured_to_structured_resource", + run_async: bool = False, +) -> DltResource: + """ + Converts unstructured data from a specified data item to structured data using provided queries. + + Args: + queries (Dict[str, str]): A dictionary of queries to be applied to the unstructured data during processing. + Each query maps a field name to a query string that specifies how to process the field. + openai_api_key (str): The API key for the OpenAI API. If provided, it sets the `OPENAI_API_KEY` environment variable. + Defaults to the value of `dlt.secrets.value`. + vectorstore (str): Vector database type, e.g. "chroma", "weaviate" (expects environment variable `WEAVIATE_URL`) + or "elastic_search" (expects environment variable `ELASTICSEARCH_URL`). Defaults to "chroma". + table_name (str): The name of the table associated with the resource. Defaults to "unstructured_to_structured_resource". + run_async (bool): Whether to run the conversion asynchronously. Defaults to False. + + Returns: + DltResource: A resource-transformer object representing the conversion of unstructured data to structured data. + + """ + if openai_api_key: + os.environ["OPENAI_API_KEY"] = openai_api_key + if queries is None: + queries = dict(INVOICE_QUERIES) + + return dlt.transformer( + convert_data, + name=table_name, + write_disposition="merge", + merge_key="metadata__data_hash", + primary_key="metadata__data_hash", + )(queries, vectorstore, run_async) + + +def convert_data( + unstructured_item: TDataItem, + queries: Dict[str, str], + vectorstore: str = "chroma", + run_async: bool = False, +) -> TDataItem: + """ + Converts unstructured data item to structured data item using provided queries. + + Args: + unstructured_item (TDataItem): The data item containing unstructured data to be converted. + queries (Dict[str, str]): A dictionary of queries to be applied to the unstructured data during processing. + Each query maps a field name to a query string that specifies how to process the field. + vectorstore (str): Vector database type, e.g. "chroma", "weaviate" or "elastic_search". Default to "chroma". + run_async (bool): Whether to run the conversion asynchronously. Defaults to False. + Returns: + TDataItem: The structured data item resulting from the conversion. + + """ + if unstructured_item.get("file_path") is None: + return None + try: + if run_async: + logger.info("Run conversion asynchronously.") + response = asyncio.run( + aprocess_file_to_structured( + unstructured_item["file_path"], + queries, + vectorstore_mapping[vectorstore], + ) + ) + else: + response = process_file_to_structured( + unstructured_item["file_path"], + queries, + vectorstore_mapping[vectorstore], + ) + response["file_path"] = unstructured_item.pop("file_path") + response["metadata"] = unstructured_item + yield response + + except ValueError as error: + logger.warning( + f"File {unstructured_item['file_path']} has unsupported format: {error}" + ) diff --git a/omniload/source/unstructured/async_index.py b/omniload/source/unstructured/async_index.py new file mode 100644 index 000000000..c0ad8316f --- /dev/null +++ b/omniload/source/unstructured/async_index.py @@ -0,0 +1,56 @@ +# Copyright 2022-2026 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, List, Optional + +from langchain.base_language import BaseLanguageModel +from langchain.chains.retrieval_qa.base import RetrievalQA +from langchain.indexes.vectorstore import ( + VectorstoreIndexCreator, + VectorStoreIndexWrapper, +) +from langchain.llms.openai import OpenAI +from langchain.schema import Document + + +class AVectorStoreIndexWrapper(VectorStoreIndexWrapper): + """Async wrapper around a vectorstore for easy access.""" + + def __init__( + self, + **kwargs: Any, + ): + super().__init__(**kwargs) + + async def aquery( + self, question: str, llm: Optional[BaseLanguageModel] = None, **kwargs: Any + ) -> str: + """Query the vectorstore.""" + llm = llm or OpenAI(temperature=0) # type: ignore[call-arg] + chain = RetrievalQA.from_chain_type( + llm, retriever=self.vectorstore.as_retriever(), **kwargs + ) + return str(await chain.arun(question)) + + +class AVectorstoreIndexCreator(VectorstoreIndexCreator): + """Async logic for creating indexes.""" + + def from_documents(self, documents: List[Document]) -> AVectorStoreIndexWrapper: + """Create a vectorstore index from documents.""" + sub_docs = self.text_splitter.split_documents(documents) + vectorstore = self.vectorstore_cls.from_documents( + sub_docs, self.embedding, **self.vectorstore_kwargs + ) + return AVectorStoreIndexWrapper(vectorstore=vectorstore) diff --git a/omniload/source/unstructured/helpers.py b/omniload/source/unstructured/helpers.py new file mode 100644 index 000000000..222cb2531 --- /dev/null +++ b/omniload/source/unstructured/helpers.py @@ -0,0 +1,105 @@ +# Copyright 2022-2026 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. + +import asyncio +from typing import Any, Dict, List, Mapping, Tuple, Type, Union + +from langchain.document_loaders import UnstructuredFileLoader +from langchain.indexes import VectorstoreIndexCreator +from langchain.vectorstores.base import VectorStore +from langchain.vectorstores.chroma import Chroma +from langchain.vectorstores.elastic_vector_search import ElasticVectorSearch +from langchain.vectorstores.weaviate import Weaviate + +from .async_index import AVectorstoreIndexCreator + +vectorstore_mapping: Mapping[str, Type[VectorStore]] = { + "chroma": Chroma, + "elastic_search": ElasticVectorSearch, + "weaviate": Weaviate, +} + + +def safely_query_index(index: Any, query: str) -> Any: + answer = index.query(query) + return answer.strip() + + +async def asafely_query_index(index: Any, query: str) -> Any: + answer = await index.aquery(query) + return answer.strip() + + +async def aprocess_file_to_structured( + file_path: Union[str, List[str]], + queries: Dict[str, str], + vectorstore: Type[VectorStore] = Chroma, +) -> Dict[str, Any]: + """ + Async processes a file loaded by the specified loader and generates structured data based on provided queries. + + Args: + file_path (Union[str, List[str]]): filepath to the file with unstructured data. + queries (Dict[str, str]): A dictionary of queries to be applied to the loaded file. + Each query maps a field name to a query string that specifies how to process the field. + vectorstore (Type[VectorStore]): Vector database type. Subclass of VectorStore. Default to Chroma. + + Returns: + Dict[str, str]: A dictionary containing the processed structured data from the loaded file. + The dictionary includes a "file_path" key with the path of the loaded file and + additional keys corresponding to the queried fields and their processed values. + """ + loader = UnstructuredFileLoader(file_path) + index = AVectorstoreIndexCreator(vectorstore_cls=vectorstore).from_loaders([loader]) + + async def mark(key: str, question: str) -> Tuple[str, str]: + return key, await asafely_query_index(index, question) + + response = { + key: result + for key, result in await asyncio.gather( + *(mark(key, question) for key, question in queries.items()) + ) + } + + return response + + +def process_file_to_structured( + file_path: Union[str, List[str]], + queries: Dict[str, str], + vectorstore: Type[VectorStore] = Chroma, +) -> Dict[str, Any]: + """ + Processes a file loaded by the specified loader and generates structured data based on provided queries. + + Args: + file_path (Union[str, List[str]]): filepath to the file with unstructured data. + queries (Dict[str, str]): A dictionary of queries to be applied to the loaded file. + Each query maps a field name to a query string that specifies how to process the field. + vectorstore (Type[VectorStore]): Vector database type. Subclass of VectorStore. Default to Chroma. + + Returns: + Dict[str, str]: A dictionary containing the processed structured data from the loaded file. + The dictionary includes a "file_path" key with the path of the loaded file and + additional keys corresponding to the queried fields and their processed values. + """ + loader = UnstructuredFileLoader(file_path) + index = VectorstoreIndexCreator(vectorstore_cls=vectorstore).from_loaders([loader]) + response = {} + + for k, query in queries.items(): + response[k] = safely_query_index(index, query) + + return response diff --git a/omniload/source/unstructured/settings.py b/omniload/source/unstructured/settings.py new file mode 100644 index 000000000..3028c604b --- /dev/null +++ b/omniload/source/unstructured/settings.py @@ -0,0 +1,22 @@ +# Copyright 2022-2026 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. + +INVOICE_QUERIES = { + "recipient_company_name": "Who is the recipient of the invoice? Just return the name. If you don't know, then return None", + "invoice_amount": "What is the total amount of the invoice? Just return the amount as decimal number, no currency or text. If you don't know, then return None", + "invoice_date": "What is the date of the invoice? Just return the date. If you don't know, then return None", + "invoice_number": "What is the invoice number? Just return the number. If you don't know, then return None", + "service_description": "What is the description of the service that this invoice is for? Just return the description. If you don't know, then return None", + "phone_number": "What is the company phone number? Just return the phone number. If you don't know, then return None", +} diff --git a/omniload/source/workable/adapter.py b/omniload/source/workable/adapter.py new file mode 100644 index 000000000..cb6a5d61e --- /dev/null +++ b/omniload/source/workable/adapter.py @@ -0,0 +1,133 @@ +# Copyright 2022-2026 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 Workable API and dlt to load data such as Candidates, Jobs, Events, etc. to the database.""" + +import logging +from typing import Any, Iterable, Optional + +import dlt +from dlt.sources import DltResource, TDataItem, TDataItems +from pendulum import DateTime + +from .settings import DEFAULT_DETAILS, DEFAULT_ENDPOINTS +from .workable_client import WorkableClient + +logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO) + +""" +To set up authorization for the Workable API, you will need an API access token and your subdomain name. + +Follow these steps: + +1. Log in to your Workable account and navigate to the API section of the Integrations page. +2. Click the "Generate new token" button to create a new API access token. +3. Copy the generated token to your clipboard. +4. Open .dlt/secrets.toml and write the token in the section [sources.workable] +5. To access data from your own Workable account, set the subdomain in the URL + to your subdomain name, for example: https://yoursubdomain.workable.com/api/v3/. +6. Write your subdomain name in .dlt/config.toml in the section [sources.workable] +""" + + +@dlt.source(name="workable") +def workable_source( + access_token: str = dlt.secrets.value, + subdomain: str = dlt.config.value, + start_date: Optional[DateTime] = None, + load_details: bool = False, +) -> Iterable[DltResource]: + """ + Retrieves data from the Workable API for the specified endpoints (DEFAULT_ENDPOINTS + 'candidates'). + For almost all endpoints, Workable API responses do not provide keys "updated_at", + so in most cases we are forced to load the date in 'replace' mode. + 'Candidates' are the only endpoints that have a key 'updated_at', which means that we can update the data incrementally. + + Resources that depend on another resource are implemented as transformers, + so they can re-use the original resource data without re-downloading. + + Args: + access_token (str): The API access token for authentication. Defaults to the value in the `dlt.secrets` object. + subdomain (str): The subdomain name for the Workable account. Defaults to the value in the `dlt.config` object. + start_date (Optional[DateTime]): An optional start date to limit the data retrieved. Defaults to January 1, 2000. + It does not affect dependent resources (jobs_activities, candidates_activities, etc). + load_details (bool): A boolean flag enables loading data from endpoints that depend on the main endpoints. + """ + resources = {} + workable = WorkableClient(access_token, subdomain, start_date=start_date) + params = {"created_after": workable.start_date_iso} + + # create resource for each endpoint + for endpoint in DEFAULT_ENDPOINTS: + # define resource + @dlt.resource(name=endpoint, write_disposition="replace") + def endpoint_resource(endpoint: str = endpoint) -> Iterable[TDataItem]: + logging.info( + f"Loading data from '{endpoint}' by 'created_at' in 'replace' mode." + ) + yield workable.pagination(endpoint=endpoint, params=params) + + # save for later and yield + resources[endpoint] = endpoint_resource + yield endpoint_resource + + @dlt.resource(name="candidates", write_disposition="merge", primary_key="id") + def candidates_resource( + updated_at: Optional[Any] = dlt.sources.incremental( + "updated_at", initial_value=workable.start_date_iso + ), + ) -> Iterable[TDataItem]: + """ + The 'updated_at' parameter is managed by the dlt.sources.incremental method. + This function is suitable only for the 'candidates' endpoint in incremental mode. + """ + logging.info( + "Fetching data from 'candidates' by 'updated_at'. Loading modified and new data." + ) + yield workable.pagination( + endpoint="candidates", params={"updated_after": updated_at.last_value} + ) + + yield candidates_resource + + if load_details: + + def _get_details( + page: Iterable[TDataItem], + main_endpoint: str, + sub_endpoint_name: str, + code_key: str, + ) -> Iterable[TDataItem]: + for item in page: + yield workable.details_from_endpoint( + main_endpoint, item[code_key], sub_endpoint_name + ) + + # A transformer functions that yield the activities, questions, etc. for each job. + for sub_endpoint in DEFAULT_DETAILS["jobs"]: + logging.info( + f"Loading additional data for 'jobs' from '{sub_endpoint}' in 'replace' mode." + ) + yield resources["jobs"] | dlt.transformer( + name=f"jobs_{sub_endpoint}", write_disposition="replace" + )(_get_details)("jobs", sub_endpoint, "shortcode") + + # A transformer functions that yield the activities and offers for each candidate. + for sub_endpoint in DEFAULT_DETAILS["candidates"]: + logging.info( + f"Loading additional data for 'candidates' from '{sub_endpoint}' in 'append' mode." + ) + yield candidates_resource | dlt.transformer( + name=f"candidates_{sub_endpoint}", write_disposition="append" + )(_get_details)("candidates", sub_endpoint, "id") diff --git a/omniload/source/workable/client.py b/omniload/source/workable/client.py new file mode 100644 index 000000000..14ed33692 --- /dev/null +++ b/omniload/source/workable/client.py @@ -0,0 +1,110 @@ +# Copyright 2022-2026 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. + +import logging +import time +from typing import Any, Dict, Iterable, Optional + +import pendulum +from dlt.common.typing import TDataItem +from dlt.sources.helpers import requests + + +class WorkableClient: + def __init__( + self, + access_token: str, + subdomain: str, + start_date: Optional[pendulum.DateTime] = None, + ): + self.access_token = access_token + self.subdomain = subdomain + # Authorization headers. Content-Type is not necessary, + # but should workable start providing alternate + # content types such as XML this won't break + self.headers = { + "accept": "application/json", + "Authorization": "Bearer " + access_token, + } + # Base URL Endpoint for all API requests + self.base_url = f"https://{subdomain}.workable.com/spi/v3" + # Increase the default limit for downloading lists from 50 to 100, + # so we need to make fewer requests + self.default_limit = 100 + self.params = {"limit": self.default_limit} + # Initialize an ISO-formatted start date. + # If the `start_date` argument is not provided, it defaults to January 1, 2000. + self.default_start_date = pendulum.datetime(2000, 1, 1).isoformat() + self.start_date_iso = ( + start_date.isoformat() + if start_date is not None + else self.default_start_date + ) + + def _request_with_rate_limit(self, url: str, **kwargs: Any) -> requests.Response: + """ + Handling rate limits in HTTP requests and ensuring + that the client doesn't exceed the limit set by the server + """ + while True: + try: + response = requests.get(url, **kwargs) + return response + except requests.HTTPError as e: + logging.warning("Rate limited. Waiting to retry...") + seconds_to_wait = ( + int(e.response.headers["X-Rate-Limit-Reset"]) + - pendulum.now().timestamp() + ) + time.sleep(seconds_to_wait) + + def pagination( + self, + endpoint: str, + custom_url: Optional[str] = None, + params: Optional[Dict[str, Any]] = None, + ) -> Iterable[TDataItem]: + """ + Queries an API endpoint using pagination and returns the results as a generator. + Args: + endpoint (str): The API endpoint to query. + custom_url (str, optional): A custom URL to use instead of the default base URL. + params (dict, optional): A dictionary of query parameters to include in the API request. + + """ + base_url = self.base_url if custom_url is None else custom_url + url = f"{base_url}/{endpoint}" + + if params is not None: + self.params.update(params) + + has_more = True + while has_more: + response = self._request_with_rate_limit( + url, headers=self.headers, params=self.params + ) + response_json = response.json() + paging = response_json.get("paging") + if paging is not None: + url = paging["next"] + else: + has_more = False + + yield response_json.get(endpoint) + + def details_from_endpoint( + self, main_endpoint: str, code: str, dependent_endpoint: str + ) -> Iterable[TDataItem]: + custom_url = f"{self.base_url}/{main_endpoint}/{code}" + return self.pagination(dependent_endpoint, custom_url=custom_url) diff --git a/omniload/source/workable/settings.py b/omniload/source/workable/settings.py new file mode 100644 index 000000000..9dee97570 --- /dev/null +++ b/omniload/source/workable/settings.py @@ -0,0 +1,44 @@ +# Copyright 2022-2026 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. + +"""Workable source settings and constants""" + +# define which endpoints to load +DEFAULT_ENDPOINTS = ( + "members", + "recruiters", + "stages", + "requisitions", + "jobs", + "custom_attributes", + "events", +) + +# define which sub endpoints to load for each main endpoint if details +# are requested +DEFAULT_DETAILS = { + "candidates": ( + "activities", + "offer", + ), + "jobs": ( + "activities", + "application_form", + "questions", + "stages", + "custom_attributes", + "members", + "recruiters", + ), +} diff --git a/omniload/source/zendesk/adapter.py b/omniload/source/zendesk/adapter.py index 088fa09a0..cd7dfc479 100644 --- a/omniload/source/zendesk/adapter.py +++ b/omniload/source/zendesk/adapter.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,22 +17,17 @@ """ from itertools import chain -from typing import Iterable, Iterator, Optional, Union +from typing import Iterable, Iterator, Optional import dlt from dlt.common import pendulum -from dlt.common.time import ensure_pendulum_datetime_utc +from dlt.common.time import ensure_pendulum_datetime from dlt.common.typing import TAnyDateTime, TDataItem, TDataItems from dlt.sources import DltResource -from omniload.error import MissingValueError - +from .helpers import make_date_ranges from .helpers.api_helpers import process_ticket, process_ticket_field -from .helpers.credentials import ( - TZendeskCredentials, - ZendeskCredentialsOAuth, - ZendeskCredentialsToken, -) +from .helpers.credentials import TZendeskCredentials, ZendeskCredentialsOAuth from .helpers.talk_api import PaginationType, ZendeskAPIClient from .settings import ( CUSTOM_FIELDS_STATE_KEY, @@ -44,10 +39,10 @@ ) -@dlt.source(max_table_nesting=0) +@dlt.source(max_table_nesting=2) def zendesk_talk( credentials: TZendeskCredentials = dlt.secrets.value, - start_date: TAnyDateTime = DEFAULT_START_DATE, + start_date: Optional[TAnyDateTime] = DEFAULT_START_DATE, end_date: Optional[TAnyDateTime] = None, ) -> Iterable[DltResource]: """ @@ -69,8 +64,8 @@ def zendesk_talk( # use the credentials to authenticate with the ZendeskClient zendesk_client = ZendeskAPIClient(credentials) - start_date_obj = ensure_pendulum_datetime_utc(start_date) - end_date_obj = ensure_pendulum_datetime_utc(end_date) if end_date else None + start_date_obj = ensure_pendulum_datetime(start_date) + end_date_obj = ensure_pendulum_datetime(end_date) if end_date else None # regular endpoints for key, talk_endpoint, item_name, cursor_paginated in TALK_ENDPOINTS: @@ -147,18 +142,13 @@ def talk_incremental_resource( Yields: TDataItem: Dictionary containing the data from the endpoint. """ - if updated_at is None or updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Zendesk") - # send the request and process it for page in zendesk_client.get_pages( talk_endpoint, talk_endpoint_name, PaginationType.START_TIME, params={ - "start_time": ensure_pendulum_datetime_utc( - updated_at.last_value - ).int_timestamp + "start_time": ensure_pendulum_datetime(updated_at.last_value).int_timestamp }, ): yield page @@ -166,10 +156,10 @@ def talk_incremental_resource( return -@dlt.source(max_table_nesting=0) +@dlt.source(max_table_nesting=2) def zendesk_chat( - credentials: Union[ZendeskCredentialsOAuth, ZendeskCredentialsToken], - start_date: TAnyDateTime = DEFAULT_START_DATE, + credentials: ZendeskCredentialsOAuth = dlt.secrets.value, + start_date: Optional[TAnyDateTime] = DEFAULT_START_DATE, end_date: Optional[TAnyDateTime] = None, ) -> Iterable[DltResource]: """ @@ -192,8 +182,8 @@ def zendesk_chat( # Authenticate zendesk_client = ZendeskAPIClient(credentials, url_prefix="https://www.zopim.com") - start_date_obj = ensure_pendulum_datetime_utc(start_date) - end_date_obj = ensure_pendulum_datetime_utc(end_date) if end_date else None + start_date_obj = ensure_pendulum_datetime(start_date) + end_date_obj = ensure_pendulum_datetime(end_date) if end_date else None yield dlt.resource(chats_table_resource, name="chats", write_disposition="merge")( zendesk_client, @@ -220,15 +210,12 @@ def chats_table_resource( Yields: dict: A dictionary representing each row of data. """ - if update_timestamp is None or update_timestamp.last_value is None: - raise MissingValueError("update_timestamp.last_value", "Zendesk") - chat_pages = zendesk_client.get_pages( "/api/v2/incremental/chats", "chats", PaginationType.START_TIME, params={ - "start_time": ensure_pendulum_datetime_utc( + "start_time": ensure_pendulum_datetime( update_timestamp.last_value ).int_timestamp, "fields": "chats(*)", @@ -241,12 +228,12 @@ def chats_table_resource( return -@dlt.source(max_table_nesting=0) +@dlt.source(max_table_nesting=2) def zendesk_support( - credentials: TZendeskCredentials, + credentials: TZendeskCredentials = dlt.secrets.value, load_all: bool = True, pivot_ticket_fields: bool = True, - start_date: TAnyDateTime = DEFAULT_START_DATE, + start_date: Optional[TAnyDateTime] = DEFAULT_START_DATE, end_date: Optional[TAnyDateTime] = None, ) -> Iterable[DltResource]: """ @@ -269,11 +256,8 @@ def zendesk_support( Sequence[DltResource]: Multiple dlt resources. """ - if start_date is None: - raise MissingValueError("start_date", "Zendesk") - - start_date_obj = ensure_pendulum_datetime_utc(start_date) - end_date_obj = ensure_pendulum_datetime_utc(end_date) if end_date else None + start_date_obj = ensure_pendulum_datetime(start_date) + end_date_obj = ensure_pendulum_datetime(end_date) if end_date else None start_date_ts = start_date_obj.int_timestamp start_date_iso_str = start_date_obj.isoformat() @@ -291,8 +275,6 @@ def ticket_events( initial_value=start_date_ts, end_value=end_date_ts, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), ) -> Iterator[TDataItem]: # URL For ticket events @@ -327,8 +309,6 @@ def ticket_table( initial_value=start_date_obj, end_value=end_date_obj, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), ) -> Iterator[TDataItem]: """ @@ -345,8 +325,6 @@ def ticket_table( Yields: TDataItem: Dictionary containing the ticket data. """ - if updated_at.last_value is None: - raise MissingValueError("updated_at.last_value", "Zendesk") # grab the custom fields from dlt state if any if pivot_fields: load_ticket_fields_state(zendesk_client) @@ -377,8 +355,6 @@ def ticket_metric_table( initial_value=start_date_iso_str, end_value=end_date_iso_str, allow_external_schedulers=True, - range_end="closed", - range_start="closed", ), ) -> Iterator[TDataItem]: """ @@ -395,16 +371,12 @@ def ticket_metric_table( TDataItem: Dictionary containing the ticket metric event data. """ # "https://example.zendesk.com/api/v2/incremental/ticket_metric_events?start_time=1332034771" - if time.last_value is None: - raise MissingValueError("time.last_value", "Zendesk") metric_event_pages = zendesk_client.get_pages( "/api/v2/incremental/ticket_metric_events", "ticket_metric_events", PaginationType.CURSOR, params={ - "start_time": ensure_pendulum_datetime_utc( - time.last_value - ).int_timestamp, + "start_time": ensure_pendulum_datetime(time.last_value).int_timestamp, }, ) for page in metric_event_pages: diff --git a/omniload/source/zendesk/settings.py b/omniload/source/zendesk/settings.py index 2e8c18857..9e9475cfd 100644 --- a/omniload/source/zendesk/settings.py +++ b/omniload/source/zendesk/settings.py @@ -1,4 +1,4 @@ -# Copyright 2022-2025 ScaleVector +# Copyright 2022-2026 ScaleVector # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,10 +16,9 @@ from dlt.common import pendulum -DEFAULT_START_DATE = pendulum.datetime(year=2024, month=10, day=3) - -INCREMENTAL_PAGE_SIZE = 1000 +DEFAULT_START_DATE = pendulum.datetime(year=2000, month=1, day=1) PAGE_SIZE = 100 +INCREMENTAL_PAGE_SIZE = 1000 CUSTOM_FIELDS_STATE_KEY = "ticket_custom_fields_v2" @@ -39,20 +38,34 @@ SUPPORT_EXTRA_ENDPOINTS = [ ("activities", "/api/v2/activities.json", None, True), ("automations", "/api/v2/automations.json", None, True), + ("custom_agent_roles", "/api/v2/custom_roles.json", "custom_roles", False), + ("dynamic_content", "/api/v2/dynamic_content/items.json", "items", True), + ("group_memberships", "/api/v2/group_memberships.json", None, True), + ("job_status", "/api/v2/job_statuses.json", "job_statuses", True), ("macros", "/api/v2/macros.json", None, True), + ("organization_fields", "/api/v2/organization_fields.json", None, True), + ("organization_memberships", "/api/v2/organization_memberships.json", None, True), ("recipient_addresses", "/api/v2/recipient_addresses.json", None, True), ("requests", "/api/v2/requests.json", None, True), + ("satisfaction_ratings", "/api/v2/satisfaction_ratings.json", None, True), + ("sharing_agreements", "/api/v2/sharing_agreements.json", None, False), + ("skips", "/api/v2/skips.json", None, True), + ("suspended_tickets", "/api/v2/suspended_tickets.json", None, True), ("targets", "/api/v2/targets.json", None, False), ("ticket_forms", "/api/v2/ticket_forms.json", None, False), ("ticket_metrics", "/api/v2/ticket_metrics.json", None, True), ("triggers", "/api/v2/triggers.json", None, True), ("user_fields", "/api/v2/user_fields.json", None, True), + ("views", "/api/v2/views.json", None, True), + ("tags", "/api/v2/tags.json", None, True), ] TALK_ENDPOINTS = [ ("calls", "/api/v2/channels/voice/calls", None, False), ("addresses", "/api/v2/channels/voice/addresses", None, False), + ("greeting_categories", "/api/v2/channels/voice/greeting_categories", None, False), ("greetings", "/api/v2/channels/voice/greetings", None, False), + ("ivrs", "/api/v2/channels/voice/ivr", None, False), ("phone_numbers", "/api/v2/channels/voice/phone_numbers", None, False), ("settings", "/api/v2/channels/voice/settings", None, False), ("lines", "/api/v2/channels/voice/lines", None, False), diff --git a/omniload/util/log.py b/omniload/util/log.py index c4f5c28d4..9b98e7fd0 100644 --- a/omniload/util/log.py +++ b/omniload/util/log.py @@ -11,7 +11,10 @@ def setup_logging( level=logging.INFO, verbose: bool = False, debug: bool = False, width: int = 36 ): - if os.environ.get("DEBUG"): + + debug = debug or os.environ.get("DEBUG", False) + + if debug: level = logging.DEBUG reset = escape_codes["reset"] diff --git a/omniload/util/vs.py b/omniload/util/vs.py new file mode 100644 index 000000000..1c1424dbf --- /dev/null +++ b/omniload/util/vs.py @@ -0,0 +1,285 @@ +import logging +import shutil +from contextlib import chdir +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import platformdirs +from git import Repo + +from omniload.util.log import setup_logging + +logger = logging.getLogger(__name__) + + +FILE_HEADER = b""" +# Copyright 2022-2026 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. + +""".lstrip() + + +@dataclass +class Rule: + source: str + target: Optional[str] + category: Optional[str] = None + use_tests: Optional[bool] = True + + +class VerifiedSourcesRecipe(list): + url: str = "https://github.com/dlt-hub/verified-sources.git" + code_path: str = "sources" + test_path: str = "tests" + + def __init__(self): + super().__init__() + self.append(Rule(source="airtable", target="airtable", category="saas")) + self.append(Rule(source="asana_dlt", target="asana", category="saas")) + self.append( + Rule(source="bing_webmaster", target="bing_webmaster", category="saas") + ) + self.append(Rule(source="chess", target="chess", category="saas")) + self.append(Rule(source="facebook_ads", target="facebook_ads", category="saas")) + self.append(Rule(source="filesystem", target=None)) + self.append(Rule(source="freshdesk", target="freshdesk", category="saas")) + self.append(Rule(source="github", target="github", category="saas")) + self.append(Rule(source="google_ads", target="google_ads", category="saas")) + self.append( + Rule(source="google_analytics", target="google_analytics", category="saas") + ) + self.append( + Rule(source="google_sheets", target="google_sheets", category="saas") + ) + self.append(Rule(source="hubspot", target="hubspot", category="saas")) + self.append(Rule(source="inbox", target="imap", category="protocol")) + self.append(Rule(source="jira", target="jira", category="saas")) + self.append(Rule(source="kafka", target="kafka", category="stream")) + self.append(Rule(source="kinesis", target="kinesis", category="stream")) + self.append(Rule(source="matomo", target="matomo", category="saas")) + self.append(Rule(source="mongodb", target="mongodb", category="database")) + self.append(Rule(source="mux", target="mux", category="saas")) + self.append(Rule(source="notion", target="notion", category="saas")) + self.append(Rule(source="personio", target="personio", category="saas")) + self.append( + Rule(source="pg_replication", target="pg_replication", category="database") + ) + self.append(Rule(source="pipedrive", target="pipedrive", category="saas")) + self.append(Rule(source="pokemon", target="pokemon", category="saas")) + self.append(Rule(source="rest_api", target=None)) + self.append(Rule(source="salesforce", target="salesforce", category="saas")) + self.append(Rule(source="scraping", target="scrapy", category="tool")) + self.append(Rule(source="shopify_dlt", target="shopify", category="saas")) + self.append(Rule(source="slack", target="slack", category="saas")) + self.append(Rule(source="sql_database", target=None)) + self.append(Rule(source="strapi", target="strapi", category="saas")) + self.append(Rule(source="stripe_analytics", target="stripe", category="saas")) + self.append( + Rule(source="unstructured_data", target="unstructured", category="tool") + ) + self.append( + Rule( + source="unstructured_data/google_drive", + target="google_drive", + category="storage", + use_tests=False, + ) + ) + self.append(Rule(source="workable", target="workable", category="saas")) + self.append(Rule(source="zendesk", target="zendesk", category="saas")) + + +class VerifiedSourcesSync: + def __init__(self, target_path): + self.target_path = target_path + self.workdir = platformdirs.user_cache_path("omniload") / "dlt-verified-sources" + self.workdir.mkdir(parents=True, exist_ok=True) + logger.info("Working directory: {}".format(self.workdir)) + self.repo = Repo(self.workdir) + self.tree = self.repo.head.commit.tree + self.recipe = VerifiedSourcesRecipe() + + def run(self): + self.acquire_sources() + self.process() + self.report_unmapped_modules() + + def process(self): + """Process and apply recipe rules, copying and rewriting the whole tree""" + self.copy_code() + # self.copy_tests() + + def copy_code(self): + """Copy code from `verified-sources` into omniload""" + + omniload_source_path = self.target_path / "omniload" / "source" + logger.info(f"Copy code to {omniload_source_path}") + + for rule in self.effective_rules: + module_path = self.tree / self.recipe.code_path / rule.source + + # Process Python files only. + # TODO: How to use the README.md files? + def tree_filter(item, _): + if item.name.endswith(".py"): + return True + return False + + selected_files = list(module_path.traverse(depth=1, predicate=tree_filter)) + selected_file_names = [item.name for item in selected_files] + logger.debug( + 'Upstream files in module "%s": %s', + module_path.name, + selected_file_names, + ) + + for blob in selected_files: + payload = blob.data_stream.read() + file_name = blob.name + source_path = blob.path + + # Rule: Rename `__init__.py` to `adapter.py`. + target_file_name = file_name.replace("__init__.py", "adapter.py") + + # Rule: Rename `*_client.py` to `client.py`. + if file_name.endswith("_client.py"): + target_file_name = "client.py" + + # Rule: Skip `setup_script_gcp_oauth.py`. + if file_name.startswith("setup_script_"): + continue + + target_path = omniload_source_path / rule.target / target_file_name + # logger.debug("%s -> %s", source_path, target_path) + Path(target_path).parent.mkdir(parents=True, exist_ok=True) + Path(target_path).write_bytes(FILE_HEADER + payload) + + def copy_tests(self): + """Copy tests from `verified-sources` into omniload""" + + omniload_tests_path = self.target_path / "tests" / "dlt" + logger.info(f"Copy tests to {omniload_tests_path}") + + for rule in self.effective_rules: + if not rule.use_tests: + continue + + module_path = self.tree / self.recipe.test_path / rule.source + + def tree_filter(item, _): + if item.name == "__init__.py": + return False + if item.name.endswith(".py"): + return True + return False + + selected_files = list(module_path.traverse(depth=1, predicate=tree_filter)) + selected_file_names = [item.name for item in selected_files] + logger.debug( + 'Upstream files in module "%s": %s', + module_path.name, + selected_file_names, + ) + + use_directory = False + if len(selected_files) > 1: + use_directory = True + + for blob in selected_files: + payload = blob.data_stream.read() + file_name = blob.name + source_path = blob.path + + if use_directory: + target_path = omniload_tests_path / rule.target / file_name + else: + # Rule: For single-file tests, rename `test_*_{source}.py` to `test_{source}.py`. + if file_name.startswith("test_"): + file_name = f"test_{rule.target}.py" + target_path = omniload_tests_path / file_name + + # logger.debug("%s -> %s", source_path, target_path) + Path(target_path).parent.mkdir(parents=True, exist_ok=True) + Path(target_path).write_bytes(FILE_HEADER + payload) + + @property + def skipped_modules(self): + return [ + "airtable", + "asana", + "chess", + "kafka", + "stripe", + ] + + @property + def effective_rules(self): + """Effective recipe rules to apply""" + rules = [] + for rule in self.recipe: + if rule.target is None: + logger.info(f"Skipping module {rule.source}") + continue + rules.append(rule) + return rules + + def acquire_sources(self): + """Acquire source code repository""" + if len(list(self.workdir.iterdir())) == 0: + with chdir(self.workdir): + Repo.clone_from(self.recipe.url, self.workdir) + + def report_unmapped_modules(self): + """Inform the user about new modules""" + modules = self.unmapped_module_names + if modules: + logger.warning("Those modules are not mapped yet: %s", modules) + else: + logger.info("All existing modules are properly mapped") + + @property + def unmapped_module_names(self): + """Module names not mapped by rules""" + return sorted(set(self.upstream_module_names) - set(self.mapped_module_names)) + + @property + def mapped_module_names(self): + """Module names covered by recipe rules""" + return [rule.source for rule in self.recipe] + + @property + def upstream_module_names(self): + """Module names in the verified-sources repository""" + names = [] + for item in self.tree / self.recipe.code_path: + # Only process directories, skip files and other items. + if item.type != "tree": + continue + # Skip special directories. + if item.name in [".dlt"]: + continue + names.append(item.name) + return names + + +if __name__ == "__main__": + """ + uv pip install gitpython platformdirs + python -m omniload.util.vs + """ + setup_logging(debug=True) + target_path = Path.cwd() + engine = VerifiedSourcesSync(target_path=target_path) + engine.run()