robosystems-client 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

robosystems_client/extensions/table_ingest_client.py

@@ -1,462 +0,0 @@
-"""Table Ingest Client for RoboSystems API
-
-Simplifies uploading Parquet files to staging tables and ingesting them into graphs.
-"""
-
-from dataclasses import dataclass
-from io import BytesIO
-from pathlib import Path
-from typing import Dict, Any, Optional, Callable, List, Union, BinaryIO
-import json
-import logging
-import httpx
-
-from ..api.tables.get_upload_url import (
-  sync_detailed as get_upload_url,
-)
-from ..api.tables.update_file_status import (
-  sync_detailed as update_file_status,
-)
-from ..api.tables.list_tables import (
-  sync_detailed as list_tables,
-)
-from ..api.tables.ingest_tables import (
-  sync_detailed as ingest_tables,
-)
-from ..models.file_upload_request import FileUploadRequest
-from ..models.file_status_update import FileStatusUpdate
-from ..models.bulk_ingest_request import BulkIngestRequest
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class UploadOptions:
-  """Options for file upload operations"""
-
-  on_progress: Optional[Callable[[str], None]] = None
-  fix_localstack_url: bool = True  # Auto-fix LocalStack URLs for localhost
-  file_name: Optional[str] = None  # Override file name (useful for buffer uploads)
-
-
-@dataclass
-class IngestOptions:
-  """Options for table ingestion operations"""
-
-  ignore_errors: bool = True
-  rebuild: bool = False
-  on_progress: Optional[Callable[[str], None]] = None
-
-
-@dataclass
-class UploadResult:
-  """Result from file upload operation"""
-
-  file_id: str
-  file_size: int
-  row_count: int
-  table_name: str
-  file_name: str
-  success: bool = True
-  error: Optional[str] = None
-
-
-@dataclass
-class TableInfo:
-  """Information about a staging table"""
-
-  table_name: str
-  row_count: int
-  file_count: int
-  total_size_bytes: int
-
-
-class TableIngestClient:
-  """Enhanced table ingest client with simplified upload workflow"""
-
-  def __init__(self, config: Dict[str, Any]):
-    self.config = config
-    self.base_url = config["base_url"]
-    self.headers = config.get("headers", {})
-    self.token = config.get("token")
-    # Create httpx client for S3 uploads
-    self._http_client = httpx.Client(timeout=120.0)
-
-  def upload_parquet_file(
-    self,
-    graph_id: str,
-    table_name: str,
-    file_or_buffer: Union[Path, str, BytesIO, BinaryIO],
-    options: Optional[UploadOptions] = None,
-  ) -> UploadResult:
-    """
-    Upload a Parquet file to a staging table.
-
-    This method handles the complete 3-step upload process:
-    1. Get presigned upload URL
-    2. Upload file to S3
-    3. Mark file as 'uploaded' (backend validates, calculates size/row count)
-
-    Args:
-        graph_id: The graph ID
-        table_name: Name of the staging table
-        file_or_buffer: Path to the Parquet file or BytesIO/BinaryIO buffer
-        options: Upload options
-
-    Returns:
-        UploadResult with upload details (size/row count calculated by backend)
-    """
-    if options is None:
-      options = UploadOptions()
-
-    # Auto-detect if this is a file path or buffer
-    is_buffer = isinstance(file_or_buffer, (BytesIO, BinaryIO)) or hasattr(
-      file_or_buffer, "read"
-    )
-
-    # Initialize file_path for type checking
-    file_path: Optional[Path] = None
-
-    if is_buffer:
-      # Handle buffer upload
-      file_name = options.file_name or "data.parquet"
-    else:
-      # Handle file path upload
-      file_path = Path(file_or_buffer)
-      file_name = file_path.name
-      if not file_path.exists():
-        return UploadResult(
-          file_id="",
-          file_size=0,
-          row_count=0,
-          table_name=table_name,
-          file_name=file_name,
-          success=False,
-          error=f"File not found: {file_path}",
-        )
-
-    try:
-      # Import client here to avoid circular imports
-      from ..client import AuthenticatedClient
-
-      # Create authenticated client with X-API-Key
-      # The token is extracted from X-API-Key header in extensions.py
-      if not self.token:
-        return UploadResult(
-          file_id="",
-          file_size=0,
-          row_count=0,
-          table_name=table_name,
-          file_name=file_name,
-          success=False,
-          error="No API key provided. Set X-API-Key in headers.",
-        )
-
-      client = AuthenticatedClient(
-        base_url=self.base_url,
-        token=self.token,
-        prefix="",  # No prefix for X-API-Key
-        auth_header_name="X-API-Key",  # Use X-API-Key header instead of Authorization
-        headers=self.headers,
-      )
-
-      # Step 1: Get presigned upload URL
-      if options.on_progress:
-        options.on_progress(
-          f"Getting upload URL for {file_name} -> table '{table_name}'..."
-        )
-
-      upload_request = FileUploadRequest(
-        file_name=file_name, content_type="application/x-parquet"
-      )
-
-      kwargs = {
-        "graph_id": graph_id,
-        "table_name": table_name,
-        "client": client,
-        "body": upload_request,
-      }
-
-      response = get_upload_url(**kwargs)
-
-      if not response.parsed:
-        error_msg = f"Failed to get upload URL (status: {response.status_code})"
-        if hasattr(response, "content"):
-          try:
-            error_detail = json.loads(response.content)
-            error_msg = f"{error_msg}: {error_detail}"
-          except (json.JSONDecodeError, ValueError):
-            error_msg = f"{error_msg}: {response.content[:200]}"
-
-        return UploadResult(
-          file_id="",
-          file_size=0,
-          row_count=0,
-          table_name=table_name,
-          file_name=file_name,
-          success=False,
-          error=error_msg,
-        )
-
-      upload_url = response.parsed.upload_url
-      file_id = response.parsed.file_id
-
-      # Fix LocalStack URL if needed
-      if options.fix_localstack_url and "localstack:4566" in upload_url:
-        upload_url = upload_url.replace("localstack:4566", "localhost:4566")
-        logger.debug("Fixed LocalStack URL for localhost access")
-
-      # Step 2: Upload file to S3
-      if options.on_progress:
-        options.on_progress(f"Uploading {file_name} to S3...")
-
-      # Read file content - handle both paths and buffers
-      if is_buffer:
-        # Read from buffer
-        if hasattr(file_or_buffer, "getvalue"):
-          file_content = file_or_buffer.getvalue()
-        else:
-          # BinaryIO or file-like object
-          file_or_buffer.seek(0)
-          file_content = file_or_buffer.read()
-      else:
-        # Read from file path
-        if file_path is None:
-          raise ValueError("file_path should not be None when not using buffer")
-        with open(file_path, "rb") as f:
-          file_content = f.read()
-
-      s3_response = self._http_client.put(
-        upload_url,
-        content=file_content,
-        headers={"Content-Type": "application/x-parquet"},
-      )
-      s3_response.raise_for_status()
-
-      # Step 3: Mark file as uploaded (backend validates and calculates size/row count)
-      if options.on_progress:
-        options.on_progress(f"Marking {file_name} as uploaded...")
-
-      status_update = FileStatusUpdate(status="uploaded")
-
-      kwargs = {
-        "graph_id": graph_id,
-        "file_id": file_id,
-        "client": client,
-        "body": status_update,
-      }
-
-      update_response = update_file_status(**kwargs)
-
-      if not update_response.parsed:
-        logger.error(
-          f"No parsed response from update_file_status. Status code: {update_response.status_code}"
-        )
-        return UploadResult(
-          file_id=file_id,
-          file_size=0,
-          row_count=0,
-          table_name=table_name,
-          file_name=file_name,
-          success=False,
-          error="Failed to complete file upload",
-        )
-
-      response_data = update_response.parsed
-
-      if isinstance(response_data, dict):
-        file_size = response_data.get("file_size_bytes", 0)
-        row_count = response_data.get("row_count", 0)
-      elif hasattr(response_data, "additional_properties"):
-        file_size = response_data.additional_properties.get("file_size_bytes", 0)
-        row_count = response_data.additional_properties.get("row_count", 0)
-      else:
-        file_size = getattr(response_data, "file_size_bytes", 0)
-        row_count = getattr(response_data, "row_count", 0)
-
-      if options.on_progress:
-        options.on_progress(
-          f"✅ Uploaded {file_name} ({file_size:,} bytes, {row_count:,} rows)"
-        )
-
-      return UploadResult(
-        file_id=file_id,
-        file_size=file_size,
-        row_count=row_count,
-        table_name=table_name,
-        file_name=file_name,
-        success=True,
-      )
-
-    except Exception as e:
-      logger.error(f"Upload failed for {file_name}: {e}")
-      return UploadResult(
-        file_id="",
-        file_size=0,
-        row_count=0,
-        table_name=table_name,
-        file_name=file_name,
-        success=False,
-        error=str(e),
-      )
-
-  def list_staging_tables(self, graph_id: str) -> List[TableInfo]:
-    """
-    List all staging tables in a graph.
-
-    Args:
-        graph_id: The graph ID
-
-    Returns:
-        List of TableInfo objects
-    """
-    try:
-      from ..client import AuthenticatedClient
-
-      if not self.token:
-        logger.error("No API key provided")
-        return []
-
-      client = AuthenticatedClient(
-        base_url=self.base_url,
-        token=self.token,
-        prefix="",
-        auth_header_name="X-API-Key",
-        headers=self.headers,
-      )
-
-      kwargs = {"graph_id": graph_id, "client": client}
-
-      response = list_tables(**kwargs)
-
-      if not response.parsed:
-        logger.error("Failed to list tables")
-        return []
-
-      tables = []
-      for table_data in response.parsed.tables:
-        tables.append(
-          TableInfo(
-            table_name=table_data.table_name,
-            row_count=table_data.row_count,
-            file_count=table_data.file_count,
-            total_size_bytes=table_data.total_size_bytes,
-          )
-        )
-
-      return tables
-
-    except Exception as e:
-      logger.error(f"Failed to list tables: {e}")
-      return []
-
-  def ingest_all_tables(
-    self, graph_id: str, options: Optional[IngestOptions] = None
-  ) -> Dict[str, Any]:
-    """
-    Ingest all staging tables into the graph.
-
-    Args:
-        graph_id: The graph ID
-        options: Ingest options
-
-    Returns:
-        Dictionary with ingestion results
-    """
-    if options is None:
-      options = IngestOptions()
-
-    try:
-      from ..client import AuthenticatedClient
-
-      if not self.token:
-        return {"success": False, "error": "No API key provided"}
-
-      client = AuthenticatedClient(
-        base_url=self.base_url,
-        token=self.token,
-        prefix="",
-        auth_header_name="X-API-Key",
-        headers=self.headers,
-      )
-
-      if options.on_progress:
-        options.on_progress("Starting table ingestion...")
-
-      ingest_request = BulkIngestRequest(
-        ignore_errors=options.ignore_errors, rebuild=options.rebuild
-      )
-
-      kwargs = {
-        "graph_id": graph_id,
-        "client": client,
-        "body": ingest_request,
-      }
-
-      response = ingest_tables(**kwargs)
-
-      if not response.parsed:
-        return {"success": False, "error": "Failed to ingest tables"}
-
-      result = {
-        "success": True,
-        "operation_id": getattr(response.parsed, "operation_id", None),
-        "message": getattr(response.parsed, "message", "Ingestion started"),
-      }
-
-      if options.on_progress:
-        options.on_progress("✅ Table ingestion completed")
-
-      return result
-
-    except Exception as e:
-      logger.error(f"Failed to ingest tables: {e}")
-      return {"success": False, "error": str(e)}
-
-  def upload_and_ingest(
-    self,
-    graph_id: str,
-    table_name: str,
-    file_path: Path,
-    upload_options: Optional[UploadOptions] = None,
-    ingest_options: Optional[IngestOptions] = None,
-  ) -> Dict[str, Any]:
-    """
-    Convenience method to upload a file and immediately ingest it.
-
-    Args:
-        graph_id: The graph ID
-        table_name: Name of the staging table
-        file_path: Path to the Parquet file
-        upload_options: Upload options
-        ingest_options: Ingest options
-
-    Returns:
-        Dictionary with upload and ingest results
-    """
-    # Upload the file
-    upload_result = self.upload_parquet_file(
-      graph_id, table_name, file_path, upload_options
-    )
-
-    if not upload_result.success:
-      return {
-        "success": False,
-        "upload": upload_result,
-        "ingest": None,
-        "error": upload_result.error,
-      }
-
-    # Ingest the table
-    ingest_result = self.ingest_all_tables(graph_id, ingest_options)
-
-    return {
-      "success": upload_result.success and ingest_result.get("success", False),
-      "upload": upload_result,
-      "ingest": ingest_result,
-    }
-
-  def close(self):
-    """Close HTTP client connections"""
-    if self._http_client:
-      self._http_client.close()

robosystems_client/models/bulk_ingest_response.py

@@ -1,137 +0,0 @@
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, TypeVar
-
-from attrs import define as _attrs_define
-from attrs import field as _attrs_field
-
-if TYPE_CHECKING:
-  from ..models.table_ingest_result import TableIngestResult
-
-
-T = TypeVar("T", bound="BulkIngestResponse")
-
-
-@_attrs_define
-class BulkIngestResponse:
-  """
-  Attributes:
-      status (str): Overall ingestion status
-      graph_id (str): Graph database identifier
-      total_tables (int): Total number of tables processed
-      successful_tables (int): Number of successfully ingested tables
-      failed_tables (int): Number of failed table ingestions
-      skipped_tables (int): Number of skipped tables (no files)
-      total_rows_ingested (int): Total rows ingested across all tables
-      total_execution_time_ms (float): Total execution time in milliseconds
-      results (list['TableIngestResult']): Per-table ingestion results
-  """
-
-  status: str
-  graph_id: str
-  total_tables: int
-  successful_tables: int
-  failed_tables: int
-  skipped_tables: int
-  total_rows_ingested: int
-  total_execution_time_ms: float
-  results: list["TableIngestResult"]
-  additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
-
-  def to_dict(self) -> dict[str, Any]:
-    status = self.status
-
-    graph_id = self.graph_id
-
-    total_tables = self.total_tables
-
-    successful_tables = self.successful_tables
-
-    failed_tables = self.failed_tables
-
-    skipped_tables = self.skipped_tables
-
-    total_rows_ingested = self.total_rows_ingested
-
-    total_execution_time_ms = self.total_execution_time_ms
-
-    results = []
-    for results_item_data in self.results:
-      results_item = results_item_data.to_dict()
-      results.append(results_item)
-
-    field_dict: dict[str, Any] = {}
-    field_dict.update(self.additional_properties)
-    field_dict.update(
-      {
-        "status": status,
-        "graph_id": graph_id,
-        "total_tables": total_tables,
-        "successful_tables": successful_tables,
-        "failed_tables": failed_tables,
-        "skipped_tables": skipped_tables,
-        "total_rows_ingested": total_rows_ingested,
-        "total_execution_time_ms": total_execution_time_ms,
-        "results": results,
-      }
-    )
-
-    return field_dict
-
-  @classmethod
-  def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
-    from ..models.table_ingest_result import TableIngestResult
-
-    d = dict(src_dict)
-    status = d.pop("status")
-
-    graph_id = d.pop("graph_id")
-
-    total_tables = d.pop("total_tables")
-
-    successful_tables = d.pop("successful_tables")
-
-    failed_tables = d.pop("failed_tables")
-
-    skipped_tables = d.pop("skipped_tables")
-
-    total_rows_ingested = d.pop("total_rows_ingested")
-
-    total_execution_time_ms = d.pop("total_execution_time_ms")
-
-    results = []
-    _results = d.pop("results")
-    for results_item_data in _results:
-      results_item = TableIngestResult.from_dict(results_item_data)
-
-      results.append(results_item)
-
-    bulk_ingest_response = cls(
-      status=status,
-      graph_id=graph_id,
-      total_tables=total_tables,
-      successful_tables=successful_tables,
-      failed_tables=failed_tables,
-      skipped_tables=skipped_tables,
-      total_rows_ingested=total_rows_ingested,
-      total_execution_time_ms=total_execution_time_ms,
-      results=results,
-    )
-
-    bulk_ingest_response.additional_properties = d
-    return bulk_ingest_response
-
-  @property
-  def additional_keys(self) -> list[str]:
-    return list(self.additional_properties.keys())
-
-  def __getitem__(self, key: str) -> Any:
-    return self.additional_properties[key]
-
-  def __setitem__(self, key: str, value: Any) -> None:
-    self.additional_properties[key] = value
-
-  def __delitem__(self, key: str) -> None:
-    del self.additional_properties[key]
-
-  def __contains__(self, key: str) -> bool:
-    return key in self.additional_properties

robosystems_client/models/table_ingest_result.py

@@ -1,107 +0,0 @@
-from collections.abc import Mapping
-from typing import Any, TypeVar, Union, cast
-
-from attrs import define as _attrs_define
-from attrs import field as _attrs_field
-
-from ..types import UNSET, Unset
-
-T = TypeVar("T", bound="TableIngestResult")
-
-
-@_attrs_define
-class TableIngestResult:
-  """
-  Attributes:
-      table_name (str): Table name
-      status (str): Ingestion status (success/failed/skipped)
-      rows_ingested (Union[Unset, int]): Number of rows ingested Default: 0.
-      execution_time_ms (Union[Unset, float]): Ingestion time in milliseconds Default: 0.0.
-      error (Union[None, Unset, str]): Error message if failed
-  """
-
-  table_name: str
-  status: str
-  rows_ingested: Union[Unset, int] = 0
-  execution_time_ms: Union[Unset, float] = 0.0
-  error: Union[None, Unset, str] = UNSET
-  additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
-
-  def to_dict(self) -> dict[str, Any]:
-    table_name = self.table_name
-
-    status = self.status
-
-    rows_ingested = self.rows_ingested
-
-    execution_time_ms = self.execution_time_ms
-
-    error: Union[None, Unset, str]
-    if isinstance(self.error, Unset):
-      error = UNSET
-    else:
-      error = self.error
-
-    field_dict: dict[str, Any] = {}
-    field_dict.update(self.additional_properties)
-    field_dict.update(
-      {
-        "table_name": table_name,
-        "status": status,
-      }
-    )
-    if rows_ingested is not UNSET:
-      field_dict["rows_ingested"] = rows_ingested
-    if execution_time_ms is not UNSET:
-      field_dict["execution_time_ms"] = execution_time_ms
-    if error is not UNSET:
-      field_dict["error"] = error
-
-    return field_dict
-
-  @classmethod
-  def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
-    d = dict(src_dict)
-    table_name = d.pop("table_name")
-
-    status = d.pop("status")
-
-    rows_ingested = d.pop("rows_ingested", UNSET)
-
-    execution_time_ms = d.pop("execution_time_ms", UNSET)
-
-    def _parse_error(data: object) -> Union[None, Unset, str]:
-      if data is None:
-        return data
-      if isinstance(data, Unset):
-        return data
-      return cast(Union[None, Unset, str], data)
-
-    error = _parse_error(d.pop("error", UNSET))
-
-    table_ingest_result = cls(
-      table_name=table_name,
-      status=status,
-      rows_ingested=rows_ingested,
-      execution_time_ms=execution_time_ms,
-      error=error,
-    )
-
-    table_ingest_result.additional_properties = d
-    return table_ingest_result
-
-  @property
-  def additional_keys(self) -> list[str]:
-    return list(self.additional_properties.keys())
-
-  def __getitem__(self, key: str) -> Any:
-    return self.additional_properties[key]
-
-  def __setitem__(self, key: str, value: Any) -> None:
-    self.additional_properties[key] = value
-
-  def __delitem__(self, key: str) -> None:
-    del self.additional_properties[key]
-
-  def __contains__(self, key: str) -> bool:
-    return key in self.additional_properties