Add Iceberg export commands; change refresh interval to int

Author: kesmit13

singlestoredb/fusion/handlers/export.py

@@ -6,6 +6,7 @@
 from typing import Optional
 
 from .. import result
+from ...management.export import _get_exports
 from ...management.export import ExportService
 from ...management.export import ExportStatus
 from ..handler import SQLHandler
@@ -187,7 +188,8 @@ def _start_export(params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         order_by=order_by or None,
         properties=json.loads(params['properties']) if params['properties'] else None,
         incremental=params.get('incremental', False),
-        refresh_interval=refresh_interval_delta,
+        refresh_interval=int(refresh_interval_delta.total_seconds())
+        if refresh_interval_delta is not None else None,
     ).start()
 
     res = FusionSQLResult()
@@ -387,6 +389,25 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
 StartIncrementalExport.register(overwrite=True)
 
 
+def _format_status(export_id: str, status: ExportStatus) -> Optional[FusionSQLResult]:
+    """Return the status of an export operation."""
+    info = status._info()
+
+    res = FusionSQLResult()
+    res.add_field('ExportID', result.STRING)
+    res.add_field('Status', result.STRING)
+    res.add_field('Message', result.STRING)
+    res.set_rows([
+        (
+            export_id,
+            info.get('status', 'Unknown'),
+            info.get('statusMsg', ''),
+        ),
+    ])
+
+    return res
+
+
 class ShowExport(SQLHandler):
     """
     SHOW EXPORT export_id;
@@ -400,31 +421,92 @@ class ShowExport(SQLHandler):
 
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         wsg = get_workspace_group({})
-        out = ExportStatus(params['export_id'], wsg)
+        return _format_status(
+            params['export_id'], ExportStatus(params['export_id'], wsg),
+        )
 
-        status = out._info()
+
+ShowExport.register(overwrite=True)
+
+
+class ShowExports(SQLHandler):
+    """
+    SHOW EXPORTS [ scope ];
+
+    # Location of the export
+    scope = FOR '<scope>'
+
+    """
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wsg = get_workspace_group({})
+
+        exports = _get_exports(wsg, params.get('scope', 'all'))
 
         res = FusionSQLResult()
         res.add_field('ExportID', result.STRING)
         res.add_field('Status', result.STRING)
         res.add_field('Message', result.STRING)
         res.set_rows([
             (
-                params['export_id'],
-                status.get('status', 'Unknown'),
-                status.get('statusMsg', ''),
-            ),
+                info['egressID'],
+                info.get('status', 'Unknown'),
+                info.get('statusMsg', ''),
+            )
+            for info in [x._info() for x in exports]
         ])
 
         return res
 
 
-ShowExport.register(overwrite=True)
+ShowExports.register(overwrite=True)
+
+
+class SuspendExport(SQLHandler):
+    """
+    SUSPEND EXPORT export_id;
+
+    # ID of export
+    export_id = '<export-id>'
+
+    """
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wsg = get_workspace_group({})
+        service = ExportService.from_export_id(wsg, params['export_id'])
+        return _format_status(params['export_id'], service.suspend())
+
+
+SuspendExport.register(overwrite=True)
+
+
+class ResumeExport(SQLHandler):
+    """
+    RESUME EXPORT export_id;
+
+    # ID of export
+    export_id = '<export-id>'
+
+    """
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wsg = get_workspace_group({})
+        service = ExportService.from_export_id(wsg, params['export_id'])
+        return _format_status(params['export_id'], service.resume())
+
+
+ResumeExport.register(overwrite=True)
 
 
-class StopExport(SQLHandler):
+class DropExport(SQLHandler):
     """
-    STOP EXPORT export_id;
+    DROP EXPORT export_id;
 
     # ID of export
     export_id = '<export-id>'
@@ -436,8 +518,8 @@ class StopExport(SQLHandler):
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         wsg = get_workspace_group({})
         service = ExportService.from_export_id(wsg, params['export_id'])
-        service.stop()
+        service.drop()
         return None
 
 
-StopExport.register(overwrite=True)
+DropExport.register(overwrite=True)

singlestoredb/management/export.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import copy
-import datetime
 import json
 from typing import Any
 from typing import Dict
@@ -29,7 +28,7 @@ class ExportService(object):
     order_by: Optional[List[Dict[str, Dict[str, str]]]]
     properties: Optional[Dict[str, Any]]
     incremental: bool
-    refresh_interval: Optional[datetime.timedelta]
+    refresh_interval: Optional[int]
     export_id: Optional[str]
 
     def __init__(
@@ -43,7 +42,7 @@ def __init__(
         partition_by: Optional[List[Dict[str, str]]] = None,
         order_by: Optional[List[Dict[str, Dict[str, str]]]] = None,
         incremental: bool = False,
-        refresh_interval: Optional[datetime.timedelta] = None,
+        refresh_interval: Optional[int] = None,
         properties: Optional[Dict[str, Any]] = None,
     ):
         #: Workspace group
@@ -156,7 +155,7 @@ def start(self, tags: Optional[List[str]] = None) -> 'ExportStatus':
                     sortOrderSpec=sort_order_spec,
                     properties=self.properties,
                     incremental=self.incremental or None,
-                    refreshInterval=int(self.refresh_interval.total_seconds())
+                    refreshInterval=self.refresh_interval
                     if self.refresh_interval is not None else None,
                 ).items() if v is not None
             },
@@ -166,8 +165,8 @@ def start(self, tags: Optional[List[str]] = None) -> 'ExportStatus':
 
         return ExportStatus(self.export_id, self.workspace_group)
 
-    def stop(self) -> 'ExportStatus':
-        """Stop the export process."""
+    def suspend(self) -> 'ExportStatus':
+        """Suspend the export process."""
         if self._manager is None:
             raise ManagementError(
                 msg='No workspace manager is associated with this object.',
@@ -179,12 +178,50 @@ def stop(self) -> 'ExportStatus':
             )
 
         self._manager._post(
-            f'workspaceGroups/{self.workspace_group.id}/egress/stopTableEgress',
+            f'workspaceGroups/{self.workspace_group.id}/egress/suspendTableEgress',
             json=dict(egressID=self.export_id),
         )
 
         return ExportStatus(self.export_id, self.workspace_group)
 
+    def resume(self) -> 'ExportStatus':
+        """Resume the export process."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        if self.export_id is None:
+            raise ManagementError(
+                msg='Export ID is not set. You must start the export first.',
+            )
+
+        self._manager._post(
+            f'workspaceGroups/{self.workspace_group.id}/egress/resumeTableEgress',
+            json=dict(egressID=self.export_id),
+        )
+
+        return ExportStatus(self.export_id, self.workspace_group)
+
+    def drop(self) -> None:
+        """Drop the export process."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        if self.export_id is None:
+            raise ManagementError(
+                msg='Export ID is not set. You must start the export first.',
+            )
+
+        self._manager._post(
+            f'workspaceGroups/{self.workspace_group.id}/egress/dropTableEgress',
+            json=dict(egressID=self.export_id),
+        )
+
+        return None
+
     def status(self) -> ExportStatus:
         """Get the status of the export process."""
         if self._manager is None:
@@ -238,3 +275,21 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         return self.status
+
+
+def _get_exports(
+    workspace_group: WorkspaceGroup,
+    scope: str = 'all',
+) -> List[ExportStatus]:
+    """Get all exports in the workspace group."""
+    if workspace_group._manager is None:
+        raise ManagementError(
+            msg='No workspace manager is associated with this object.',
+        )
+
+    out = workspace_group._manager._get(
+        f'workspaceGroups/{workspace_group.id}/egress/tableEgressStatus',
+        json=dict(scope=scope),
+    )
+
+    return out.json()