omniload 0.0.0.dev0__py3-none-any.whl

omniload/src/filesystem/readers.py

@@ -0,0 +1,187 @@
+# Copyright 2022-2025 ScaleVector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Any, Iterator, List, Optional
+
+from dlt.common import json
+from dlt.common.typing import copy_sig
+from dlt.sources import DltResource, DltSource, TDataItems
+from dlt.sources.filesystem import FileItemDict
+
+from .helpers import fetch_arrow, fetch_json
+
+
+def _read_csv(
+    items: Iterator[FileItemDict], chunksize: int = 10000, **pandas_kwargs: Any
+) -> Iterator[TDataItems]:
+    """Reads csv file with Pandas chunk by chunk.
+
+    Args:
+        chunksize (int): Number of records to read in one chunk
+        **pandas_kwargs: Additional keyword arguments passed to Pandas.read_csv
+    Returns:
+        TDataItem: The file content
+    """
+    import pandas as pd
+
+    # apply defaults to pandas kwargs
+    kwargs = {**{"header": "infer", "chunksize": chunksize}, **pandas_kwargs}
+
+    for file_obj in items:
+        # Here we use pandas chunksize to read the file in chunks and avoid loading the whole file
+        # in memory.
+        with file_obj.open() as file:
+            for df in pd.read_csv(file, **kwargs):
+                yield df.to_dict(orient="records")
+
+
+def _read_csv_headless(
+    items: Iterator[FileItemDict],
+    chunksize: int = 10000,
+    column_names: Optional[List[str]] = None,
+    **pandas_kwargs: Any,
+) -> Iterator[TDataItems]:
+    """Reads csv file without headers, using provided column names or generating them.
+
+    Args:
+        chunksize (int): Number of records to read in one chunk
+        column_names (list[str], optional): Column names for the CSV. If not provided,
+            columns will be named unknown_col_0, unknown_col_1, etc.
+        **pandas_kwargs: Additional keyword arguments passed to Pandas.read_csv
+    Returns:
+        TDataItem: The file content
+    """
+    import pandas as pd
+
+    for file_obj in items:
+        with file_obj.open() as file:
+            # Determine column names
+            if column_names:
+                names = column_names
+            else:
+                # Count columns from first row
+                first_row = pd.read_csv(file, header=None, nrows=1)
+                num_columns = len(first_row.columns)
+                names = [f"unknown_col_{i}" for i in range(num_columns)]
+                file.seek(0)  # Reset file pointer after reading first row
+
+            kwargs = {
+                **{"header": None, "names": names, "chunksize": chunksize},
+                **pandas_kwargs,
+            }
+
+            for df in pd.read_csv(file, **kwargs):
+                yield df.to_dict(orient="records")
+
+
+def _read_jsonl(
+    items: Iterator[FileItemDict], chunksize: int = 1000
+) -> Iterator[TDataItems]:
+    """Reads jsonl file content and extract the data.
+
+    Args:
+        chunksize (int, optional): The number of JSON lines to load and yield at once, defaults to 1000
+
+    Returns:
+        TDataItem: The file content
+    """
+    for file_obj in items:
+        with file_obj.open() as f:
+            lines_chunk = []
+            for line in f:
+                lines_chunk.append(json.loadb(line))
+                if len(lines_chunk) >= chunksize:
+                    yield lines_chunk
+                    lines_chunk = []
+        if lines_chunk:
+            yield lines_chunk
+
+
+def _read_parquet(
+    items: Iterator[FileItemDict],
+    chunksize: int = 10,
+) -> Iterator[TDataItems]:
+    """Reads parquet file content and extract the data.
+
+    Args:
+        chunksize (int, optional): The number of files to process at once, defaults to 10.
+
+    Returns:
+        TDataItem: The file content
+    """
+    from pyarrow import parquet as pq
+
+    for file_obj in items:
+        with file_obj.open() as f:
+            parquet_file = pq.ParquetFile(f)
+            for rows in parquet_file.iter_batches(batch_size=chunksize):
+                yield rows.to_pylist()
+
+
+def _read_csv_duckdb(
+    items: Iterator[FileItemDict],
+    chunk_size: Optional[int] = 5000,
+    use_pyarrow: bool = False,
+    **duckdb_kwargs: Any,
+) -> Iterator[TDataItems]:
+    """A resource to extract data from the given CSV files.
+
+    Uses DuckDB engine to import and cast CSV data.
+
+    Args:
+        items (Iterator[FileItemDict]): CSV files to read.
+        chunk_size (Optional[int]):
+            The number of rows to read at once. Defaults to 5000.
+        use_pyarrow (bool):
+            Whether to use `pyarrow` to read the data and designate
+            data schema. If set to False (by default), JSON is used.
+        duckdb_kwargs (Dict):
+            Additional keyword arguments to pass to the `read_csv()`.
+
+    Returns:
+        Iterable[TDataItem]: Data items, read from the given CSV files.
+    """
+    import duckdb
+
+    helper = fetch_arrow if use_pyarrow else fetch_json
+
+    for item in items:
+        with item.open() as f:
+            file_data = duckdb.from_csv_auto(f, **duckdb_kwargs)  # type: ignore
+
+            yield from helper(file_data, chunk_size)
+
+
+if TYPE_CHECKING:
+
+    class ReadersSource(DltSource):
+        """This is a typing stub that provides docstrings and signatures to the resources in `readers" source"""
+
+        @copy_sig(_read_csv)
+        def read_csv(self) -> DltResource: ...
+
+        @copy_sig(_read_csv_headless)
+        def read_csv_headless(self) -> DltResource: ...
+
+        @copy_sig(_read_jsonl)
+        def read_jsonl(self) -> DltResource: ...
+
+        @copy_sig(_read_parquet)
+        def read_parquet(self) -> DltResource: ...
+
+        @copy_sig(_read_csv_duckdb)
+        def read_csv_duckdb(self) -> DltResource: ...
+
+else:
+    ReadersSource = DltSource

omniload/src/filters.py

@@ -0,0 +1,62 @@
+def cast_set_to_list(row):
+    # this handles just the sqlalchemy backend for now
+    if isinstance(row, dict):
+        for key in row.keys():
+            if isinstance(row[key], set):
+                row[key] = list(row[key])
+    return row
+
+
+def cast_spanner_types(row):
+    if not isinstance(row, dict):
+        return row
+
+    from google.cloud.spanner_v1.data_types import JsonObject
+
+    for key in row.keys():
+        if isinstance(row[key], JsonObject):
+            import json
+
+            row[key] = json.loads(row[key].serialize())
+    return row
+
+
+def handle_mysql_empty_dates(row):
+    # MySQL returns empty dates as 0000-00-00, which is not a valid date, we handle them here.
+    if not isinstance(row, dict):
+        return row
+
+    for key in row.keys():
+        if not isinstance(row[key], str):
+            continue
+
+        if row[key] == "0000-00-00":
+            from datetime import date
+
+            row[key] = date(1970, 1, 1)
+
+        elif row[key] == "0000-00-00 00:00:00":
+            from datetime import datetime
+
+            row[key] = datetime(1970, 1, 1, 0, 0, 0)
+    return row
+
+
+def table_adapter_exclude_columns(cols: list[str]):
+    from dlt.common.libs.sql_alchemy import Table
+
+    def excluder(table: Table):
+        cols_to_remove = [col for col in table._columns if col.name in cols]  # type: ignore
+        for col in cols_to_remove:
+            table._columns.remove(col)  # type: ignore
+
+    return excluder
+
+
+def create_masking_filter(mask_configs: list[str]):
+    from omniload.src.masking import create_masking_mapper
+
+    if not mask_configs:
+        return lambda x: x
+
+    return create_masking_mapper(mask_configs)

omniload/src/fireflies/__init__.py

@@ -0,0 +1,151 @@
+from typing import Iterable, List, Optional
+
+import dlt
+import pendulum
+from dlt.common.time import ensure_pendulum_datetime
+from dlt.common.typing import TDataItem
+from dlt.sources import DltResource
+
+from .helpers import FirefliesAPI
+
+
+@dlt.source(name="fireflies", max_table_nesting=0)
+def fireflies_source(
+    api_key: str,
+    start_datetime: Optional[pendulum.DateTime],
+    end_datetime: Optional[pendulum.DateTime],
+    granularity: Optional[str] = None,
+) -> List[DltResource]:
+    fireflies_api = FirefliesAPI(api_key=api_key)
+
+    start_datetime = (
+        ensure_pendulum_datetime(start_datetime) if start_datetime else None
+    )
+    end_datetime = ensure_pendulum_datetime(end_datetime) if end_datetime else None
+
+    # Select fetch method based on granularity
+    def get_analytics_fetcher():
+        if granularity == "DAY":
+            return fireflies_api.fetch_analytics_daily
+        elif granularity == "HOUR":
+            return fireflies_api.fetch_analytics_hourly
+        elif granularity == "MONTH":
+            return fireflies_api.fetch_analytics_monthly
+        else:
+            return fireflies_api.fetch_analytics
+
+    @dlt.resource(write_disposition="replace")
+    def active_meetings() -> Iterable[TDataItem]:
+        for page in fireflies_api.fetch_active_meetings():
+            for item in page:
+                yield item
+
+    @dlt.resource(
+        write_disposition="merge",
+        primary_key=["start_time", "end_time"],
+    )
+    def analytics(
+        updated_at: dlt.sources.incremental[
+            pendulum.DateTime
+        ] = dlt.sources.incremental(
+            "end_time",
+            initial_value=start_datetime
+            if start_datetime
+            else pendulum.datetime(1970, 1, 1, tz="UTC"),
+            end_value=end_datetime if end_datetime else None,
+            range_end="closed" if end_datetime else "open",
+            range_start="closed",
+        ),
+    ) -> Iterable[TDataItem]:
+        from_date_dt = updated_at.last_value
+        to_date_dt = (
+            updated_at.end_value if updated_at.end_value else pendulum.now(tz="UTC")
+        )
+
+        from_date_iso = from_date_dt.to_iso8601_string() if from_date_dt else None
+        to_date_iso = to_date_dt.to_iso8601_string() if to_date_dt else None
+
+        fetch_method = get_analytics_fetcher()
+        for page in fetch_method(
+            from_date=from_date_iso,
+            to_date=to_date_iso,
+        ):
+            for item in page:
+                if "end_time" in item and isinstance(item["end_time"], str):
+                    item["end_time"] = pendulum.parse(item["end_time"])
+                yield item
+
+    @dlt.resource(write_disposition="replace")
+    def channels() -> Iterable[TDataItem]:
+        for page in fireflies_api.fetch_channels():
+            for item in page:
+                yield item
+
+    @dlt.resource(write_disposition="replace")
+    def users() -> Iterable[TDataItem]:
+        for page in fireflies_api.fetch_users():
+            for item in page:
+                yield item
+
+    @dlt.resource(write_disposition="replace")
+    def user_groups() -> Iterable[TDataItem]:
+        for page in fireflies_api.fetch_user_groups():
+            for item in page:
+                yield item
+
+    @dlt.resource(
+        write_disposition="merge",
+        primary_key="id",
+    )
+    def transcripts(
+        updated_at: dlt.sources.incremental[
+            pendulum.DateTime
+        ] = dlt.sources.incremental(
+            "date",
+            initial_value=start_datetime
+            if start_datetime
+            else pendulum.datetime(1970, 1, 1, tz="UTC"),
+            end_value=end_datetime if end_datetime else None,
+            range_end="closed" if end_datetime else "open",
+            range_start="closed",
+        ),
+    ) -> Iterable[TDataItem]:
+        from_date_dt = updated_at.last_value
+        to_date_dt = (
+            updated_at.end_value if updated_at.end_value else pendulum.now(tz="UTC")
+        )
+
+        from_date_iso = from_date_dt.to_iso8601_string() if from_date_dt else None
+        to_date_iso = to_date_dt.to_iso8601_string() if to_date_dt else None
+
+        for page in fireflies_api.fetch_transcripts(
+            from_date=from_date_iso,
+            to_date=to_date_iso,
+        ):
+            for item in page:
+                if "date" in item and isinstance(item["date"], (int, float)):
+                    item["date"] = pendulum.from_timestamp(item["date"] / 1000)
+                yield item
+
+    @dlt.resource(write_disposition="replace")
+    def bites() -> Iterable[TDataItem]:
+        for page in fireflies_api.fetch_bites():
+            for item in page:
+                yield item
+
+    @dlt.resource(write_disposition="replace")
+    def contacts() -> Iterable[TDataItem]:
+        for page in fireflies_api.fetch_contacts():
+            for item in page:
+                yield item
+
+    return [
+        active_meetings,
+        analytics,
+        channels,
+        users,
+        transcripts,
+        user_groups,
+        bites,
+        contacts,
+    ]