robosystems-client 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

robosystems_client/extensions/copy_client.py

@@ -0,0 +1,469 @@
+"""Enhanced Copy Client with SSE support
+
+Provides intelligent data copy operations with progress monitoring.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, Any, Optional, Callable, Union, List
+from enum import Enum
+import time
+import logging
+
+from ..api.copy.copy_data_to_graph import sync_detailed as copy_data_to_graph
+from ..models.s3_copy_request import S3CopyRequest
+from ..models.url_copy_request import URLCopyRequest
+from ..models.data_frame_copy_request import DataFrameCopyRequest
+from ..models.copy_response import CopyResponse
+from ..models.copy_response_status import CopyResponseStatus
+from ..models.s3_copy_request_file_format import S3CopyRequestFileFormat
+from .sse_client import SSEClient, AsyncSSEClient, SSEConfig, EventType
+
+logger = logging.getLogger(__name__)
+
+
+class CopySourceType(Enum):
+  """Types of copy sources"""
+
+  S3 = "s3"
+  URL = "url"
+  DATAFRAME = "dataframe"
+
+
+@dataclass
+class CopyOptions:
+  """Options for copy operations"""
+
+  on_progress: Optional[Callable[[str, Optional[float]], None]] = None
+  on_queue_update: Optional[Callable[[int, int], None]] = None
+  on_warning: Optional[Callable[[str], None]] = None
+  timeout: Optional[int] = None
+  test_mode: Optional[bool] = None
+
+
+@dataclass
+class CopyResult:
+  """Result from copy operation"""
+
+  status: str  # 'completed', 'failed', 'partial', 'accepted'
+  rows_imported: Optional[int] = None
+  rows_skipped: Optional[int] = None
+  bytes_processed: Optional[int] = None
+  execution_time_ms: Optional[float] = None
+  warnings: Optional[List[str]] = None
+  error: Optional[str] = None
+  operation_id: Optional[str] = None
+  sse_url: Optional[str] = None
+  message: Optional[str] = None
+
+
+@dataclass
+class CopyStatistics:
+  """Statistics from copy operation"""
+
+  total_rows: int
+  imported_rows: int
+  skipped_rows: int
+  bytes_processed: int
+  duration: float  # seconds
+  throughput: float  # rows per second
+
+
+class CopyClient:
+  """Enhanced copy client with SSE streaming support"""
+
+  def __init__(self, config: Dict[str, Any]):
+    self.config = config
+    self.base_url = config["base_url"]
+    self.sse_client: Optional[SSEClient] = None
+
+    # Get client authentication if provided
+    self.auth_token = config.get("auth_token")
+    self.api_key = config.get("api_key")
+
+  def copy_from_s3(
+    self, graph_id: str, request: S3CopyRequest, options: Optional[CopyOptions] = None
+  ) -> CopyResult:
+    """Copy data from S3 to graph database"""
+    return self._execute_copy(graph_id, request, CopySourceType.S3, options)
+
+  def copy_from_url(
+    self, graph_id: str, request: URLCopyRequest, options: Optional[CopyOptions] = None
+  ) -> CopyResult:
+    """Copy data from URL to graph database (when available)"""
+    return self._execute_copy(graph_id, request, CopySourceType.URL, options)
+
+  def copy_from_dataframe(
+    self,
+    graph_id: str,
+    request: DataFrameCopyRequest,
+    options: Optional[CopyOptions] = None,
+  ) -> CopyResult:
+    """Copy data from DataFrame to graph database (when available)"""
+    return self._execute_copy(graph_id, request, CopySourceType.DATAFRAME, options)
+
+  def _execute_copy(
+    self,
+    graph_id: str,
+    request: Union[S3CopyRequest, URLCopyRequest, DataFrameCopyRequest],
+    source_type: CopySourceType,
+    options: Optional[CopyOptions] = None,
+  ) -> CopyResult:
+    """Execute copy operation with automatic SSE monitoring for long-running operations"""
+    if options is None:
+      options = CopyOptions()
+
+    start_time = time.time()
+
+    # Import client here to avoid circular imports
+    from ..client import AuthenticatedClient
+
+    # Create authenticated client
+    client = AuthenticatedClient(
+      base_url=self.base_url,
+      token=self.auth_token,
+      headers={"X-API-Key": self.api_key} if self.api_key else None,
+    )
+
+    try:
+      # Execute the copy request
+      response = copy_data_to_graph(graph_id=graph_id, client=client, body=request)
+
+      if response.parsed:
+        response_data: CopyResponse = response.parsed
+
+        # Check if this is an accepted (async) operation
+        if (
+          response_data.status == CopyResponseStatus.ACCEPTED
+          and response_data.operation_id
+        ):
+          # This is a long-running operation with SSE monitoring
+          if options.on_progress:
+            options.on_progress("Copy operation started. Monitoring progress...", None)
+
+          # If SSE URL is provided, use it for monitoring
+          if response_data.sse_url:
+            return self._monitor_copy_operation(
+              response_data.operation_id, options, start_time
+            )
+
+          # Otherwise return the accepted response
+          return CopyResult(
+            status="accepted",
+            operation_id=response_data.operation_id,
+            sse_url=response_data.sse_url,
+            message=response_data.message,
+          )
+
+        # This is a synchronous response - operation completed immediately
+        return self._build_copy_result(response_data, time.time() - start_time)
+      else:
+        return CopyResult(
+          status="failed",
+          error="No response data received",
+          execution_time_ms=(time.time() - start_time) * 1000,
+        )
+
+    except Exception as e:
+      return CopyResult(
+        status="failed",
+        error=str(e),
+        execution_time_ms=(time.time() - start_time) * 1000,
+      )
+
+  def _monitor_copy_operation(
+    self, operation_id: str, options: CopyOptions, start_time: float
+  ) -> CopyResult:
+    """Monitor a copy operation using SSE"""
+    timeout_ms = options.timeout or 3600000  # Default 1 hour for copy operations
+    timeout_time = time.time() + (timeout_ms / 1000)
+
+    result = CopyResult(status="failed")
+    warnings: List[str] = []
+
+    # Set up SSE connection
+    sse_config = SSEConfig(base_url=self.base_url, timeout=timeout_ms // 1000)
+    sse_client = SSEClient(sse_config)
+
+    try:
+      sse_client.connect(operation_id)
+
+      # Set up event handlers
+      def on_queue_update(data):
+        if options.on_queue_update:
+          position = data.get("position", data.get("queue_position", 0))
+          wait_time = data.get("estimated_wait_seconds", 0)
+          options.on_queue_update(position, wait_time)
+
+      def on_progress(data):
+        if options.on_progress:
+          message = data.get("message", data.get("status", "Processing..."))
+          progress_percent = data.get("progress_percent", data.get("progress"))
+          options.on_progress(message, progress_percent)
+
+        # Check for warnings in progress updates
+        if "warnings" in data and data["warnings"]:
+          warnings.extend(data["warnings"])
+          if options.on_warning:
+            for warning in data["warnings"]:
+              options.on_warning(warning)
+
+      def on_completed(data):
+        nonlocal result
+        completion_data = data.get("result", data)
+        result = CopyResult(
+          status=completion_data.get("status", "completed"),
+          rows_imported=completion_data.get("rows_imported"),
+          rows_skipped=completion_data.get("rows_skipped"),
+          bytes_processed=completion_data.get("bytes_processed"),
+          execution_time_ms=(time.time() - start_time) * 1000,
+          warnings=warnings if warnings else completion_data.get("warnings"),
+          message=completion_data.get("message"),
+        )
+
+      def on_error(data):
+        nonlocal result
+        result = CopyResult(
+          status="failed",
+          error=data.get("message", data.get("error", "Copy operation failed")),
+          execution_time_ms=(time.time() - start_time) * 1000,
+          warnings=warnings if warnings else None,
+        )
+
+      def on_cancelled(data):
+        nonlocal result
+        result = CopyResult(
+          status="failed",
+          error="Copy operation cancelled",
+          execution_time_ms=(time.time() - start_time) * 1000,
+          warnings=warnings if warnings else None,
+        )
+
+      # Register event handlers
+      sse_client.on(EventType.QUEUE_UPDATE.value, on_queue_update)
+      sse_client.on(EventType.OPERATION_PROGRESS.value, on_progress)
+      sse_client.on(EventType.OPERATION_COMPLETED.value, on_completed)
+      sse_client.on(EventType.OPERATION_ERROR.value, on_error)
+      sse_client.on(EventType.OPERATION_CANCELLED.value, on_cancelled)
+
+      # Listen for events until completion or timeout
+      while time.time() < timeout_time:
+        sse_client.listen(timeout=1)  # Process events for 1 second
+
+        # Check if operation is complete
+        if result.status in ["completed", "failed", "partial"]:
+          break
+
+      if time.time() >= timeout_time:
+        result = CopyResult(
+          status="failed",
+          error=f"Copy operation timeout after {timeout_ms}ms",
+          execution_time_ms=(time.time() - start_time) * 1000,
+        )
+
+    finally:
+      sse_client.close()
+
+    return result
+
+  def _build_copy_result(
+    self, response_data: CopyResponse, execution_time: float
+  ) -> CopyResult:
+    """Build copy result from response data"""
+    return CopyResult(
+      status=response_data.status.value,
+      rows_imported=response_data.rows_imported,
+      rows_skipped=response_data.rows_skipped,
+      bytes_processed=response_data.bytes_processed,
+      execution_time_ms=response_data.execution_time_ms or (execution_time * 1000),
+      warnings=response_data.warnings,
+      message=response_data.message,
+      error=str(response_data.error_details) if response_data.error_details else None,
+    )
+
+  def calculate_statistics(self, result: CopyResult) -> Optional[CopyStatistics]:
+    """Calculate copy statistics from result"""
+    if result.status == "failed" or not result.rows_imported:
+      return None
+
+    total_rows = (result.rows_imported or 0) + (result.rows_skipped or 0)
+    duration = (result.execution_time_ms or 0) / 1000  # Convert to seconds
+    throughput = (result.rows_imported or 0) / duration if duration > 0 else 0
+
+    return CopyStatistics(
+      total_rows=total_rows,
+      imported_rows=result.rows_imported or 0,
+      skipped_rows=result.rows_skipped or 0,
+      bytes_processed=result.bytes_processed or 0,
+      duration=duration,
+      throughput=throughput,
+    )
+
+  def copy_s3(
+    self,
+    graph_id: str,
+    table_name: str,
+    s3_path: str,
+    access_key_id: str,
+    secret_access_key: str,
+    region: str = "us-east-1",
+    file_format: Optional[str] = None,
+    ignore_errors: bool = False,
+  ) -> CopyResult:
+    """Convenience method for simple S3 copy with default options"""
+
+    # Map string format to enum
+    format_enum = S3CopyRequestFileFormat.PARQUET
+    if file_format:
+      format_map = {
+        "csv": S3CopyRequestFileFormat.CSV,
+        "parquet": S3CopyRequestFileFormat.PARQUET,
+        "json": S3CopyRequestFileFormat.JSON,
+        "delta": S3CopyRequestFileFormat.DELTA,
+        "iceberg": S3CopyRequestFileFormat.ICEBERG,
+      }
+      format_enum = format_map.get(file_format.lower(), S3CopyRequestFileFormat.PARQUET)
+
+    request = S3CopyRequest(
+      table_name=table_name,
+      s3_path=s3_path,
+      s3_access_key_id=access_key_id,
+      s3_secret_access_key=secret_access_key,
+      s3_region=region,
+      file_format=format_enum,
+      ignore_errors=ignore_errors,
+    )
+
+    return self.copy_from_s3(graph_id, request)
+
+  def monitor_multiple_copies(
+    self, operation_ids: List[str], options: Optional[CopyOptions] = None
+  ) -> Dict[str, CopyResult]:
+    """Monitor multiple copy operations concurrently"""
+    results = {}
+    for operation_id in operation_ids:
+      result = self._monitor_copy_operation(
+        operation_id, options or CopyOptions(), time.time()
+      )
+      results[operation_id] = result
+    return results
+
+  def batch_copy_from_s3(
+    self, graph_id: str, copies: List[Dict[str, Any]]
+  ) -> List[CopyResult]:
+    """Batch copy multiple tables from S3"""
+    results = []
+    for copy_config in copies:
+      request = copy_config.get("request")
+      options = copy_config.get("options")
+      if request:
+        result = self.copy_from_s3(graph_id, request, options)
+        results.append(result)
+    return results
+
+  def copy_with_retry(
+    self,
+    graph_id: str,
+    request: Union[S3CopyRequest, URLCopyRequest, DataFrameCopyRequest],
+    source_type: CopySourceType,
+    max_retries: int = 3,
+    options: Optional[CopyOptions] = None,
+  ) -> CopyResult:
+    """Copy with retry logic for transient failures"""
+    if options is None:
+      options = CopyOptions()
+
+    last_error: Optional[Exception] = None
+    attempt = 0
+
+    while attempt < max_retries:
+      attempt += 1
+
+      try:
+        result = self._execute_copy(graph_id, request, source_type, options)
+
+        # If successful or partially successful, return
+        if result.status in ["completed", "partial"]:
+          return result
+
+        # If failed, check if it's retryable
+        if result.status == "failed":
+          is_retryable = self._is_retryable_error(result.error)
+          if not is_retryable or attempt == max_retries:
+            return result
+
+          # Wait before retry with exponential backoff
+          wait_time = min(1000 * (2 ** (attempt - 1)), 30000) / 1000
+          if options.on_progress:
+            options.on_progress(
+              f"Retrying copy operation (attempt {attempt}/{max_retries}) in {wait_time}s...",
+              None,
+            )
+          time.sleep(wait_time)
+
+      except Exception as e:
+        last_error = e
+
+        if attempt == max_retries:
+          raise last_error
+
+        # Wait before retry
+        wait_time = min(1000 * (2 ** (attempt - 1)), 30000) / 1000
+        if options.on_progress:
+          options.on_progress(
+            f"Retrying after error (attempt {attempt}/{max_retries}) in {wait_time}s...",
+            None,
+          )
+        time.sleep(wait_time)
+
+    raise last_error or Exception("Copy operation failed after all retries")
+
+  def _is_retryable_error(self, error: Optional[str]) -> bool:
+    """Check if an error is retryable"""
+    if not error:
+      return False
+
+    retryable_patterns = [
+      "timeout",
+      "network",
+      "connection",
+      "temporary",
+      "unavailable",
+      "rate limit",
+      "throttl",
+    ]
+
+    lower_error = error.lower()
+    return any(pattern in lower_error for pattern in retryable_patterns)
+
+  def close(self):
+    """Cancel any active SSE connections"""
+    if self.sse_client:
+      self.sse_client.close()
+      self.sse_client = None
+
+
+class AsyncCopyClient:
+  """Async version of CopyClient for async/await usage"""
+
+  def __init__(self, config: Dict[str, Any]):
+    self.config = config
+    self.base_url = config["base_url"]
+    self.sse_client: Optional[AsyncSSEClient] = None
+
+    # Get client authentication if provided
+    self.auth_token = config.get("auth_token")
+    self.api_key = config.get("api_key")
+
+  async def copy_from_s3(
+    self, graph_id: str, request: S3CopyRequest, options: Optional[CopyOptions] = None
+  ) -> CopyResult:
+    """Copy data from S3 to graph database asynchronously"""
+    # Async implementation would go here
+    # For now, this is a placeholder
+    raise NotImplementedError("Async copy client not yet implemented")
+
+  async def close(self):
+    """Close any active connections"""
+    if self.sse_client:
+      await self.sse_client.close()
+      self.sse_client = None

robosystems_client/extensions/extensions.py

@@ -8,6 +8,7 @@ from typing import Dict, Any, Optional, Callable
 
 from .query_client import QueryClient
 from .operation_client import OperationClient
+from .copy_client import CopyClient
 from .sse_client import SSEClient
 
 
@@ -39,6 +40,7 @@ class RoboSystemsExtensions:
     }
 
     # Initialize clients
+    self.copy = CopyClient(self.config)
     self.query = QueryClient(self.config)
     self.operations = OperationClient(self.config)
 
@@ -67,6 +69,7 @@ class RoboSystemsExtensions:
 
   def close(self):
     """Clean up all active connections"""
+    self.copy.close()
     self.query.close()
     self.operations.close_all()
 
@@ -93,6 +96,20 @@ class RoboSystemsExtensions:
     """Cancel an operation using the operation client"""
     return self.operations.cancel_operation(operation_id)
 
+  def copy_from_s3(
+    self,
+    graph_id: str,
+    table_name: str,
+    s3_path: str,
+    access_key_id: str,
+    secret_access_key: str,
+    **kwargs,
+  ):
+    """Copy data from S3 using the copy client"""
+    return self.copy.copy_s3(
+      graph_id, table_name, s3_path, access_key_id, secret_access_key, **kwargs
+    )
+
 
 class AsyncRoboSystemsExtensions:
   """Async version of the extensions class"""

robosystems_client/models/__init__.py

@@ -14,7 +14,6 @@ from .auth_response_user import AuthResponseUser
 from .available_extension import AvailableExtension
 from .available_extensions_response import AvailableExtensionsResponse
 from .backup_create_request import BackupCreateRequest
-from .backup_export_request import BackupExportRequest
 from .backup_list_response import BackupListResponse
 from .backup_response import BackupResponse
 from .backup_restore_request import BackupRestoreRequest
@@ -34,6 +33,9 @@ from .connection_provider_info_provider import ConnectionProviderInfoProvider
 from .connection_response import ConnectionResponse
 from .connection_response_metadata import ConnectionResponseMetadata
 from .connection_response_provider import ConnectionResponseProvider
+from .copy_response import CopyResponse
+from .copy_response_error_details_type_0 import CopyResponseErrorDetailsType0
+from .copy_response_status import CopyResponseStatus
 from .create_api_key_request import CreateAPIKeyRequest
 from .create_api_key_response import CreateAPIKeyResponse
 from .create_connection_request import CreateConnectionRequest
@@ -41,7 +43,6 @@ from .create_connection_request_provider import CreateConnectionRequestProvider
 from .create_graph_request import CreateGraphRequest
 from .create_subgraph_request import CreateSubgraphRequest
 from .create_subgraph_request_metadata_type_0 import CreateSubgraphRequestMetadataType0
-from .credit_check_request import CreditCheckRequest
 from .credit_summary import CreditSummary
 from .credit_summary_response import CreditSummaryResponse
 from .credits_summary_response import CreditsSummaryResponse
@@ -56,6 +57,8 @@ from .custom_schema_definition_relationships_item import (
 )
 from .cypher_query_request import CypherQueryRequest
 from .cypher_query_request_parameters_type_0 import CypherQueryRequestParametersType0
+from .data_frame_copy_request import DataFrameCopyRequest
+from .data_frame_copy_request_format import DataFrameCopyRequestFormat
 from .database_health_response import DatabaseHealthResponse
 from .database_info_response import DatabaseInfoResponse
 from .delete_subgraph_request import DeleteSubgraphRequest
@@ -90,6 +93,9 @@ from .get_current_graph_bill_response_getcurrentgraphbill import (
 from .get_graph_billing_history_response_getgraphbillinghistory import (
   GetGraphBillingHistoryResponseGetgraphbillinghistory,
 )
+from .get_graph_limits_response_getgraphlimits import (
+  GetGraphLimitsResponseGetgraphlimits,
+)
 from .get_graph_monthly_bill_response_getgraphmonthlybill import (
   GetGraphMonthlyBillResponseGetgraphmonthlybill,
 )
@@ -125,9 +131,6 @@ from .health_status import HealthStatus
 from .health_status_details_type_0 import HealthStatusDetailsType0
 from .http_validation_error import HTTPValidationError
 from .initial_entity_data import InitialEntityData
-from .kuzu_backup_health_response_kuzubackuphealth import (
-  KuzuBackupHealthResponseKuzubackuphealth,
-)
 from .link_token_request import LinkTokenRequest
 from .link_token_request_options_type_0 import LinkTokenRequestOptionsType0
 from .link_token_request_provider_type_0 import LinkTokenRequestProviderType0
@@ -166,6 +169,9 @@ from .repository_credits_response import RepositoryCreditsResponse
 from .repository_plan import RepositoryPlan
 from .repository_type import RepositoryType
 from .response_mode import ResponseMode
+from .s3_copy_request import S3CopyRequest
+from .s3_copy_request_file_format import S3CopyRequestFileFormat
+from .s3_copy_request_s3_url_style_type_0 import S3CopyRequestS3UrlStyleType0
 from .schema_export_response import SchemaExportResponse
 from .schema_export_response_data_stats_type_0 import SchemaExportResponseDataStatsType0
 from .schema_export_response_schema_definition_type_0 import (
@@ -210,7 +216,9 @@ from .transaction_summary_response import TransactionSummaryResponse
 from .update_api_key_request import UpdateAPIKeyRequest
 from .update_password_request import UpdatePasswordRequest
 from .update_user_request import UpdateUserRequest
-from .upgrade_subscription_request import UpgradeSubscriptionRequest
+from .url_copy_request import URLCopyRequest
+from .url_copy_request_file_format import URLCopyRequestFileFormat
+from .url_copy_request_headers_type_0 import URLCopyRequestHeadersType0
 from .user_analytics_response import UserAnalyticsResponse
 from .user_analytics_response_api_usage import UserAnalyticsResponseApiUsage
 from .user_analytics_response_graph_usage import UserAnalyticsResponseGraphUsage
@@ -247,7 +255,6 @@ __all__ = (
   "AvailableExtension",
   "AvailableExtensionsResponse",
   "BackupCreateRequest",
-  "BackupExportRequest",
   "BackupListResponse",
   "BackupResponse",
   "BackupRestoreRequest",
@@ -263,6 +270,9 @@ __all__ = (
   "ConnectionResponse",
   "ConnectionResponseMetadata",
   "ConnectionResponseProvider",
+  "CopyResponse",
+  "CopyResponseErrorDetailsType0",
+  "CopyResponseStatus",
   "CreateAPIKeyRequest",
   "CreateAPIKeyResponse",
   "CreateConnectionRequest",
@@ -270,7 +280,6 @@ __all__ = (
   "CreateGraphRequest",
   "CreateSubgraphRequest",
   "CreateSubgraphRequestMetadataType0",
-  "CreditCheckRequest",
   "CreditsSummaryResponse",
   "CreditsSummaryResponseCreditsByAddonItem",
   "CreditSummary",
@@ -283,6 +292,8 @@ __all__ = (
   "CypherQueryRequestParametersType0",
   "DatabaseHealthResponse",
   "DatabaseInfoResponse",
+  "DataFrameCopyRequest",
+  "DataFrameCopyRequestFormat",
   "DeleteSubgraphRequest",
   "DeleteSubgraphResponse",
   "DetailedTransactionsResponse",
@@ -299,6 +310,7 @@ __all__ = (
   "GetCurrentAuthUserResponseGetcurrentauthuser",
   "GetCurrentGraphBillResponseGetcurrentgraphbill",
   "GetGraphBillingHistoryResponseGetgraphbillinghistory",
+  "GetGraphLimitsResponseGetgraphlimits",
   "GetGraphMonthlyBillResponseGetgraphmonthlybill",
   "GetGraphSchemaInfoResponseGetgraphschemainfo",
   "GetGraphUsageDetailsResponseGetgraphusagedetails",
@@ -320,7 +332,6 @@ __all__ = (
   "HealthStatusDetailsType0",
   "HTTPValidationError",
   "InitialEntityData",
-  "KuzuBackupHealthResponseKuzubackuphealth",
   "LinkTokenRequest",
   "LinkTokenRequestOptionsType0",
   "LinkTokenRequestProviderType0",
@@ -351,6 +362,9 @@ __all__ = (
   "RepositoryPlan",
   "RepositoryType",
   "ResponseMode",
+  "S3CopyRequest",
+  "S3CopyRequestFileFormat",
+  "S3CopyRequestS3UrlStyleType0",
   "SchemaExportResponse",
   "SchemaExportResponseDataStatsType0",
   "SchemaExportResponseSchemaDefinitionType0",
@@ -385,7 +399,9 @@ __all__ = (
   "UpdateAPIKeyRequest",
   "UpdatePasswordRequest",
   "UpdateUserRequest",
-  "UpgradeSubscriptionRequest",
+  "URLCopyRequest",
+  "URLCopyRequestFileFormat",
+  "URLCopyRequestHeadersType0",
   "UserAnalyticsResponse",
   "UserAnalyticsResponseApiUsage",
   "UserAnalyticsResponseGraphUsage",