vedana-etl 0.1.0.dev3__py3-none-any.whl

vedana_etl/store.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple
+
+import pandas as pd
+from grist_api import GristDocAPI
+from sqlalchemy import Column
+
+from datapipe.run_config import RunConfig
+from datapipe.store.database import MetaKey
+from datapipe.store.table_store import TableStore, TableStoreCaps
+from datapipe.types import DataDF, DataSchema, IndexDF, MetaSchema
+
+
+@dataclass(frozen=True)
+class _GristAuth:
+    server: str
+    api_key: str
+    doc_id: str
+
+
+def _row_dicts_from_df(df: DataDF, columns: Sequence[str]) -> List[Dict[str, Any]]:
+    # Ensure JSON-serializable values: convert NaN to None
+    clean_df = df.copy()
+    if not clean_df.empty:
+        clean_df = clean_df.where(pd.notna(clean_df), None)
+    return [
+        {col: row.get(col) for col in columns}
+        for row in clean_df.to_dict(orient="records")  # type: ignore
+    ]
+
+
+def _pk_tuple_from_fields(fields: Dict[str, Any], pk_columns: Sequence[str]) -> Tuple[Any, ...]:
+    return tuple(fields.get(col) for col in pk_columns)
+
+
+class GristStore(TableStore):
+    """Grist document table backend.
+    Provided table is expected to exist and its schema is expected to match the provided schema.
+    Based on Grist Python Client: https://github.com/gristlabs/py_grist_api
+    """
+
+    caps = TableStoreCaps(
+        supports_delete=True,
+        supports_get_schema=True,
+        supports_read_all_rows=True,
+        supports_read_nonexistent_rows=True,
+        supports_read_meta_pseudo_df=True,
+    )
+
+    def __init__(
+        self,
+        server: str,
+        api_key: str,
+        doc_id: str,
+        table: str,
+        data_sql_schema: List[Column],
+        page_size: int = 1000,
+    ) -> None:
+        super().__init__()
+        self._auth = _GristAuth(server=server.rstrip("/"), api_key=api_key, doc_id=doc_id)
+        self.table = table
+        self.data_sql_schema = data_sql_schema
+        self._pk_columns: List[str] = [c.name for c in self.data_sql_schema if c.primary_key]
+        self._value_columns: List[str] = [c.name for c in self.data_sql_schema if not c.primary_key]
+        self._page_size = page_size
+
+        self._api = GristDocAPI(self._auth.doc_id, api_key=self._auth.api_key, server=self._auth.server)
+
+    # Basic metadata interface
+    def get_schema(self) -> DataSchema:
+        return self.data_sql_schema
+
+    def get_primary_schema(self) -> DataSchema:
+        return [c for c in self.data_sql_schema if c.primary_key]
+
+    def get_meta_schema(self) -> MetaSchema:
+        meta_key_prop = MetaKey.get_property_name()
+        return [column for column in self.data_sql_schema if hasattr(column, meta_key_prop)]
+
+    # Endpoints (for low-level access via client.call)
+    @property
+    def _records_path(self) -> str:
+        return f"tables/{self.table}/records"
+
+    # Low-level helpers
+    def _iter_all_records(self) -> Iterator[Dict[str, Any]]:
+        offset = 0
+        while True:
+            records = (
+                self._api.call(self._records_path, {"limit": self._page_size, "offset": offset}, method="GET")
+                .json()
+                .get("records")
+            )
+            if not records:
+                break
+            for rec in records:
+                yield rec
+            if len(records) < self._page_size:
+                break
+            offset += self._page_size
+
+    def _map_pk_to_record_id(self) -> Dict[Tuple[Any, ...], int]:
+        mapping: Dict[Tuple[Any, ...], int] = {}
+        for rec in self._iter_all_records():
+            rec_id_any = rec.get("id")
+            if not isinstance(rec_id_any, int):
+                continue
+            rec_id: int = rec_id_any
+            fields: Dict[str, Any] = rec.get("fields", {})
+            pk = _pk_tuple_from_fields(fields, self._pk_columns)
+            mapping[pk] = rec_id
+        return mapping
+
+    def _fetch_all_rows(self) -> List[Dict[str, Any]]:
+        data = self._api.fetch_table(self.table)
+        return [r._asdict() for r in data]  # namedtuples -> dict
+
+    # CRUD
+    def insert_rows(self, df: DataDF) -> None:
+        if df.empty:
+            return
+        all_columns = self._pk_columns + self._value_columns
+        rows = _row_dicts_from_df(df, all_columns)
+        self._api.add_records(self.table, rows)
+
+    def update_rows(self, df: DataDF) -> None:
+        if df.empty:
+            return
+
+        # Upsert using sync_table helper
+        all_columns = self._pk_columns + self._value_columns
+        new_data = [type("Row", (), row) for row in _row_dicts_from_df(df, all_columns)]
+        key_cols = [(c, c) for c in self._pk_columns]
+        other_cols = [(c, c) for c in self._value_columns]
+        self._api.sync_table(self.table, new_data, key_cols, other_cols)
+
+    def delete_rows(self, idx: IndexDF) -> None:
+        if idx is None or idx.empty:
+            return
+
+        # Map keys to record ids, then delete via SQL
+        pk_to_id = self._map_pk_to_record_id()
+
+        ids_to_delete: List[int] = []
+        for row in idx[self._pk_columns].to_dict(orient="records"):
+            pk = tuple(row.get(col) for col in self._pk_columns)
+            rec_id = pk_to_id.get(pk)
+            if rec_id is not None:
+                ids_to_delete.append(rec_id)
+
+        if not ids_to_delete:
+            return
+
+        self._api.delete_records(self.table, ids_to_delete)
+
+    def read_rows(self, idx: Optional[IndexDF] = None) -> DataDF:
+        if idx is None:  # Read full table
+            rows = self._fetch_all_rows()
+            if rows:
+                return pd.DataFrame.from_records(rows)[[c.name for c in self.data_sql_schema]]
+            return pd.DataFrame(columns=[c.name for c in self.data_sql_schema])
+
+        if idx.empty:
+            return pd.DataFrame(columns=[c.name for c in self.data_sql_schema])
+
+        # Build a local index of existing rows and filter
+        wanted: set[Tuple[Any, ...]] = set()
+        for row in idx[self._pk_columns].to_dict(orient="records"):
+            wanted.add(tuple(row.get(col) for col in self._pk_columns))
+
+        rows = self._fetch_all_rows()
+        if not rows:
+            return pd.DataFrame(columns=[c.name for c in self.data_sql_schema])
+        df_all = pd.DataFrame.from_records(rows)
+        if self._pk_columns:
+            tuples = list(map(tuple, df_all[self._pk_columns].astype(object).itertuples(index=False, name=None)))
+            mask = [t in wanted for t in tuples]
+            df_sel = df_all.loc[mask]
+        else:
+            df_sel = df_all.iloc[0:0]
+        if df_sel.empty:
+            return pd.DataFrame(columns=[c.name for c in self.data_sql_schema])
+        return df_sel[[c.name for c in self.data_sql_schema]]
+
+    def read_rows_meta_pseudo_df(
+        self,
+        chunksize: int = 1000,
+        run_config: Optional[RunConfig] = None,
+    ) -> Iterator[DataDF]:
+        # Stream records in pages and yield as DataFrames
+        buffer: List[Dict[str, Any]] = []
+
+        def flush() -> Iterator[DataDF]:
+            nonlocal buffer
+            if buffer:
+                df = pd.DataFrame.from_records(buffer)
+                buffer = []
+                yield df
+
+        rows = self._fetch_all_rows()
+        for row in rows:
+            buffer.append(row)
+            if len(buffer) >= chunksize:
+                yield from flush()
+
+        yield from flush()

vedana_etl-0.1.0.dev3.dist-info/METADATA

@@ -0,0 +1,51 @@
+Metadata-Version: 2.4
+Name: vedana-etl
+Version: 0.1.0.dev3
+Summary: Pipeline template for Vedana
+Author-email: Andrey Tatarinov <a@tatarinov.co>, Timur Sheydaev <tsheyd@epoch8.co>
+Requires-Python: >=3.12
+Requires-Dist: alembic>=1.16.1
+Requires-Dist: datapipe-app>=0.5.4
+Requires-Dist: datapipe-core>=0.14.3
+Requires-Dist: grist-api>=0.1.1
+Requires-Dist: neo4j>=5.28.1
+Requires-Dist: openai>=2.8.0
+Requires-Dist: pandas>=1.2.0
+Requires-Dist: pgvector>=0.4.2
+Requires-Dist: pytest>=8.4.1
+Requires-Dist: requests>=2.32.4
+Requires-Dist: sqlalchemy>=2.0.41
+Requires-Dist: vedana-core>=0.5.0
+Description-Content-Type: text/markdown
+
+# Basic pipeline for all vedana projects.
+
+This pipeline:
+
+- Parses Grist Data & Data Model
+- Ensures that Memgraph index/vector index structure is in sync with data model
+- Updates Memgraph database in incremental fashion
+
+To add steps:
+1. Pass extra transformations to [get_pipeline](src/pipeline.py)
+2. Create new app configuration from [app.py](src/app.py)
+
+## Pipeline Labels Hierarchy
+
+### Pipeline
+
+`labels=("pipeline", "pipeline_name")` defines a set of operations as standalone, sort of like a DAG in Airflow
+or a Dagster Job. Its purpose is to be able to render it as a separate tab on the ETL page of Backoffice in order to
+look at it independently of other transformations
+
+### Stage
+
+`labels=("stage", "stage_name")` defines a stage of `pipeline`. Currently, stages are useful for creating and managing
+observability features, such as [main dashboard's](/libs/vedana-backoffice/vedana_backoffice/pages/main_dashboard.py)
+Ingest table, which displays DataTable's of all transformations with `labels=("stage", "extract")`.
+Stages are also useful when running the pipeline manually.
+
+### Flow
+
+`labels=("flow", "flow_name")` helps execute a `pipeline` (or possibly several pipelines) in a nice fashion,
+used in defining cron jobs, etc.

vedana_etl-0.1.0.dev3.dist-info/RECORD

@@ -0,0 +1,13 @@
+vedana_etl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vedana_etl/app.py,sha256=-nD70kEUdMYpZSt6gW0rCUNCakCnbM47McD7E_VGlGQ,318
+vedana_etl/catalog.py,sha256=oNDSXQcOUG-D4IfleGgWGhe4H2oCpmG0Hxdj1TU2ekA,7432
+vedana_etl/config.py,sha256=z9L_EY_UsHK5lfmaNSkTj-f9Kxx43hZfZ5hURDZSnOs,603
+vedana_etl/pipeline.py,sha256=AcDcWBgauucIPe9mnL8Uar6Vv3jUqBXg0kINLwR5s9I,4145
+vedana_etl/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vedana_etl/schemas.py,sha256=OaV7zop6bkzIHZQh56oeH29wz7KVtm7wSEvBMJTsaIA,948
+vedana_etl/settings.py,sha256=9phXeq3-YlTsUnBkMvoMemzMXAQjMEYrouEpoXSbUxk,612
+vedana_etl/steps.py,sha256=27mS155c7CkuZPKQlao1cSnLhOy9qLja6MQ7Oqa2tKY,26081
+vedana_etl/store.py,sha256=5JbxJJ3Yhbzw_klZOwtImdW9jGLOiMU6VAGHXDNxf6Q,7386
+vedana_etl-0.1.0.dev3.dist-info/METADATA,sha256=mSiht1wRjN_cFkTzGBzQsH3Zndtm4JgOv2G_2KKDMug,1873
+vedana_etl-0.1.0.dev3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+vedana_etl-0.1.0.dev3.dist-info/RECORD,,

vedana_etl-0.1.0.dev3.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any