[OSO-731] add DataAnalytics feature to pyoso (#4723)

* feat(pyoso): add DataAnalytics feature to pyoso

* chore(pyoso): bump package version

* refactor(pyoso/analytics): make variables snake case

* feat(pyoso): add and use pydantic models for JSON validation

* fix(pyoso): fix README

Author: IcaroG

uv.lock

@@ -4,9 +4,11 @@ _WARNING: THIS IS A WORK IN PROGRESS_
 
 `pyoso` is a Python package for fetching models and metrics from OSO. This package provides an easy-to-use interface to interact with oso and retrieve valuable data for analysis and monitoring.
 
-## Current Features
+## Features
 
-- Execute custom SQL queries for analyzing the OSO dataset
+- Execute custom SQL queries for analyzing the OSO dataset.
+- Inspect data dependencies and freshness with an analytics tree.
+- Semantic modeling layer to build and execute complex queries (optional).
 
 ## Installation
 
@@ -28,12 +30,13 @@ This will include the `oso_semantic` package for building semantic models and qu
 
 ## Usage
 
-Here is a basic example of how to use `pyoso`:
+Here is a basic example of how to use `pyoso` to fetch data directly into a pandas DataFrame:
 
 ```python
+import os
 from pyoso import Client
 
-# Initialize the client
+# Initialize the client with an API key
 os.environ["OSO_API_KEY"] = 'your_api_key'
 client = Client()
 
@@ -44,6 +47,34 @@ artifacts = client.to_pandas(query)
 print(artifacts)
 ```
 
+### Inspecting Data Dependencies
+
+For more advanced use cases, the `client.query()` method returns a `QueryResponse` object that contains both the data and analytics metadata. This allows you to inspect the dependency tree of the data sources used in your query.
+
+```python
+import os
+from pyoso import Client
+
+# Initialize the client
+os.environ["OSO_API_KEY"] = "your_api_key"
+client = Client()
+
+# Execute a query to get a QueryResponse object
+query = "SELECT * FROM artifacts_v1 LIMIT 5"
+response = client.query(query)
+
+# You can still get the DataFrame as before
+df = response.to_pandas()
+print("\n--- Query Data ---")
+print(df)
+
+# Now, inspect the analytics to see the dependency tree
+print("\n--- Data Dependency Tree ---")
+response.analytics.print_tree()
+```
+
+This will output a tree structure showing how the final `artifacts_v1` table was constructed from its upstream dependencies, helping you understand the data's origin and freshness.
+
 ## Documentation
 
 For detailed documentation about the OSO dataset, please refer to the [official documentation](https://docs.opensource.observer/docs/integrate/datasets/).

warehouse/pyoso/README.md

@@ -7,5 +7,6 @@
 except PackageNotFoundError:
     __version__ = "unknown"
 
-from .client import Client
+from .analytics import DataAnalytics
+from .client import Client, QueryResponse
 from .exceptions import HTTPError, OsoError

warehouse/pyoso/pyoso/__init__.py

@@ -0,0 +1,139 @@
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+
+class PartitionStatusRange(BaseModel):
+    end_key: str = Field(alias="endKey")
+    start_key: str = Field(alias="startKey")
+    status: str
+
+
+class PartitionStatus(BaseModel):
+    num_failed: int = Field(alias="numFailed")
+    num_materialized: int = Field(alias="numMaterialized")
+    num_materializing: int = Field(alias="numMaterializing")
+    num_partitions: int = Field(alias="numPartitions")
+    ranges: list[PartitionStatusRange]
+
+
+class MaterializationStatus(BaseModel):
+    partition_status: PartitionStatus | None = Field(
+        default=None, alias="partitionStatus"
+    )
+    latest_materialization: datetime | None = Field(
+        default=None, alias="latestMaterialization"
+    )
+
+
+class DataStatus(BaseModel):
+    key: str
+    status: MaterializationStatus
+    dependencies: list[str]
+
+
+class DataAnalytics:
+    """Container for analytics data with tree-structured display methods."""
+
+    def __init__(self, analytics_data: dict[str, DataStatus]):
+        self._analytics_data = analytics_data
+        self._root_keys = self._calculate_root_keys()
+
+    def _calculate_root_keys(self) -> list[str]:
+        """Calculate root nodes (nodes that are not dependencies of others)."""
+        all_dependencies = set()
+        for data_status in self._analytics_data.values():
+            all_dependencies.update(data_status.dependencies)
+
+        root_keys = [
+            k for k in self._analytics_data.keys() if k not in all_dependencies
+        ]
+
+        # If no clear roots, return all keys
+        if not root_keys:
+            root_keys = list(self._analytics_data.keys())
+
+        return root_keys
+
+    def __iter__(self):
+        """Iterate over the analytics data keys."""
+        return iter(self._analytics_data.keys())
+
+    def __contains__(self, key: str) -> bool:
+        """Check if a key exists in the analytics data."""
+        return key in self._analytics_data
+
+    def __len__(self):
+        """Get the number of analytics entries."""
+        return len(self._analytics_data)
+
+    @property
+    def root_keys(self) -> list[str]:
+        """Get the root keys (top-level nodes in the dependency tree)."""
+        return self._root_keys.copy()
+
+    def get(self, key: str) -> DataStatus | None:
+        """Get analytics data for a specific key."""
+        return self._analytics_data.get(key)
+
+    def print_tree(self, key: str | None = None):
+        """Print analytics data as a tree structure.
+
+        Args:
+            key: If provided, print analytics for this specific key and its dependencies.
+                 If None, print analytics for all root keys.
+        """
+        if key is not None:
+            self._print_analytics_tree(key, set())
+        else:
+            # Print all root nodes
+            num_roots = len(self._root_keys)
+            for i, root_key in enumerate(self._root_keys):
+                self._print_analytics_tree(
+                    root_key, set(), is_last=(i == num_roots - 1)
+                )
+
+    def _print_analytics_tree(
+        self, key: str, visited: set, indent: str = "", is_last: bool = True
+    ):
+        """Recursively print analytics tree for a given key."""
+        if key in visited:
+            print(f"{indent}├──{key} (circular dependency)")
+            return
+
+        visited.add(key)
+        data_status = self._analytics_data[key]
+
+        # Print the current node with inline status information
+        status = data_status.status
+        status_parts = []
+
+        if not status.latest_materialization and not status.partition_status:
+            status_parts.append("No analytics data")
+        if status.latest_materialization:
+            status_parts.append(
+                f"Last: {status.latest_materialization.strftime('%Y-%m-%d %H:%M:%S')}"
+            )
+        if status.partition_status:
+            ps = status.partition_status
+            status_parts.append(
+                f"Partitions: {ps.num_materialized}/{ps.num_partitions}"
+            )
+
+        status_text = f" ({', '.join(status_parts)})" if status_parts else ""
+
+        # Determine the prefix for the current node
+        prefix = "└── " if is_last else "├── "
+        print(f"{indent}{prefix}{key}{status_text}")
+
+        # Print dependencies
+        if data_status.dependencies:
+            num_dependencies = len(data_status.dependencies)
+            for i, dep in enumerate(data_status.dependencies):
+                is_last_dep = i == num_dependencies - 1
+                dep_indent = indent + ("    " if is_last else "│   ")
+                self._print_analytics_tree(
+                    dep, visited.copy(), dep_indent, is_last=is_last_dep
+                )
+
+        visited.remove(key)

warehouse/pyoso/pyoso/analytics.py

@@ -1,10 +1,11 @@
 import json
 import os
-from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import pandas as pd
 import requests
+from pydantic import BaseModel
+from pyoso.analytics import DataAnalytics, DataStatus
 from pyoso.constants import DEFAULT_BASE_URL, OSO_API_KEY
 from pyoso.exceptions import OsoError, OsoHTTPError
 from pyoso.semantic import create_registry
@@ -19,20 +20,66 @@
     pass
 
 
-@dataclass
-class ClientConfig:
-    base_url: Optional[str]
+class ClientConfig(BaseModel):
+    base_url: str | None
 
 
-@dataclass
-class QueryData:
+class QueryData(BaseModel):
     columns: list[str]
     data: list[list[Any]]
 
 
+class QueryResponse:
+    """Response object containing query data and optional analytics."""
+
+    def __init__(self, data: QueryData, analytics: DataAnalytics):
+        self._data = data
+        self._analytics = analytics
+
+    def to_pandas(self) -> pd.DataFrame:
+        """Convert the query data to a pandas DataFrame."""
+        return pd.DataFrame(
+            self._data.data, columns=self._data.columns
+        ).convert_dtypes()
+
+    @property
+    def analytics(self) -> DataAnalytics:
+        """Get the analytics data."""
+        return self._analytics
+
+    @property
+    def data(self) -> QueryData:
+        """Get the raw query data."""
+        return self._data
+
+    @staticmethod
+    def from_response_chunks(response_chunks) -> "QueryResponse":
+        """Parse HTTP response chunks into QueryResponse."""
+        columns = []
+        data = []
+        analytics = {}
+
+        for chunk in response_chunks:
+            if chunk:
+                parsed_obj = json.loads(chunk)
+
+                if "assetStatus" in parsed_obj:
+                    for asset in parsed_obj["assetStatus"]:
+                        data_status = DataStatus.model_validate(asset)
+                        analytics[data_status.key] = data_status
+                elif "columns" in parsed_obj:
+                    columns.extend(parsed_obj["columns"])
+
+                if "data" in parsed_obj:
+                    data.extend(parsed_obj["data"])
+
+        query_data = QueryData(columns=columns, data=data)
+        return QueryResponse(data=query_data, analytics=DataAnalytics(analytics))
+
+
 class Client:
     def __init__(
-        self, api_key: Optional[str] = None, client_opts: Optional[ClientConfig] = None
+        self, api_key: str | None = None, client_opts: ClientConfig | None = None
     ):
         self.__api_key = api_key if api_key else os.environ.get(OSO_API_KEY)
         if not self.__api_key:
@@ -51,8 +98,12 @@ def __init__(
             )
 
     def __query(
-        self, query: str, input_dialect="trino", output_dialect="trino"
-    ) -> QueryData:
+        self,
+        query: str,
+        input_dialect="trino",
+        output_dialect="trino",
+        include_analytics: bool = True,
+    ) -> QueryResponse:
         # The following checks are only for providing better error messages as
         # the oso api does _not_ support multiple queries nor the use of
         # semicolons.
@@ -78,26 +129,22 @@ def __query(
                 json={
                     "query": query_expression.sql(dialect=output_dialect),
                     "format": "minimal",
+                    "includeAnalytics": include_analytics,
                 },
                 stream=True,
             )
             response.raise_for_status()
-            columns = []
-            data = []
-            for chunk in response.iter_lines(chunk_size=None):
-                if chunk:
-                    parsed_obj = json.loads(chunk)
-                    if "columns" in parsed_obj:
-                        columns.extend(parsed_obj["columns"])
-                    if "data" in parsed_obj:
-                        data.extend(parsed_obj["data"])
-
-            return QueryData(columns=columns, data=data)
+
+            return QueryResponse.from_response_chunks(
+                response.iter_lines(chunk_size=None)
+            )
         except requests.HTTPError as e:
             raise OsoHTTPError(e, response=e.response) from None
 
     def to_pandas(self, query: str):
-        query_data = self.__query(query)
-        return pd.DataFrame(
-            query_data.data, columns=query_data.columns
-        ).convert_dtypes()
+        query_response = self.__query(query)
+        return query_response.to_pandas()
+
+    def query(self, query: str, include_analytics: bool = True) -> QueryResponse:
+        """Execute a SQL query and return the full response including analytics data."""
+        return self.__query(query, include_analytics=include_analytics)