omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/zendesk/helpers/credentials.py

@@ -0,0 +1,68 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module handles how credentials are read in dlt sources
+"""
+
+from typing import ClassVar, List, Union
+
+import dlt
+from dlt.common.configuration import configspec
+from dlt.common.configuration.specs import CredentialsConfiguration
+from dlt.common.typing import TSecretValue
+
+
+@configspec
+class ZendeskCredentialsBase(CredentialsConfiguration):
+    """
+    The Base version of all the ZendeskCredential classes.
+    """
+
+    subdomain: str = dlt.config.value
+    __config_gen_annotations__: ClassVar[List[str]] = []
+
+
+@configspec
+class ZendeskCredentialsEmailPass(ZendeskCredentialsBase):
+    """
+    This class is used to store credentials for Email + Password Authentication
+    """
+
+    email: str = dlt.config.value
+    password: TSecretValue = dlt.secrets.value
+
+
+@configspec
+class ZendeskCredentialsOAuth(ZendeskCredentialsBase):
+    """
+    This class is used to store credentials for OAuth Token Authentication
+    """
+
+    oauth_token: TSecretValue = dlt.secrets.value
+
+
+@configspec
+class ZendeskCredentialsToken(ZendeskCredentialsBase):
+    """
+    This class is used to store credentials for Token Authentication
+    """
+
+    email: str = dlt.config.value
+    token: TSecretValue = dlt.secrets.value
+
+
+TZendeskCredentials = Union[
+    ZendeskCredentialsEmailPass, ZendeskCredentialsToken, ZendeskCredentialsOAuth
+]

omniload/src/zendesk/helpers/talk_api.py

@@ -0,0 +1,132 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+from typing import Any, Dict, Iterator, Optional, Tuple
+
+from dlt.common.typing import DictStrStr, TDataItems, TSecretValue
+from dlt.sources.helpers.requests import client
+
+from .. import settings
+from .credentials import (
+    TZendeskCredentials,
+    ZendeskCredentialsEmailPass,
+    ZendeskCredentialsOAuth,
+    ZendeskCredentialsToken,
+)
+
+
+class PaginationType(Enum):
+    OFFSET = 0
+    CURSOR = 1
+    STREAM = 2
+    START_TIME = 3
+
+
+class ZendeskAPIClient:
+    """
+    API client used to make requests to Zendesk talk, support and chat API
+    """
+
+    subdomain: str = ""
+    url: str = ""
+    headers: Optional[DictStrStr]
+    auth: Optional[Tuple[str, TSecretValue]]
+
+    def __init__(
+        self, credentials: TZendeskCredentials, url_prefix: Optional[str] = None
+    ) -> None:
+        """
+        Initializer for the API client which is then used to make API calls to the ZendeskAPI
+
+        Args:
+            credentials: ZendeskCredentials object which contains the necessary credentials to authenticate to ZendeskAPI
+        """
+        # oauth token is the preferred way to authenticate, followed by api token and then email + password combo
+        # fill headers and auth for every possibility of credentials given, raise error if credentials are of incorrect type
+        if isinstance(credentials, ZendeskCredentialsOAuth):
+            self.headers = {"Authorization": f"Bearer {credentials.oauth_token}"}
+            self.auth = None
+        elif isinstance(credentials, ZendeskCredentialsToken):
+            self.headers = None
+            self.auth = (f"{credentials.email}/token", credentials.token)
+        elif isinstance(credentials, ZendeskCredentialsEmailPass):
+            self.auth = (credentials.email, credentials.password)
+            self.headers = None
+        else:
+            raise TypeError(
+                "Wrong credentials type provided to ZendeskAPIClient. The credentials need to be of type: ZendeskCredentialsOAuth, ZendeskCredentialsToken or ZendeskCredentialsEmailPass"
+            )
+
+        # If url_prefix is set it overrides the default API URL (e.g. chat api uses zopim.com domain)
+        if url_prefix:
+            self.url = url_prefix
+        else:
+            self.subdomain = credentials.subdomain
+            self.url = f"https://{self.subdomain}.zendesk.com"
+
+    def get_pages(
+        self,
+        endpoint: str,
+        data_point_name: str,
+        pagination: PaginationType,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> Iterator[TDataItems]:
+        """
+        Makes a request to a paginated endpoint and returns a generator of data items per page.
+
+        Args:
+            endpoint: The url to the endpoint, e.g. /api/v2/calls
+            data_point_name: The key which data items are nested under in the response object (e.g. calls)
+            params: Optional dict of query params to include in the request
+            pagination: Type of pagination type used by endpoint
+
+        Returns:
+            Generator of pages, each page is a list of dict data items
+        """
+        # update the page size to enable cursor pagination
+        params = params or {}
+        if pagination == PaginationType.CURSOR:
+            params["page[size]"] = settings.PAGE_SIZE
+        elif pagination == PaginationType.STREAM:
+            params["per_page"] = settings.INCREMENTAL_PAGE_SIZE
+        elif pagination == PaginationType.START_TIME:
+            params["limit"] = settings.INCREMENTAL_PAGE_SIZE
+
+        # make request and keep looping until there is no next page
+        get_url = f"{self.url}{endpoint}"
+        while get_url:
+            response = client.get(
+                get_url, headers=self.headers, auth=self.auth, params=params
+            )
+            response.raise_for_status()
+            response_json = response.json()
+            result = response_json[data_point_name]
+            yield result
+
+            get_url = None
+            if pagination == PaginationType.CURSOR:
+                if response_json["meta"]["has_more"]:
+                    get_url = response_json["links"]["next"]
+            elif pagination == PaginationType.OFFSET:
+                get_url = response_json.get("next_page", None)
+            elif pagination == PaginationType.STREAM:
+                # See https://developer.zendesk.com/api-reference/ticketing/ticket-management/incremental_exports/#json-format
+                if not response_json["end_of_stream"]:
+                    get_url = response_json["next_page"]
+            elif pagination == PaginationType.START_TIME:
+                if response_json["count"] > 0:
+                    get_url = response_json["next_page"]
+
+            params = {}

omniload/src/zendesk/settings.py

@@ -0,0 +1,71 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Zendesk source settings and constants"""
+
+from dlt.common import pendulum
+
+DEFAULT_START_DATE = pendulum.datetime(year=2024, month=10, day=3)
+
+INCREMENTAL_PAGE_SIZE = 1000
+PAGE_SIZE = 100
+
+
+CUSTOM_FIELDS_STATE_KEY = "ticket_custom_fields_v2"
+
+# Tuples of (Resource name, endpoint URL, data_key, supports pagination)
+# data_key is the key which data list is nested under in responses
+# if the data key is None it is assumed to be the same as the resource name
+# The last element of the tuple says if endpoint supports cursor pagination
+SUPPORT_ENDPOINTS = [
+    ("users", "/api/v2/users.json", "users", True),
+    ("sla_policies", "/api/v2/slas/policies.json", None, False),
+    ("groups", "/api/v2/groups.json", None, True),
+    ("organizations", "/api/v2/organizations.json", None, True),
+    ("brands", "/api/v2/brands.json", None, True),
+]
+
+SUPPORT_EXTRA_ENDPOINTS = [
+    ("activities", "/api/v2/activities.json", None, True),
+    ("automations", "/api/v2/automations.json", None, True),
+    ("macros", "/api/v2/macros.json", None, True),
+    ("recipient_addresses", "/api/v2/recipient_addresses.json", None, True),
+    ("requests", "/api/v2/requests.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),
+]
+
+TALK_ENDPOINTS = [
+    ("calls", "/api/v2/channels/voice/calls", None, False),
+    ("addresses", "/api/v2/channels/voice/addresses", None, False),
+    ("greetings", "/api/v2/channels/voice/greetings", 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),
+    ("agents_activity", "/api/v2/channels/voice/stats/agents_activity", None, False),
+    (
+        "current_queue_activity",
+        "/api/v2/channels/voice/stats/current_queue_activity",
+        None,
+        False,
+    ),
+]
+
+INCREMENTAL_TALK_ENDPOINTS = {
+    "calls": "/api/v2/channels/voice/stats/incremental/calls.json",
+    "legs": "/api/v2/channels/voice/stats/incremental/legs.json",
+}

omniload/src/zoom/__init__.py

@@ -0,0 +1,99 @@
+from typing import Any, Dict, Iterable, Sequence
+
+import dlt
+import pendulum
+from dlt.common.typing import TAnyDateTime, TDataItem
+from dlt.sources import DltResource
+
+from .helpers import ZoomClient
+
+
+@dlt.source(name="zoom", max_table_nesting=0)
+def zoom_source(
+    client_id: str,
+    client_secret: str,
+    account_id: str,
+    start_date: pendulum.DateTime,
+    end_date: pendulum.DateTime | None = None,
+) -> Sequence[DltResource]:
+    """Create a Zoom source with meetings resource for all users in the account."""
+    client = ZoomClient(
+        client_id=client_id,
+        client_secret=client_secret,
+        account_id=account_id,
+    )
+
+    @dlt.resource(write_disposition="merge", primary_key="id")
+    def meetings(
+        datetime: dlt.sources.incremental[TAnyDateTime] = dlt.sources.incremental(
+            "start_time",
+            initial_value=start_date.isoformat(),
+            end_value=end_date.isoformat() if end_date is not None else None,
+            range_start="closed",
+            range_end="closed",
+        ),
+    ) -> Iterable[TDataItem]:
+        if datetime.last_value:
+            start_dt = pendulum.parse(datetime.last_value)
+        else:
+            start_dt = pendulum.parse(start_date)
+
+        if end_date is None:
+            end_dt = pendulum.now("UTC")
+        else:
+            end_dt = pendulum.parse(datetime.end_value)
+
+        base_params: Dict[str, Any] = {
+            "type": "scheduled",
+            "page_size": 300,
+            "from": start_dt.to_date_string(),
+            "to": end_dt.to_date_string(),
+        }
+
+        for user in client.get_users():
+            user_id = user["id"]
+            yield from client.get_meetings(user_id, base_params)
+
+    @dlt.resource(write_disposition="merge", primary_key="id")
+    def users() -> Iterable[TDataItem]:
+        yield from client.get_users()
+
+    @dlt.resource(write_disposition="merge", primary_key="id")
+    def participants(
+        datetime: dlt.sources.incremental[TAnyDateTime] = dlt.sources.incremental(
+            "join_time",
+            initial_value=start_date.isoformat(),
+            end_value=end_date.isoformat() if end_date is not None else None,
+            range_start="closed",
+            range_end="closed",
+        ),
+    ) -> Iterable[TDataItem]:
+        if datetime.last_value:
+            start_dt = pendulum.parse(datetime.last_value)
+        else:
+            start_dt = pendulum.parse(start_date)
+
+        if end_date is None:
+            end_dt = pendulum.now("UTC")
+        else:
+            end_dt = pendulum.parse(datetime.end_value)
+
+        participant_params: Dict[str, Any] = {
+            "page_size": 300,
+        }
+        meeting_params = {
+            "type": "previous_meetings",
+            "page_size": 300,
+        }
+        for user in client.get_users():
+            user_id = user["id"]
+            for meeting in client.get_meetings(user_id=user_id, params=meeting_params):
+                meeting_id = meeting["id"]
+                yield from client.get_participants(
+                    meeting_id=meeting_id,
+                    params=participant_params,
+                    start_date=start_dt,
+                    end_date=end_dt,
+                )
+
+    return meetings, users, participants

omniload/src/zoom/helpers.py

@@ -0,0 +1,102 @@
+import time
+from typing import Any, Dict, Iterator, Optional
+
+import pendulum
+
+from omniload.src.http_client import create_client
+
+
+class ZoomClient:
+    """Minimal Zoom API client supporting Server-to-Server OAuth."""
+
+    def __init__(
+        self,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        account_id: Optional[str] = None,
+    ) -> None:
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.account_id = account_id
+        self.token_expires_at: float = 0
+        self.base_url = "https://api.zoom.us/v2"
+        self.session = create_client()
+        self._refresh_access_token()
+
+    def _refresh_access_token(self) -> None:
+        token_url = "https://zoom.us/oauth/token"
+        auth = (self.client_id, self.client_secret)
+        resp = self.session.post(
+            token_url,
+            params={"grant_type": "account_credentials", "account_id": self.account_id},
+            auth=auth,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        self.access_token = data.get("access_token")
+        self.token_expires_at = time.time() + data.get("expires_in", 3600)
+
+    def _ensure_token(self) -> None:
+        if self.access_token is None or self.token_expires_at <= time.time():
+            self._refresh_access_token()
+
+    def _headers(self) -> Dict[str, str]:
+        self._ensure_token()
+        return {
+            "Authorization": f"Bearer {self.access_token}",
+            "Accept": "application/json",
+        }
+
+    def get_users(self) -> Iterator[Dict[str, Any]]:
+        url = f"{self.base_url}/users"
+
+        params = {"page_size": 1000}
+        while True:
+            response = self.session.get(url, headers=self._headers(), params=params)
+            response.raise_for_status()
+            data = response.json()
+            for user in data.get("users", []):
+                yield user
+            token = data.get("next_page_token")
+            if not token:
+                break
+            params["next_page_token"] = token
+
+    # https://developers.zoom.us/docs/api/rest/reference/zoom-api/methods/#operation/meetings
+    def get_meetings(
+        self, user_id: str, params: Dict[str, Any]
+    ) -> Iterator[Dict[str, Any]]:
+        url = f"{self.base_url}/users/{user_id}/meetings"
+        while True:
+            response = self.session.get(url, headers=self._headers(), params=params)
+            response.raise_for_status()
+            data = response.json()
+            for item in data.get("meetings", []):
+                item["zoom_user_id"] = user_id
+                yield item
+            token = data.get("next_page_token")
+            if not token:
+                break
+            params["next_page_token"] = token
+
+    # https://developers.zoom.us/docs/api/rest/reference/zoom-api/methods/#operation/reportMeetingParticipants
+    def get_participants(
+        self,
+        meeting_id: str,
+        params: Dict[str, Any],
+        start_date: pendulum.DateTime,
+        end_date: pendulum.DateTime,
+    ) -> Iterator[Dict[str, Any]]:
+        url = f"{self.base_url}/report/meetings/{meeting_id}/participants"
+        while True:
+            response = self.session.get(url, headers=self._headers(), params=params)
+            response.raise_for_status()
+            data = response.json()
+            for item in data.get("participants", []):
+                join_time = pendulum.parse(item["join_time"])
+                if join_time >= start_date and join_time <= end_date:
+                    yield item
+            token = data.get("next_page_token")
+            if not token:
+                break
+            params["next_page_token"] = token

omniload/testdata/.gitignore

@@ -0,0 +1,2 @@
+*.db
+*.duckdb

omniload/testdata/create_replace.csv

@@ -0,0 +1,21 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-19","True","AGILENT TECHNOLOGIES INC"
+"AA","2024-04-19","True","ALCOA CORP"
+"AAA","2024-04-19","True","ALTERNATIVE ACCESS FIRST PRI"
+"AAAU","2024-04-19","True","GOLDMAN SACHS PHYSICAL GOLD"
+"AACG","2024-04-19","True","ATA CREATIVITY GLOBAL - ADR"
+"AACI","2024-04-19","True","ARMADA ACQUISITION CORP I"
+"AACIU","2024-04-19","True","ARMADA ACQUISITION CORP I"
+"AACIW","2024-04-19","True",""
+"AACT","2024-04-19","True","ARES ACQUISITION CORP II"
+"AACT+","2024-04-19","True",""
+"AACT=","2024-04-19","True","ARES ACQUISITION CORP II"
+"AADI","2024-04-19","True","AADI BIOSCIENCE INC"
+"AADR","2024-04-19","True","ADVISORSHARES DORSEY WRIGHT"
+"AAGR","2024-04-19","True","AFRICAN AGRICULTURE HOLDINGS"
+"AAGRW","2024-04-19","True",""
+"AAL","2024-04-19","True","AMERICAN AIRLINES GROUP INC"
+"AAMC","2024-04-19","True","ALTISOURCE ASSET MANAGEMENT"
+"AAME","2024-04-19","True","ATLANTIC AMERICAN CORP"
+"AAN","2024-04-19","True","AARON'S CO INC/THE"
+"AAOI","2024-04-19","True","APPLIED OPTOELECTRONICS INC"

omniload/testdata/delete_insert_expected.csv

@@ -0,0 +1,6 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-20","True","AGILENT TECHNOLOGIES INC____updated"
+"AA","2024-04-19","True","ALCOA CORP____updated"
+"AAA","2024-04-21","True","ALTERNATIVE ACCESS FIRST PRI____updated"
+"AAAU","2024-04-22","True","GOLDMAN SACHS PHYSICAL GOLD____updated"
+"B","2024-04-18","True","SOME TECHNOLOGIES INC"

omniload/testdata/delete_insert_part1.csv

@@ -0,0 +1,5 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-19","True","AGILENT TECHNOLOGIES INC"
+"AA","2024-04-19","True","ALCOA CORP"
+"AAA","2024-04-19","True","ALTERNATIVE ACCESS FIRST PRI"
+"B","2024-04-18","True","SOME TECHNOLOGIES INC"

omniload/testdata/delete_insert_part2.csv

@@ -0,0 +1,6 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-20","True","AGILENT TECHNOLOGIES INC____updated"
+"AA","2024-04-19","True","ALCOA CORP____updated"
+"AAA","2024-04-21","True","ALTERNATIVE ACCESS FIRST PRI____updated"
+"AAAU","2024-04-22","True","GOLDMAN SACHS PHYSICAL GOLD____updated"
+"BBB","2024-04-18","True","SOME CORP____updated"

omniload/testdata/merge_expected.csv

@@ -0,0 +1,5 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-20","True","AGILENT TECHNOLOGIES INC____updated"
+"AA","2024-04-19","True","ALCOA CORP"
+"AAA","2024-04-21","True","ALTERNATIVE ACCESS FIRST PRI____updated"
+"AAAU","2024-04-22","True","GOLDMAN SACHS PHYSICAL GOLD____updated"

omniload/testdata/merge_part1.csv

@@ -0,0 +1,4 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-19","True","AGILENT TECHNOLOGIES INC"
+"AA","2024-04-19","True","ALCOA CORP"
+"AAA","2024-04-19","True","ALTERNATIVE ACCESS FIRST PRI"

omniload/testdata/merge_part2.csv

@@ -0,0 +1,5 @@
+"symbol","date","isEnabled","name"
+"A","2024-04-20","True","AGILENT TECHNOLOGIES INC____updated"
+"AA","2024-04-19","True","ALCOA CORP____updated"
+"AAA","2024-04-21","True","ALTERNATIVE ACCESS FIRST PRI____updated"
+"AAAU","2024-04-22","True","GOLDMAN SACHS PHYSICAL GOLD____updated"

omniload/tests/unit/test_smartsheets.py

@@ -0,0 +1,133 @@
+import sys
+import unittest
+from unittest.mock import patch
+
+import smartsheet  # type: ignore
+from smartsheet.models import Cell, Column, Row, Sheet  # type: ignore
+
+from omniload.src.smartsheets import _get_sheet_data, smartsheet_source
+
+
+def pp(x):
+    print(x, file=sys.stderr)
+
+
+class TestSmartsheetSource(unittest.TestCase):
+    @patch("omniload.src.smartsheets.smartsheet.Smartsheet")
+    def test_smartsheet_source_success(self, mock_smartsheet_client):
+        # Mock Smartsheet client and its methods
+        mock_client_instance = mock_smartsheet_client.return_value
+
+        # Mock sheet details response
+        mock_sheet_details = Sheet(
+            {
+                "id": 123,
+                "name": "Test Sheet 1",
+                "columns": [
+                    Column(
+                        {"id": 1, "title": "Col A", "type": "TEXT_NUMBER", "index": 0}
+                    ),
+                    Column(
+                        {"id": 2, "title": "Col B", "type": "TEXT_NUMBER", "index": 1}
+                    ),
+                ],
+                "rows": [
+                    Row(
+                        {
+                            "id": 101,
+                            "sheetId": 123,
+                            "cells": [
+                                Cell({"columnId": 1, "value": "r1c1"}),
+                                Cell({"columnId": 2, "value": "r1c2"}),
+                            ],
+                        }
+                    ),
+                    Row(
+                        {
+                            "id": 102,
+                            "sheetId": 123,
+                            "cells": [
+                                Cell({"columnId": 1, "value": "r2c1"}),
+                                Cell({"columnId": 2, "value": "r2c2"}),
+                            ],
+                        }
+                    ),
+                ],
+            }
+        )
+        mock_client_instance.Sheets.get_sheet.return_value = mock_sheet_details
+
+        resource = smartsheet_source(access_token="test_token", sheet_id="123")
+        data = list(resource)
+        self.assertEqual(len(data), 2)
+        self.assertEqual(data[0], {"_row_id": 101, "Col A": "r1c1", "Col B": "r1c2"})
+        self.assertEqual(data[1], {"_row_id": 102, "Col A": "r2c1", "Col B": "r2c2"})
+
+        mock_smartsheet_client.assert_called_once_with("test_token")
+        mock_client_instance.Sheets.get_sheet.assert_any_call(
+            123, include=["objectValue"]
+        )  # for resource name
+        mock_client_instance.Sheets.get_sheet.assert_any_call(
+            123
+        )  # for _get_sheet_data
+
+    @patch("omniload.src.smartsheets.smartsheet.Smartsheet")
+    def test_smartsheet_source_api_error(self, mock_smartsheet_client):
+        mock_client_instance = mock_smartsheet_client.return_value
+        mock_client_instance.Sheets.get_sheet.side_effect = (
+            smartsheet.exceptions.ApiError("API Error", 500)
+        )
+
+        with self.assertRaises(smartsheet.exceptions.ApiError):
+            source = smartsheet_source(access_token="test_token", sheet_id="123")
+            # Consume the generator to trigger the API call
+            list(source)
+
+    def test_get_sheet_data(self):
+        mock_sheet = Sheet(
+            {
+                "id": 456,
+                "name": "Data Sheet",
+                "columns": [
+                    Column(
+                        {"id": 10, "title": "ID", "type": "TEXT_NUMBER", "index": 0}
+                    ),
+                    Column(
+                        {"id": 20, "title": "Value", "type": "TEXT_NUMBER", "index": 1}
+                    ),
+                ],
+                "rows": [
+                    Row(
+                        {
+                            "id": 201,
+                            "sheetId": 456,
+                            "cells": [
+                                Cell({"columnId": 10, "value": 1}),
+                                Cell({"columnId": 20, "value": "Alpha"}),
+                            ],
+                        }
+                    ),
+                    Row(
+                        {
+                            "id": 202,
+                            "sheetId": 456,
+                            "cells": [
+                                Cell({"columnId": 10, "value": 2}),
+                                Cell({"columnId": 20, "value": "Beta"}),
+                            ],
+                        }
+                    ),
+                ],
+            }
+        )
+
+        data_generator = _get_sheet_data(mock_sheet)
+        data = list(data_generator)
+
+        self.assertEqual(len(data), 2)
+        self.assertEqual(data[0], {"_row_id": 201, "ID": 1, "Value": "Alpha"})
+        self.assertEqual(data[1], {"_row_id": 202, "ID": 2, "Value": "Beta"})
+
+
+if __name__ == "__main__":
+    unittest.main()