robosystems-client 0.2.7__py3-none-any.whl → 0.2.8__py3-none-any.whl

robosystems_client/extensions/query_client.py

@@ -141,39 +141,100 @@ class QueryClient:
       if hasattr(response, "parsed") and response.parsed:
         response_data = response.parsed
 
-        # Handle both dict and object responses
+        # Handle both dict and attrs object responses
         if isinstance(response_data, dict):
           # Response is a plain dict
           data = response_data
+          has_data = "data" in data and "columns" in data
         else:
-          # Response is an object with additional_properties
-          data = (
-            response_data.additional_properties
-            if hasattr(response_data, "additional_properties")
-            else response_data
-          )
+          # Response is an attrs object - check for attributes directly
+          data = response_data
+          has_data = hasattr(data, "data") and hasattr(data, "columns")
 
         # Check if this is an immediate response
-        if "data" in data and "columns" in data:
+        if has_data:
+          # Extract data from either dict or attrs object
+          if isinstance(data, dict):
+            result_data = data["data"]
+            result_columns = data["columns"]
+            result_row_count = data.get("row_count", len(data["data"]))
+            result_execution_time = data.get("execution_time_ms", 0)
+            result_timestamp = data.get("timestamp", datetime.now().isoformat())
+          else:
+            # attrs object - access attributes directly
+            from ..types import UNSET
+
+            raw_data = data.data if data.data is not UNSET else []
+            # Convert data items to dicts if they're objects
+            result_data = []
+            for item in raw_data:
+              if hasattr(item, "to_dict"):
+                result_data.append(item.to_dict())
+              elif hasattr(item, "additional_properties"):
+                result_data.append(item.additional_properties)
+              else:
+                result_data.append(item)
+            result_columns = data.columns if data.columns is not UNSET else []
+            result_row_count = (
+              data.row_count if data.row_count is not UNSET else len(result_data)
+            )
+            result_execution_time = (
+              data.execution_time_ms if data.execution_time_ms is not UNSET else 0
+            )
+            result_timestamp = (
+              data.timestamp
+              if data.timestamp is not UNSET
+              else datetime.now().isoformat()
+            )
+
           return QueryResult(
-            data=data["data"],
-            columns=data["columns"],
-            row_count=data.get("row_count", len(data["data"])),
-            execution_time_ms=data.get("execution_time_ms", 0),
+            data=result_data,
+            columns=result_columns,
+            row_count=result_row_count,
+            execution_time_ms=result_execution_time,
             graph_id=graph_id,
-            timestamp=data.get("timestamp", datetime.now().isoformat()),
+            timestamp=result_timestamp,
           )
 
         # Check if this is a queued response
-        if data.get("status") == "queued" and "operation_id" in data:
-          queued_response = QueuedQueryResponse(
-            status=data["status"],
-            operation_id=data["operation_id"],
-            queue_position=data.get("queue_position", 0),
-            estimated_wait_seconds=data.get("estimated_wait_seconds", 0),
-            message=data.get("message", "Query queued"),
+        is_queued = False
+        queued_response = None
+
+        if isinstance(data, dict):
+          is_queued = data.get("status") == "queued" and "operation_id" in data
+          if is_queued:
+            queued_response = QueuedQueryResponse(
+              status=data["status"],
+              operation_id=data["operation_id"],
+              queue_position=data.get("queue_position", 0),
+              estimated_wait_seconds=data.get("estimated_wait_seconds", 0),
+              message=data.get("message", "Query queued"),
+            )
+        else:
+          is_queued = (
+            hasattr(data, "status")
+            and hasattr(data, "operation_id")
+            and getattr(data, "status", None) == "queued"
           )
+          if is_queued:
+            from ..types import UNSET
+
+            queued_response = QueuedQueryResponse(
+              status=data.status,
+              operation_id=data.operation_id,
+              queue_position=data.queue_position
+              if hasattr(data, "queue_position") and data.queue_position is not UNSET
+              else 0,
+              estimated_wait_seconds=data.estimated_wait_seconds
+              if hasattr(data, "estimated_wait_seconds")
+              and data.estimated_wait_seconds is not UNSET
+              else 0,
+              message=data.message
+              if hasattr(data, "message") and data.message is not UNSET
+              else "Query queued",
+            )
 
+        if is_queued and queued_response:
           # Notify about queue status
           if options.on_queue_update:
             options.on_queue_update(

{robosystems_client-0.2.7.dist-info → robosystems_client-0.2.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robosystems-client
-Version: 0.2.7
+Version: 0.2.8
 Summary: Python Client for RoboSystems financial graph database API
 Author: RFS LLC
 License: MIT

{robosystems_client-0.2.7.dist-info → robosystems_client-0.2.8.dist-info}/RECORD

@@ -135,7 +135,7 @@ robosystems_client/extensions/dataframe_utils.py,sha256=gK1bgkVqBF0TvWVdGQvqWrt-
 robosystems_client/extensions/extensions.py,sha256=ROnCobUek4Dke9dVx2sTzNKhz309NOG40EDSYHtNmWs,6257
 robosystems_client/extensions/graph_client.py,sha256=OBi0xj0SLIRKLeSu_DiGt2ZakCmhggvNrMP3jdRfEgQ,10326
 robosystems_client/extensions/operation_client.py,sha256=B1qju-wWQrnrnVJixKGgsA_KEInviwJwdlJxzm_i7P0,13359
-robosystems_client/extensions/query_client.py,sha256=-4g1tE3QMM6WggV01W8U9jTJV6hL-00V1T_tvcKb3d0,16949
+robosystems_client/extensions/query_client.py,sha256=cX3e8EBoTeg4Lwm6edJYRULM2UmGpfqNX3f48S8TQbE,19430
 robosystems_client/extensions/sse_client.py,sha256=XvQIq3JQ0Yiax11E7cwclhupShYOpEMURM2cYQodiz8,15058
 robosystems_client/extensions/table_ingest_client.py,sha256=H-gJBsYshEfiJ3KTIE3yjspUjQWLIYFKuttrmfRjFdk,12873
 robosystems_client/extensions/token_utils.py,sha256=qCK_s1vBzRnSYwtgncPZRLJVIw3WXmzqNTWjdEEpdgs,10899
@@ -375,7 +375,7 @@ robosystems_client/models/user_usage_response_graphs.py,sha256=xAH-ZnhaUfWQ_2EpZ
 robosystems_client/models/user_usage_summary_response.py,sha256=4hthwTH7bXyzdYlHoekDYOgDLI-stGRH507Bl2rUjYA,3655
 robosystems_client/models/user_usage_summary_response_usage_vs_limits.py,sha256=XrZnRcy1nD3xtKX4svbww7QfEHrN7_XIfeL9j5ZMbyQ,1298
 robosystems_client/models/validation_error.py,sha256=R77OuQG2nJ3WDFfY--xbEhg6x1D7gAAp_1UdnG8Ka2A,1949
-robosystems_client-0.2.7.dist-info/METADATA,sha256=d8108sRLbIc6bJBbGx_pDcnj21HAejrSzGFFowIYasE,3903
-robosystems_client-0.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-robosystems_client-0.2.7.dist-info/licenses/LICENSE,sha256=LjFqQPU4eQh7jAQ04SmE9eC0j74HCdXvzbo0hjW4mWo,1063
-robosystems_client-0.2.7.dist-info/RECORD,,
+robosystems_client-0.2.8.dist-info/METADATA,sha256=AShLe4Foo6kgRdRz6Hsu0VcRWgwyPr7ZVu3XTKWswxI,3903
+robosystems_client-0.2.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+robosystems_client-0.2.8.dist-info/licenses/LICENSE,sha256=LjFqQPU4eQh7jAQ04SmE9eC0j74HCdXvzbo0hjW4mWo,1063
+robosystems_client-0.2.8.dist-info/RECORD,,