singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/fusion/handlers/export.py

@@ -0,0 +1,525 @@
+#!/usr/bin/env python3
+import datetime
+import json
+from typing import Any
+from typing import Dict
+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
+from ..result import FusionSQLResult
+from .utils import get_workspace_group
+
+
+class CreateClusterIdentity(SQLHandler):
+    """
+    CREATE CLUSTER IDENTITY
+        catalog
+        storage
+    ;
+
+    # Catolog
+    catalog = CATALOG { _catalog_config | _catalog_creds }
+    _catalog_config = CONFIG '<catalog-config>'
+    _catalog_creds = CREDENTIALS '<catalog-creds>'
+
+    # Storage
+    storage = LINK { _link_config | _link_creds }
+    _link_config = S3 CONFIG '<link-config>'
+    _link_creds = CREDENTIALS '<link-creds>'
+
+    Description
+    -----------
+    Create a cluster identity for allowing the export service to access
+    external cloud resources.
+
+    Arguments
+    ---------
+    * ``<catalog-config>`` and ``<catalog-creds>``: Catalog configuration
+      and credentials in JSON format.
+    * ``<link-config>`` and ``<link-creds>``: Storage link configuration
+      and credentials in JSON format.
+
+    Remarks
+    -------
+    * ``CATALOG`` specifies the details of the catalog to connect to.
+    * ``LINK`` specifies the details of the data storage to connect to.
+
+    Example
+    -------
+    The following statement creates a cluster identity for the catalog
+    and link::
+
+        CREATE CLUSTER IDENTITY
+            CATALOG CONFIG '{
+                "catalog_type": "GLUE",
+                "table_format": "ICEBERG",
+                "catalog_id": "13983498723498",
+                "catalog_region": "us-east-1"
+            }'
+            LINK S3 CONFIG '{
+                "region": "us-east-1",
+                "endpoint_url": "s3://bucket-name"
+
+            }'
+        ;
+
+    """
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        # Catalog
+        catalog_config = json.loads(params['catalog'].get('catalog_config', '{}') or '{}')
+        catalog_creds = json.loads(params['catalog'].get('catalog_creds', '{}') or '{}')
+
+        # Storage
+        storage_config = json.loads(params['storage'].get('link_config', '{}') or '{}')
+        storage_creds = json.loads(params['storage'].get('link_creds', '{}') or '{}')
+
+        storage_config['provider'] = 'S3'
+
+        wsg = get_workspace_group({})
+
+        if wsg._manager is None:
+            raise TypeError('no workspace manager is associated with workspace group')
+
+        out = ExportService(
+            wsg,
+            'none',
+            'none',
+            dict(**catalog_config, **catalog_creds),
+            dict(**storage_config, **storage_creds),
+            columns=None,
+        ).create_cluster_identity()
+
+        res = FusionSQLResult()
+        res.add_field('Identity', result.STRING)
+        res.set_rows([(out['identity'],)])
+
+        return res
+
+
+CreateClusterIdentity.register(overwrite=True)
+
+
+def _start_export(params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+    # From table
+    if isinstance(params['from_table'], str):
+        from_database = None
+        from_table = params['from_table']
+    else:
+        from_database, from_table = params['from_table']
+
+    # Catalog
+    catalog_config = json.loads(params['catalog'].get('catalog_config', '{}') or '{}')
+    catalog_creds = json.loads(params['catalog'].get('catalog_creds', '{}') or '{}')
+
+    # Storage
+    storage_config = json.loads(params['storage'].get('link_config', '{}') or '{}')
+    storage_creds = json.loads(params['storage'].get('link_creds', '{}') or '{}')
+
+    storage_config['provider'] = 'S3'
+
+    wsg = get_workspace_group({})
+
+    if from_database is None:
+        raise ValueError('database name must be specified for source table')
+
+    if wsg._manager is None:
+        raise TypeError('no workspace manager is associated with workspace group')
+
+    partition_by = []
+    if params['partition_by']:
+        for key in params['partition_by']:
+            transform = key['partition_key']['transform']['col_transform']
+            part = {}
+            part['transform'] = transform[0].lower()
+            part['name'] = transform[-1]['transform_col']
+            partition_by.append(part)
+
+    order_by = []
+    if params['order_by'] and params['order_by']['by']:
+        for key in params['order_by']['by']:
+            transform = key['transform']['col_transform']
+            order = {}
+            order['transform'] = transform[0].lower()
+            order['name'] = transform[-1]['transform_col']
+            order['direction'] = 'ascending'
+            order['null_order'] = 'nulls_first'
+            if key.get('direction'):
+                if 'desc' in key['direction'].lower():
+                    order['direction'] = 'descending'
+            if key.get('null_order'):
+                if 'last' in key['null_order'].lower():
+                    order['null_order'] = 'nulls_last'
+            order_by.append(order)
+
+    # Refresh interval
+    refresh_interval_delta = None
+    refresh_interval = params.get('refresh_interval', None)
+    if refresh_interval is not None:
+        value = int(refresh_interval['refresh_interval_value'])
+        time_unit = refresh_interval['refresh_interval_time_unit'].upper()
+        if value < 0:
+            raise ValueError('refresh interval must be greater than 0')
+        if time_unit == 'SECONDS':
+            refresh_interval_delta = datetime.timedelta(seconds=int(value))
+        elif time_unit == 'MINUTES':
+            refresh_interval_delta = datetime.timedelta(minutes=int(value))
+        elif time_unit == 'HOURS':
+            refresh_interval_delta = datetime.timedelta(hours=int(value))
+        elif time_unit == 'DAYS':
+            refresh_interval_delta = datetime.timedelta(days=int(value))
+        else:
+            raise ValueError('invalid refresh interval time unit')
+
+    out = ExportService(
+        wsg,
+        from_database,
+        from_table,
+        dict(**catalog_config, **catalog_creds),
+        dict(**storage_config, **storage_creds),
+        columns=None,
+        partition_by=partition_by or None,
+        order_by=order_by or None,
+        properties=json.loads(params['properties']) if params['properties'] else None,
+        incremental=params.get('incremental', False),
+        refresh_interval=int(refresh_interval_delta.total_seconds())
+        if refresh_interval_delta is not None else None,
+    ).start()
+
+    res = FusionSQLResult()
+    res.add_field('ExportID', result.STRING)
+    res.set_rows([(out.export_id,)])
+
+    return res
+
+
+class StartExport(SQLHandler):
+    """
+    START EXPORT
+        from_table
+        catalog
+        storage
+        [ partition_by ]
+        [ order_by ]
+        [ properties ]
+    ;
+
+    # From table
+    from_table = FROM <table>
+
+    # Transforms
+    _col_transform = { VOID | IDENTITY | YEAR | MONTH | DAY | HOUR } ( _transform_col )
+    _transform_col = <column>
+    _arg_transform = { BUCKET | TRUNCATE } ( _transform_col <comma> _transform_arg )
+    _transform_arg = <integer>
+    transform = { _col_transform | _arg_transform }
+
+    # Partitions
+    partition_by = PARTITION BY partition_key,...
+    partition_key = transform
+
+    # Sort order
+    order_by = ORDER BY sort_key,...
+    sort_key = transform [ direction ] [ null_order ]
+    direction = { ASC | DESC | ASCENDING | DESCENDING }
+    null_order = { NULLS_FIRST | NULLS_LAST }
+
+    # Properties
+    properties = PROPERTIES '<json>'
+
+    # Catolog
+    catalog = CATALOG [ _catalog_config ] [ _catalog_creds ]
+    _catalog_config = CONFIG '<catalog-config>'
+    _catalog_creds = CREDENTIALS '<catalog-creds>'
+
+    # Storage
+    storage = LINK [ _link_config ] [ _link_creds ]
+    _link_config = S3 CONFIG '<link-config>'
+    _link_creds = CREDENTIALS '<link-creds>'
+
+    Description
+    -----------
+    Start an export.
+
+    Arguments
+    ---------
+    * ``<catalog-config>`` and ``<catalog-creds>``: The catalog configuration.
+    * ``<link-config>`` and ``<link-creds>``: The storage link configuration.
+
+    Remarks
+    -------
+    * ``FROM <table>`` specifies the SingleStore table to export. The same name will
+      be used for the exported table.
+    * ``CATALOG`` specifies the details of the catalog to connect to.
+    * ``LINK`` specifies the details of the data storage to connect to.
+
+    Examples
+    --------
+    The following statement starts an export operation with the given
+    catalog and link configurations. The source table to export is
+    named "customer_data"::
+
+        START EXPORT FROM my_db.customer_data
+            CATALOG CONFIG '{
+                "catalog_type": "GLUE",
+                "table_format": "ICEBERG",
+                "catalog_id": "13983498723498",
+                "catalog_region": "us-east-1"
+            }'
+            LINK S3 CONFIG '{
+                "region": "us-east-1",
+                "endpoint_url": "s3://bucket-name"
+            }'
+        ;
+
+    """  # noqa
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        return _start_export(params)
+
+
+StartExport.register(overwrite=True)
+
+
+class StartIncrementalExport(SQLHandler):
+    """
+    START INCREMENTAL EXPORT
+        from_table
+        catalog
+        storage
+        [ partition_by ]
+        [ order_by ]
+        [ properties ]
+        [ refresh_interval ]
+    ;
+
+    # From table
+    from_table = FROM <table>
+
+    # Transforms
+    _col_transform = { VOID | IDENTITY | YEAR | MONTH | DAY | HOUR } ( _transform_col )
+    _transform_col = <column>
+    _arg_transform = { BUCKET | TRUNCATE } ( _transform_col <comma> _transform_arg )
+    _transform_arg = <integer>
+    transform = { _col_transform | _arg_transform }
+
+    # Partitions
+    partition_by = PARTITION BY partition_key,...
+    partition_key = transform
+
+    # Sort order
+    order_by = ORDER BY sort_key,...
+    sort_key = transform [ direction ] [ null_order ]
+    direction = { ASC | DESC | ASCENDING | DESCENDING }
+    null_order = { NULLS_FIRST | NULLS_LAST }
+
+    # Properties
+    properties = PROPERTIES '<json>'
+
+    # Catolog
+    catalog = CATALOG [ _catalog_config ] [ _catalog_creds ]
+    _catalog_config = CONFIG '<catalog-config>'
+    _catalog_creds = CREDENTIALS '<catalog-creds>'
+
+    # Storage
+    storage = LINK [ _link_config ] [ _link_creds ]
+    _link_config = S3 CONFIG '<link-config>'
+    _link_creds = CREDENTIALS '<link-creds>'
+
+    # Refresh interval
+    refresh_interval = REFRESH INTERVAL _refresh_interval_value _refresh_interval_time_unit
+    _refresh_interval_value = <integer>
+    _refresh_interval_time_unit = { SECONDS | MINUTES | HOURS | DAYS }
+
+    Description
+    -----------
+    Start an incremental export.
+
+    Arguments
+    ---------
+    * ``<catalog-config>`` and ``<catalog-creds>``: The catalog configuration.
+    * ``<link-config>`` and ``<link-creds>``: The storage link configuration.
+
+    Remarks
+    -------
+    * ``FROM <table>`` specifies the SingleStore table to export. The same name will
+      be used for the exported table.
+    * ``CATALOG`` specifies the details of the catalog to connect to.
+    * ``LINK`` specifies the details of the data storage to connect to.
+    * ``REFRESH INTERVAL`` specifies the interval for refreshing the
+      incremental export. The default is 1 day.
+
+    Examples
+    --------
+    The following statement starts an export operation with the given
+    catalog and link configurations. The source table to export is
+    named "customer_data"::
+
+        START INCREMENTAL EXPORT FROM my_db.customer_data
+            CATALOG CONFIG '{
+                "catalog_type": "GLUE",
+                "table_format": "ICEBERG",
+                "catalog_id": "13983498723498",
+                "catalog_region": "us-east-1"
+            }'
+            LINK S3 CONFIG '{
+                "region": "us-east-1",
+                "endpoint_url": "s3://bucket-name"
+            }'
+            REFRESH INTERVAL 24 HOURS
+        ;
+
+    """  # noqa
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        params['incremental'] = True
+        return _start_export(params)
+
+
+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;
+
+    # ID of export
+    export_id = '<export-id>'
+
+    """
+
+    _enabled = False
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wsg = get_workspace_group({})
+        return _format_status(
+            params['export_id'], ExportStatus(params['export_id'], wsg),
+        )
+
+
+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([
+            (
+                info['egressID'],
+                info.get('status', 'Unknown'),
+                info.get('statusMsg', ''),
+            )
+            for info in [x._info() for x in exports]
+        ])
+
+        return res
+
+
+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 DropExport(SQLHandler):
+    """
+    DROP 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'])
+        service.drop()
+        return None
+
+
+DropExport.register(overwrite=True)