robosystems-client 0.1.9__py3-none-any.whl

robosystems_client/extensions/auth_integration.py

@@ -0,0 +1,210 @@
+"""Authentication Integration for RoboSystems SDK Extensions
+
+Provides proper integration with the generated SDK authentication system.
+"""
+
+from typing import Dict, Any
+from ..client import Client, AuthenticatedClient
+from .extensions import RoboSystemsExtensions, RoboSystemsExtensionConfig
+
+
+class AuthenticatedExtensions(RoboSystemsExtensions):
+  """Extensions with proper authentication integration"""
+
+  def __init__(
+    self, api_key: str, config: RoboSystemsExtensionConfig = None, base_url: str = None
+  ):
+    """Initialize extensions with API key authentication
+
+    Args:
+        api_key: RoboSystems API key
+        config: Extension configuration
+        base_url: API base URL (defaults to production)
+    """
+    if config is None:
+      config = RoboSystemsExtensionConfig()
+
+    # Set base URL
+    if base_url:
+      config.base_url = base_url
+    elif not config.base_url:
+      config.base_url = "https://api.robosystems.ai"
+
+    # Add authentication headers
+    if not config.headers:
+      config.headers = {}
+    config.headers["X-API-Key"] = api_key
+    config.headers["Authorization"] = f"Bearer {api_key}"
+
+    super().__init__(config)
+
+    # Store authenticated client for SDK operations
+    self._authenticated_client = AuthenticatedClient(
+      base_url=config.base_url, token=api_key, headers=config.headers
+    )
+
+  @property
+  def authenticated_client(self) -> AuthenticatedClient:
+    """Get the authenticated client for direct SDK operations"""
+    return self._authenticated_client
+
+  def execute_cypher_query(
+    self, graph_id: str, query: str, parameters: Dict[str, Any] = None
+  ):
+    """Execute Cypher query using authenticated SDK client"""
+    from ..api.query.execute_cypher_query import sync_detailed
+    from ..models.cypher_query_request import CypherQueryRequest
+
+    request = CypherQueryRequest(query=query, parameters=parameters or {})
+
+    response = sync_detailed(
+      graph_id=graph_id, client=self._authenticated_client, body=request
+    )
+
+    if response.parsed:
+      return {
+        "data": getattr(response.parsed, "data", []),
+        "columns": getattr(response.parsed, "columns", []),
+        "row_count": getattr(response.parsed, "row_count", 0),
+        "execution_time_ms": getattr(response.parsed, "execution_time_ms", 0),
+      }
+    else:
+      raise Exception(f"Query failed: {response.status_code}")
+
+
+class CookieAuthExtensions(RoboSystemsExtensions):
+  """Extensions with cookie-based authentication"""
+
+  def __init__(
+    self,
+    cookies: Dict[str, str],
+    config: RoboSystemsExtensionConfig = None,
+    base_url: str = None,
+  ):
+    """Initialize extensions with cookie authentication
+
+    Args:
+        cookies: Authentication cookies (e.g., {'auth-token': 'token_value'})
+        config: Extension configuration
+        base_url: API base URL
+    """
+    if config is None:
+      config = RoboSystemsExtensionConfig()
+
+    if base_url:
+      config.base_url = base_url
+    elif not config.base_url:
+      config.base_url = "https://api.robosystems.ai"
+
+    super().__init__(config)
+
+    # Store cookies for requests
+    self._cookies = cookies
+
+    # Create client with cookies
+    self._client = Client(
+      base_url=config.base_url, cookies=cookies, headers=config.headers or {}
+    )
+
+  @property
+  def client(self) -> Client:
+    """Get the client for cookie-based operations"""
+    return self._client
+
+
+class TokenExtensions(RoboSystemsExtensions):
+  """Extensions with JWT/Bearer token authentication"""
+
+  def __init__(
+    self, token: str, config: RoboSystemsExtensionConfig = None, base_url: str = None
+  ):
+    """Initialize extensions with JWT token
+
+    Args:
+        token: JWT or Bearer token
+        config: Extension configuration
+        base_url: API base URL
+    """
+    if config is None:
+      config = RoboSystemsExtensionConfig()
+
+    if base_url:
+      config.base_url = base_url
+    elif not config.base_url:
+      config.base_url = "https://api.robosystems.ai"
+
+    # Add authentication headers
+    if not config.headers:
+      config.headers = {}
+    config.headers["Authorization"] = f"Bearer {token}"
+
+    super().__init__(config)
+
+    # Store authenticated client
+    self._authenticated_client = AuthenticatedClient(
+      base_url=config.base_url, token=token, headers=config.headers
+    )
+
+  @property
+  def authenticated_client(self) -> AuthenticatedClient:
+    """Get the authenticated client"""
+    return self._authenticated_client
+
+
+def create_extensions(auth_method: str, **kwargs) -> RoboSystemsExtensions:
+  """Factory function to create extensions with proper authentication
+
+  Args:
+      auth_method: 'api_key', 'cookie', or 'token'
+      **kwargs: Authentication parameters
+
+  Returns:
+      Configured extensions instance
+
+  Examples:
+      # API Key authentication
+      ext = create_extensions('api_key', api_key='your-key', base_url='https://api.robosystems.ai')
+
+      # Cookie authentication
+      ext = create_extensions('cookie', cookies={'auth-token': 'token'})
+
+      # JWT Token authentication
+      ext = create_extensions('token', token='jwt-token')
+  """
+  if auth_method == "api_key":
+    api_key = kwargs.pop("api_key")
+    return AuthenticatedExtensions(api_key, **kwargs)
+
+  elif auth_method == "cookie":
+    cookies = kwargs.pop("cookies")
+    return CookieAuthExtensions(cookies, **kwargs)
+
+  elif auth_method == "token":
+    token = kwargs.pop("token")
+    return TokenExtensions(token, **kwargs)
+
+  else:
+    raise ValueError(
+      f"Unknown auth method: {auth_method}. Use 'api_key', 'cookie', or 'token'"
+    )
+
+
+# Example usage functions
+def create_production_extensions(api_key: str) -> AuthenticatedExtensions:
+  """Create extensions for production environment"""
+  return AuthenticatedExtensions(
+    api_key=api_key,
+    config=RoboSystemsExtensionConfig(
+      base_url="https://api.robosystems.ai", max_retries=3, retry_delay=2000, timeout=60
+    ),
+  )
+
+
+def create_development_extensions(api_key: str) -> AuthenticatedExtensions:
+  """Create extensions for development environment"""
+  return AuthenticatedExtensions(
+    api_key=api_key,
+    config=RoboSystemsExtensionConfig(
+      base_url="http://localhost:8000", max_retries=5, retry_delay=1000, timeout=30
+    ),
+  )

robosystems_client/extensions/extensions.py

@@ -0,0 +1,170 @@
+"""RoboSystems SDK Extensions - Main entry point
+
+Enhanced clients with SSE support for the RoboSystems API.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, Any, Optional, Callable
+
+from .query_client import QueryClient
+from .operation_client import OperationClient
+from .sse_client import SSEClient
+
+
+@dataclass
+class RoboSystemsExtensionConfig:
+  """Configuration for RoboSystems extensions"""
+
+  base_url: Optional[str] = None
+  headers: Optional[Dict[str, str]] = None
+  max_retries: int = 5
+  retry_delay: int = 1000
+  timeout: int = 30
+
+
+class RoboSystemsExtensions:
+  """Main extensions class providing enhanced RoboSystems API functionality"""
+
+  def __init__(self, config: RoboSystemsExtensionConfig = None):
+    if config is None:
+      config = RoboSystemsExtensionConfig()
+
+    # Get base URL from config or use default
+    self.config = {
+      "base_url": config.base_url or "http://localhost:8000",
+      "headers": config.headers or {},
+      "max_retries": config.max_retries,
+      "retry_delay": config.retry_delay,
+      "timeout": config.timeout,
+    }
+
+    # Initialize clients
+    self.query = QueryClient(self.config)
+    self.operations = OperationClient(self.config)
+
+  def monitor_operation(
+    self, operation_id: str, on_progress: Optional[Callable] = None
+  ) -> Any:
+    """Convenience method to monitor any operation"""
+    from .operation_client import MonitorOptions
+
+    options = MonitorOptions(on_progress=on_progress)
+    return self.operations.monitor_operation(operation_id, options)
+
+  def create_sse_client(self) -> SSEClient:
+    """Create custom SSE client for advanced use cases"""
+    from .sse_client import SSEConfig
+
+    sse_config = SSEConfig(
+      base_url=self.config["base_url"],
+      headers=self.config["headers"],
+      max_retries=self.config["max_retries"],
+      retry_delay=self.config["retry_delay"],
+      timeout=self.config["timeout"],
+    )
+
+    return SSEClient(sse_config)
+
+  def close(self):
+    """Clean up all active connections"""
+    self.query.close()
+    self.operations.close_all()
+
+  # Convenience methods that delegate to the appropriate clients
+  def execute_query(self, graph_id: str, query: str, parameters: Dict[str, Any] = None):
+    """Execute a query using the query client"""
+    return self.query.query(graph_id, query, parameters)
+
+  def stream_query(
+    self,
+    graph_id: str,
+    query: str,
+    parameters: Dict[str, Any] = None,
+    chunk_size: int = 1000,
+  ):
+    """Stream a query using the query client"""
+    return self.query.stream_query(graph_id, query, parameters, chunk_size)
+
+  def get_operation_status(self, operation_id: str):
+    """Get operation status using the operation client"""
+    return self.operations.get_operation_status(operation_id)
+
+  def cancel_operation(self, operation_id: str):
+    """Cancel an operation using the operation client"""
+    return self.operations.cancel_operation(operation_id)
+
+
+class AsyncRoboSystemsExtensions:
+  """Async version of the extensions class"""
+
+  def __init__(self, config: RoboSystemsExtensionConfig = None):
+    if config is None:
+      config = RoboSystemsExtensionConfig()
+
+    self.config = {
+      "base_url": config.base_url or "http://localhost:8000",
+      "headers": config.headers or {},
+      "max_retries": config.max_retries,
+      "retry_delay": config.retry_delay,
+      "timeout": config.timeout,
+    }
+
+    # Initialize async clients
+    from .query_client import AsyncQueryClient
+    from .operation_client import AsyncOperationClient
+
+    self.query = AsyncQueryClient(self.config)
+    self.operations = AsyncOperationClient(self.config)
+
+  async def monitor_operation(
+    self, operation_id: str, on_progress: Optional[Callable] = None
+  ) -> Any:
+    """Convenience method to monitor any operation (async)"""
+    from .operation_client import MonitorOptions
+
+    options = MonitorOptions(on_progress=on_progress)
+    return await self.operations.monitor_operation(operation_id, options)
+
+  def create_sse_client(self):
+    """Create custom async SSE client for advanced use cases"""
+    from .sse_client import AsyncSSEClient, SSEConfig
+
+    sse_config = SSEConfig(
+      base_url=self.config["base_url"],
+      headers=self.config["headers"],
+      max_retries=self.config["max_retries"],
+      retry_delay=self.config["retry_delay"],
+      timeout=self.config["timeout"],
+    )
+
+    return AsyncSSEClient(sse_config)
+
+  async def close(self):
+    """Clean up all active connections (async)"""
+    await self.query.close()
+    await self.operations.close_all()
+
+  async def execute_query(
+    self, graph_id: str, query: str, parameters: Dict[str, Any] = None
+  ):
+    """Execute a query using the async query client"""
+    return await self.query.query(graph_id, query, parameters)
+
+  async def stream_query(
+    self,
+    graph_id: str,
+    query: str,
+    parameters: Dict[str, Any] = None,
+    chunk_size: int = 1000,
+  ):
+    """Stream a query using the async query client"""
+    async for item in self.query.stream_query(graph_id, query, parameters, chunk_size):
+      yield item
+
+  async def get_operation_status(self, operation_id: str):
+    """Get operation status using the async operation client"""
+    return await self.operations.get_operation_status(operation_id)
+
+  async def cancel_operation(self, operation_id: str):
+    """Cancel an operation using the async operation client"""
+    return await self.operations.cancel_operation(operation_id)