robosystems-client 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

robosystems_client/extensions/operation_client.py

@@ -71,7 +71,14 @@ class OperationClient:
   def __init__(self, config: Dict[str, Any]):
     self.config = config
     self.base_url = config["base_url"]
+    self.headers = config.get("headers", {})
+    # Get token from config if passed by parent
+    self.token = config.get("token")
     self.active_operations: Dict[str, SSEClient] = {}
+    # Thread safety for operations tracking
+    import threading
+
+    self._lock = threading.Lock()
 
   def monitor_operation(
     self, operation_id: str, options: MonitorOptions = None
@@ -144,7 +151,8 @@ class OperationClient:
     # Connect and monitor
     try:
       sse_client.connect(operation_id)
-      self.active_operations[operation_id] = sse_client
+      with self._lock:
+        self.active_operations[operation_id] = sse_client
 
       # Wait for completion
       import time
@@ -166,10 +174,11 @@ class OperationClient:
         time.sleep(options.poll_interval or 0.1)
 
     finally:
-      # Clean up
-      if operation_id in self.active_operations:
-        self.active_operations[operation_id].close()
-        del self.active_operations[operation_id]
+      # Clean up with thread safety
+      with self._lock:
+        if operation_id in self.active_operations:
+          self.active_operations[operation_id].close()
+          del self.active_operations[operation_id]
 
     return result
 
@@ -179,11 +188,16 @@ class OperationClient:
     from ..api.operations.get_operation_status import (
       sync_detailed as get_operation_status,
     )
-    from ..client import AuthenticatedClient
+    from ..client import Client
 
-    client = AuthenticatedClient(base_url=self.base_url)
+    # Use regular Client with headers instead of AuthenticatedClient
+    client = Client(base_url=self.base_url, headers=self.headers)
     try:
-      response = get_operation_status(operation_id=operation_id, client=client)
+      kwargs = {"operation_id": operation_id, "client": client}
+      # Only add token if it's a valid string
+      if self.token and isinstance(self.token, str) and self.token.strip():
+        kwargs["token"] = self.token
+      response = get_operation_status(**kwargs)
       if response.parsed:
         return {
           "operation_id": operation_id,
@@ -201,21 +215,27 @@ class OperationClient:
     """Cancel an operation"""
     # This would use the generated SDK to call /v1/operations/{operation_id}/cancel
     from ..api.operations.cancel_operation import sync_detailed as cancel_operation
-    from ..client import AuthenticatedClient
+    from ..client import Client
 
-    client = AuthenticatedClient(base_url=self.base_url)
+    # Use regular Client with headers instead of AuthenticatedClient
+    client = Client(base_url=self.base_url, headers=self.headers)
     try:
-      response = cancel_operation(operation_id=operation_id, client=client)
+      kwargs = {"operation_id": operation_id, "client": client}
+      # Only add token if it's a valid string
+      if self.token and isinstance(self.token, str) and self.token.strip():
+        kwargs["token"] = self.token
+      response = cancel_operation(**kwargs)
       if response.parsed:
         return response.parsed.cancelled or False
     except Exception as e:
       print(f"Failed to cancel operation {operation_id}: {e}")
       return False
 
-    # Also close any active SSE connection
-    if operation_id in self.active_operations:
-      self.active_operations[operation_id].close()
-      del self.active_operations[operation_id]
+    # Also close any active SSE connection with thread safety
+    with self._lock:
+      if operation_id in self.active_operations:
+        self.active_operations[operation_id].close()
+        del self.active_operations[operation_id]
 
     return False
 
@@ -226,15 +246,17 @@ class OperationClient:
 
   def close_all(self):
     """Close all active operation monitors"""
-    for sse_client in self.active_operations.values():
-      sse_client.close()
-    self.active_operations.clear()
+    with self._lock:
+      for sse_client in self.active_operations.values():
+        sse_client.close()
+      self.active_operations.clear()
 
   def close_operation(self, operation_id: str):
     """Close monitoring for a specific operation"""
-    if operation_id in self.active_operations:
-      self.active_operations[operation_id].close()
-      del self.active_operations[operation_id]
+    with self._lock:
+      if operation_id in self.active_operations:
+        self.active_operations[operation_id].close()
+        del self.active_operations[operation_id]
 
 
 class AsyncOperationClient:

robosystems_client/extensions/query_client.py

@@ -4,7 +4,17 @@ Provides intelligent query execution with automatic strategy selection.
 """
 
 from dataclasses import dataclass
-from typing import Dict, Any, Optional, Callable, AsyncIterator, Iterator, Union
+from typing import (
+  Dict,
+  Any,
+  Optional,
+  Callable,
+  AsyncIterator,
+  Iterator,
+  Union,
+  Generator,
+  List,
+)
 from datetime import datetime
 
 from ..api.query.execute_cypher_query import sync_detailed as execute_cypher_query
@@ -70,6 +80,9 @@ class QueryClient:
   def __init__(self, config: Dict[str, Any]):
     self.config = config
     self.base_url = config["base_url"]
+    self.headers = config.get("headers", {})
+    # Get token from config if passed by parent
+    self.token = config.get("token")
     self.sse_client: Optional[SSEClient] = None
 
   def execute_query(
@@ -85,15 +98,17 @@ class QueryClient:
     )
 
     # Execute the query through the generated client
-    from ..client import AuthenticatedClient
+    from ..client import Client
 
-    # Get client instance (you'd configure this based on your setup)
-    client = AuthenticatedClient(base_url=self.base_url)
+    # Create client with headers
+    client = Client(base_url=self.base_url, headers=self.headers)
 
     try:
-      response = execute_cypher_query(
-        graph_id=graph_id, client=client, body=query_request
-      )
+      kwargs = {"graph_id": graph_id, "client": client, "body": query_request}
+      # Only add token if it's a valid string
+      if self.token and isinstance(self.token, str) and self.token.strip():
+        kwargs["token"] = self.token
+      response = execute_cypher_query(**kwargs)
 
       # Check response type and handle accordingly
       if hasattr(response, "parsed") and response.parsed:
@@ -145,7 +160,15 @@ class QueryClient:
     except Exception as e:
       if isinstance(e, QueuedQueryError):
         raise
-      raise Exception(f"Query execution failed: {str(e)}")
+
+      error_msg = str(e)
+      # Check for authentication errors
+      if (
+        "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg.lower()
+      ):
+        raise Exception(f"Authentication failed during query execution: {error_msg}")
+      else:
+        raise Exception(f"Query execution failed: {error_msg}")
 
     # Unexpected response format
     raise Exception("Unexpected response format from query endpoint")
@@ -316,18 +339,92 @@ class QueryClient:
     cypher: str,
     parameters: Dict[str, Any] = None,
     chunk_size: int = 1000,
-  ) -> Iterator[Any]:
-    """Streaming query for large results"""
+    on_progress: Optional[Callable[[int, int], None]] = None,
+  ) -> Generator[Any, None, None]:
+    """Stream query results for large datasets with progress tracking
+
+    Args:
+        graph_id: Graph ID to query
+        cypher: Cypher query string
+        parameters: Query parameters
+        chunk_size: Number of records per chunk
+        on_progress: Callback for progress updates (current, total)
+
+    Yields:
+        Individual records from query results
+
+    Example:
+        >>> def progress(current, total):
+        ...     print(f"Processed {current}/{total} records")
+        >>> for record in query_client.stream_query(
+        ...     'graph_id',
+        ...     'MATCH (n) RETURN n',
+        ...     chunk_size=100,
+        ...     on_progress=progress
+        ... ):
+        ...     process_record(record)
+    """
     request = QueryRequest(query=cypher, parameters=parameters)
     result = self.execute_query(
       graph_id, request, QueryOptions(mode="stream", chunk_size=chunk_size)
     )
 
+    count = 0
     if isinstance(result, Iterator):
-      yield from result
+      for item in result:
+        count += 1
+        if on_progress and count % chunk_size == 0:
+          on_progress(count, None)  # Total unknown in streaming
+        yield item
     else:
       # If not streaming, yield all results at once
-      yield from result.data
+      total = len(result.data)
+      for item in result.data:
+        count += 1
+        if on_progress:
+          on_progress(count, total)
+        yield item
+
+  def query_batch(
+    self,
+    graph_id: str,
+    queries: List[str],
+    parameters_list: Optional[List[Dict[str, Any]]] = None,
+    parallel: bool = False,
+  ) -> List[Union[QueryResult, Dict[str, Any]]]:
+    """Execute multiple queries in batch
+
+    Args:
+        graph_id: Graph ID to query
+        queries: List of Cypher query strings
+        parameters_list: List of parameter dicts (one per query)
+        parallel: Execute queries in parallel (experimental)
+
+    Returns:
+        List of QueryResult objects or error dicts
+
+    Example:
+        >>> results = query_client.query_batch('graph_id', [
+        ...     'MATCH (n:Person) RETURN count(n)',
+        ...     'MATCH (c:Company) RETURN count(c)'
+        ... ])
+    """
+    if parameters_list is None:
+      parameters_list = [None] * len(queries)
+
+    if len(queries) != len(parameters_list):
+      raise ValueError("queries and parameters_list must have same length")
+
+    results = []
+    for query, params in zip(queries, parameters_list):
+      try:
+        result = self.query(graph_id, query, params)
+        results.append(result)
+      except Exception as e:
+        # Store error as result
+        results.append({"error": str(e), "query": query})
+
+    return results
 
   def close(self):
     """Cancel any active SSE connections"""

robosystems_client/extensions/tests/test_dataframe_utils.py

@@ -0,0 +1,334 @@
+"""Tests for DataFrame utilities"""
+
+import pytest
+from unittest.mock import Mock, patch
+import tempfile
+import os
+
+# Make pandas optional for tests
+try:
+  import pandas as pd
+
+  HAS_PANDAS = True
+except ImportError:
+  HAS_PANDAS = False
+  pd = None
+
+# Only run tests if pandas is available
+pytestmark = pytest.mark.skipif(not HAS_PANDAS, reason="pandas not installed")
+
+if HAS_PANDAS:
+  from robosystems_client.extensions.dataframe_utils import (
+    query_result_to_dataframe,
+    parse_datetime_columns,
+    stream_to_dataframe,
+    dataframe_to_cypher_params,
+    export_query_to_csv,
+    compare_dataframes,
+    DataFrameQueryClient,
+  )
+
+
+class TestQueryResultToDataFrame:
+  """Test converting query results to DataFrames"""
+
+  def test_query_result_to_dataframe_basic(self):
+    """Test basic conversion from query result to DataFrame"""
+    result = {
+      "data": [
+        {"name": "Alice", "age": 30},
+        {"name": "Bob", "age": 25},
+        {"name": "Charlie", "age": 35},
+      ],
+      "columns": ["name", "age"],
+      "row_count": 3,
+    }
+
+    df = query_result_to_dataframe(result)
+
+    assert len(df) == 3
+    assert list(df.columns) == ["name", "age"]
+    assert df.iloc[0]["name"] == "Alice"
+    assert df.iloc[1]["age"] == 25
+
+  def test_query_result_to_dataframe_nested(self):
+    """Test conversion with nested data"""
+    result = {
+      "data": [
+        {"name": "Alice", "company": {"name": "TechCorp", "revenue": 1000000}},
+        {"name": "Bob", "company": {"name": "StartupInc", "revenue": 500000}},
+      ],
+      "columns": ["name", "company"],
+    }
+
+    df = query_result_to_dataframe(result, normalize_nested=True)
+
+    assert "name" in df.columns
+    assert "company.name" in df.columns
+    assert "company.revenue" in df.columns
+    assert df.iloc[0]["company.name"] == "TechCorp"
+
+  def test_query_result_to_dataframe_empty(self):
+    """Test conversion of empty result"""
+    result = {"data": [], "columns": ["name", "age"], "row_count": 0}
+
+    df = query_result_to_dataframe(result)
+
+    assert len(df) == 0
+    assert list(df.columns) == ["name", "age"]
+
+  def test_query_result_to_dataframe_with_dates(self):
+    """Test conversion with date parsing"""
+    result = {
+      "data": [
+        {"name": "Alice", "created_at": "2023-01-15T10:30:00"},
+        {"name": "Bob", "created_at": "2023-02-20T14:45:00"},
+      ],
+      "columns": ["name", "created_at"],
+    }
+
+    df = query_result_to_dataframe(result, parse_dates=True)
+
+    assert pd.api.types.is_datetime64_any_dtype(df["created_at"])
+    assert df.iloc[0]["created_at"].year == 2023
+
+
+class TestParseDateTimeColumns:
+  """Test datetime parsing functionality"""
+
+  def test_parse_datetime_columns_specific(self):
+    """Test parsing specific datetime columns"""
+    df = pd.DataFrame(
+      {
+        "name": ["Alice", "Bob"],
+        "created_at": ["2023-01-15", "2023-02-20"],
+        "updated_at": ["2023-01-16T10:30:00", "2023-02-21T14:45:00"],
+        "count": [1, 2],
+      }
+    )
+
+    df = parse_datetime_columns(df, date_columns=["created_at", "updated_at"])
+
+    assert pd.api.types.is_datetime64_any_dtype(df["created_at"])
+    assert pd.api.types.is_datetime64_any_dtype(df["updated_at"])
+    assert not pd.api.types.is_datetime64_any_dtype(df["count"])
+
+  def test_parse_datetime_columns_infer(self):
+    """Test automatic datetime column inference"""
+    df = pd.DataFrame(
+      {
+        "name": ["Alice", "Bob"],
+        "timestamp": ["2023-01-15T10:30:00", "2023-02-20T14:45:00"],
+        "not_a_date": ["abc", "def"],
+      }
+    )
+
+    df = parse_datetime_columns(df, infer=True)
+
+    assert pd.api.types.is_datetime64_any_dtype(df["timestamp"])
+    assert df["not_a_date"].dtype == "object"
+
+
+class TestStreamToDataFrame:
+  """Test streaming results to DataFrame"""
+
+  def test_stream_to_dataframe_basic(self):
+    """Test converting stream to DataFrame"""
+
+    def mock_stream():
+      for i in range(10):
+        yield {"id": i, "value": i * 2}
+
+    df = stream_to_dataframe(mock_stream(), chunk_size=3)
+
+    assert len(df) == 10
+    assert df.iloc[5]["value"] == 10
+
+  def test_stream_to_dataframe_with_callback(self):
+    """Test stream with chunk callback"""
+    chunk_counts = []
+
+    def on_chunk(chunk_df, total):
+      chunk_counts.append(len(chunk_df))
+
+    def mock_stream():
+      for i in range(10):
+        yield {"id": i, "value": i * 2}
+
+    df = stream_to_dataframe(mock_stream(), chunk_size=3, on_chunk=on_chunk)
+
+    assert len(df) == 10
+    assert chunk_counts == [3, 3, 3, 1]  # 3 chunks of 3, 1 chunk of 1
+
+
+class TestDataFrameToCypherParams:
+  """Test DataFrame to Cypher parameter conversion"""
+
+  def test_dataframe_to_cypher_params(self):
+    """Test converting DataFrame to Cypher parameters"""
+    df = pd.DataFrame(
+      {
+        "name": ["Alice", "Bob", "Charlie"],
+        "age": [30, 25, 35],
+        "active": [True, False, True],
+      }
+    )
+
+    params = dataframe_to_cypher_params(df)
+
+    assert "data" in params
+    assert len(params["data"]) == 3
+    assert params["data"][0]["name"] == "Alice"
+    assert params["data"][1]["age"] == 25
+
+  def test_dataframe_to_cypher_params_with_nan(self):
+    """Test handling NaN values"""
+    df = pd.DataFrame(
+      {"name": ["Alice", "Bob"], "age": [30, pd.NA], "score": [95.5, None]}
+    )
+
+    params = dataframe_to_cypher_params(df, param_name="records")
+
+    assert "records" in params
+    assert params["records"][1]["age"] is None
+    assert params["records"][1]["score"] is None
+
+
+class TestExportQueryToCSV:
+  """Test CSV export functionality"""
+
+  @patch("robosystems_client.extensions.dataframe_utils.logger")
+  def test_export_query_to_csv(self, mock_logger):
+    """Test exporting query results to CSV"""
+    mock_client = Mock()
+
+    def mock_stream(*args, **kwargs):
+      for i in range(5):
+        yield {"id": i, "name": f"Item {i}"}
+
+    mock_client.stream_query = Mock(side_effect=mock_stream)
+
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".csv") as f:
+      temp_file = f.name
+
+    try:
+      count = export_query_to_csv(
+        mock_client, "graph_id", "MATCH (n) RETURN n", temp_file, chunk_size=2
+      )
+
+      assert count == 5
+      mock_logger.info.assert_called()
+
+      # Verify CSV content
+      df = pd.read_csv(temp_file)
+      assert len(df) == 5
+      assert df.iloc[0]["name"] == "Item 0"
+
+    finally:
+      if os.path.exists(temp_file):
+        os.unlink(temp_file)
+
+
+class TestCompareDataFrames:
+  """Test DataFrame comparison"""
+
+  def test_compare_dataframes_with_keys(self):
+    """Test comparing DataFrames with key columns"""
+    df1 = pd.DataFrame(
+      {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"], "age": [30, 25, 35]}
+    )
+
+    df2 = pd.DataFrame(
+      {"id": [1, 2, 4], "name": ["Alice", "Robert", "David"], "age": [31, 25, 40]}
+    )
+
+    diff = compare_dataframes(df1, df2, key_columns=["id"])
+
+    assert "_merge" in diff.columns
+    assert "name_old" in diff.columns
+    assert "name_new" in diff.columns
+
+  def test_compare_dataframes_without_keys(self):
+    """Test comparing DataFrames without keys"""
+    df1 = pd.DataFrame({"name": ["Alice", "Bob"], "age": [30, 25]})
+
+    df2 = pd.DataFrame({"name": ["Alice", "Charlie"], "age": [30, 35]})
+
+    diff = compare_dataframes(df1, df2)
+
+    assert len(diff) == 2  # Bob and Charlie rows
+
+
+class TestDataFrameQueryClient:
+  """Test DataFrameQueryClient class"""
+
+  def test_query_df(self):
+    """Test query_df method"""
+    mock_client = Mock()
+    mock_client.query.return_value = {
+      "data": [{"name": "Alice"}, {"name": "Bob"}],
+      "columns": ["name"],
+      "row_count": 2,
+    }
+
+    df_client = DataFrameQueryClient(mock_client)
+    df = df_client.query_df("graph_id", "MATCH (n) RETURN n")
+
+    assert len(df) == 2
+    assert df.iloc[0]["name"] == "Alice"
+    mock_client.query.assert_called_once()
+
+  def test_stream_df(self):
+    """Test stream_df method"""
+    mock_client = Mock()
+
+    def mock_stream(*args, **kwargs):
+      for i in range(3):
+        yield {"id": i, "value": i * 10}
+
+    mock_client.stream_query.return_value = mock_stream()
+
+    df_client = DataFrameQueryClient(mock_client)
+    df = df_client.stream_df("graph_id", "MATCH (n) RETURN n")
+
+    assert len(df) == 3
+    assert df.iloc[1]["value"] == 10
+
+  def test_query_batch_df(self):
+    """Test query_batch_df method"""
+    mock_client = Mock()
+    mock_client.query_batch.return_value = [
+      {"data": [{"count": 10}], "columns": ["count"]},
+      {"data": [{"count": 20}], "columns": ["count"]},
+      {"error": "Query failed", "query": "INVALID"},
+    ]
+
+    df_client = DataFrameQueryClient(mock_client)
+    dfs = df_client.query_batch_df(
+      "graph_id",
+      [
+        "MATCH (p:Person) RETURN count(p)",
+        "MATCH (c:Company) RETURN count(c)",
+        "INVALID",
+      ],
+    )
+
+    assert len(dfs) == 3
+    assert dfs[0].iloc[0]["count"] == 10
+    assert dfs[1].iloc[0]["count"] == 20
+    assert "error" in dfs[2].columns
+
+  def test_export_to_csv(self):
+    """Test export_to_csv method"""
+    mock_client = Mock()
+
+    with patch(
+      "robosystems_client.extensions.dataframe_utils.export_query_to_csv"
+    ) as mock_export:
+      mock_export.return_value = 100
+
+      df_client = DataFrameQueryClient(mock_client)
+      count = df_client.export_to_csv("graph_id", "MATCH (n) RETURN n", "output.csv")
+
+      assert count == 100
+      mock_export.assert_called_once()

robosystems_client/extensions/tests/test_integration.py

@@ -71,7 +71,7 @@ class TestAuthenticatedIntegration:
     )
     assert ext.config["headers"]["Authorization"] == "Bearer jwt_token_here"
 
-  @patch("robosystems_client.extensions.auth_integration.sync_detailed")
+  @patch("robosystems_client.api.query.execute_cypher_query.sync_detailed")
   def test_cypher_query_execution(self, mock_sync_detailed, extensions):
     """Test executing Cypher queries through authenticated client"""
     # Mock the response