diff --git a/docs/src/api.rst b/docs/src/api.rst index c99706590..b76315d46 100644 --- a/docs/src/api.rst +++ b/docs/src/api.rst @@ -239,6 +239,12 @@ object that can be used to interact with the Management API. :toctree: generated/ manage_workspaces + manage_teams + manage_users + manage_audit_logs + manage_private_connections + manage_storage_dr + manage_metrics WorkspaceManager @@ -430,6 +436,261 @@ FilesObject FilesObject.rmdir +TeamsManager +............ + +TeamsManager objects are returned by the :func:`manage_teams` function. +They allow you to create, retrieve, and manage teams in your organization. + +.. currentmodule:: singlestoredb.management.teams + +.. autosummary:: + :toctree: generated/ + + TeamsManager + TeamsManager.create_team + TeamsManager.get_team + TeamsManager.list_teams + TeamsManager.teams + TeamsManager.delete_team + TeamsManager.update_team + TeamsManager.get_team_identity_roles + + +Team +.... + +Team objects are retrieved from :meth:`TeamsManager.get_team` or by +retrieving an element from :attr:`TeamsManager.teams`. + +.. autosummary:: + :toctree: generated/ + + Team + Team.update + Team.delete + Team.refresh + Team.identity_roles + + +UsersManager +............ + +UsersManager objects are returned by the :func:`manage_users` function. +They allow you to retrieve and manage users in your organization. + +.. currentmodule:: singlestoredb.management.users + +.. autosummary:: + :toctree: generated/ + + UsersManager + UsersManager.get_user + UsersManager.get_user_identity_roles + UsersManager.create_user_invitation + UsersManager.get_user_invitation + UsersManager.list_user_invitations + UsersManager.delete_user_invitation + UsersManager.user_invitations + + +User +.... + +User objects are retrieved from :meth:`UsersManager.get_user`. + +.. autosummary:: + :toctree: generated/ + + User + User.identity_roles + + +UserInvitation +.............. + +UserInvitation objects are returned by the various UsersManager invitation methods. + +.. autosummary:: + :toctree: generated/ + + UserInvitation + + +AuditLogsManager +................ + +AuditLogsManager objects are returned by the :func:`manage_audit_logs` function. +They allow you to retrieve and analyze audit logs for your organization. + +.. currentmodule:: singlestoredb.management.audit_logs + +.. autosummary:: + :toctree: generated/ + + AuditLogsManager + AuditLogsManager.list_audit_logs + AuditLogsManager.audit_logs + AuditLogsManager.get_audit_logs_for_user + AuditLogsManager.get_audit_logs_for_resource + AuditLogsManager.get_failed_actions + AuditLogsManager.get_actions_by_type + + +AuditLog +........ + +AuditLog objects are returned by the various AuditLogsManager methods. + +.. autosummary:: + :toctree: generated/ + + AuditLog + + +PrivateConnectionsManager +......................... + +PrivateConnectionsManager objects are returned by the :func:`manage_private_connections` function. +They allow you to create and manage private connections in your organization. + +.. currentmodule:: singlestoredb.management.private_connections + +.. autosummary:: + :toctree: generated/ + + PrivateConnectionsManager + PrivateConnectionsManager.create_private_connection + PrivateConnectionsManager.get_private_connection + PrivateConnectionsManager.private_connections + PrivateConnectionsManager.delete_private_connection + PrivateConnectionsManager.update_private_connection + + +PrivateConnection +................. + +PrivateConnection objects are retrieved from :meth:`PrivateConnectionsManager.get_private_connection` +or by retrieving an element from :attr:`PrivateConnectionsManager.private_connections`. + +.. autosummary:: + :toctree: generated/ + + PrivateConnection + + +PrivateConnectionKaiInfo +........................ + +PrivateConnectionKaiInfo objects contain KAI-specific information for private connections. + +.. autosummary:: + :toctree: generated/ + + PrivateConnectionKaiInfo + + +PrivateConnectionOutboundAllowList +.................................. + +PrivateConnectionOutboundAllowList objects contain outbound allow list information for private connections. + +.. autosummary:: + :toctree: generated/ + + PrivateConnectionOutboundAllowList + + +IdentityRole +............ + +IdentityRole objects are used by both teams and users management for role information. + +.. currentmodule:: singlestoredb.management.teams + +.. autosummary:: + :toctree: generated/ + + IdentityRole + + +StorageDRManager +................ + +StorageDRManager objects are returned by the :func:`manage_storage_dr` function. +They allow you to manage storage disaster recovery for your organization. + +.. currentmodule:: singlestoredb.management.storage_dr + +.. autosummary:: + :toctree: generated/ + + StorageDRManager + StorageDRManager.get_status + StorageDRManager.get_available_regions + StorageDRManager.setup_storage_dr + StorageDRManager.start_failover + StorageDRManager.start_failback + StorageDRManager.start_pre_provision + StorageDRManager.stop_pre_provision + + +Storage DR +---------- + +Storage Disaster Recovery objects provide information about replicated databases +and disaster recovery regions. + +.. autosummary:: + :toctree: generated/ + + ReplicatedDatabase + StorageDRStatus + StorageDRRegion + StorageDRCompute + + +MetricsManager +.............. + +MetricsManager objects are returned by the :func:`manage_metrics` function. +They allow you to retrieve metrics for your organization. + +.. currentmodule:: singlestoredb.management.metrics + +.. autosummary:: + :toctree: generated/ + + MetricsManager + MetricsManager.get_workspace_group_metrics + + +Metrics +------- + +Metrics objects provide workspace group metrics and data points. + +.. autosummary:: + :toctree: generated/ + + WorkspaceGroupMetrics + MetricDataPoint + + +Billing Usage +------------- + +Billing Usage objects provide usage and billing information for workspaces. + +.. currentmodule:: singlestoredb.management.billing_usage + +.. autosummary:: + :toctree: generated/ + + UsageItem + BillingUsageItem + + Notebook Tools -------------- diff --git a/singlestoredb/__init__.py b/singlestoredb/__init__.py index 3137b66bc..77a4159fe 100644 --- a/singlestoredb/__init__.py +++ b/singlestoredb/__init__.py @@ -26,6 +26,8 @@ ) from .management import ( manage_cluster, manage_workspaces, manage_files, manage_regions, + manage_teams, manage_private_connections, manage_audit_logs, + manage_users, ) from .types import ( Date, Time, Timestamp, DateFromTicks, TimeFromTicks, TimestampFromTicks, diff --git a/singlestoredb/management/__init__.py b/singlestoredb/management/__init__.py index 8a87d2840..7ce3e64b8 100644 --- a/singlestoredb/management/__init__.py +++ b/singlestoredb/management/__init__.py @@ -1,8 +1,14 @@ #!/usr/bin/env python +from .audit_logs import manage_audit_logs from .cluster import manage_cluster from .files import manage_files from .manager import get_token +from .metrics import manage_metrics +from .private_connections import manage_private_connections from .region import manage_regions +from .storage_dr import manage_storage_dr +from .teams import manage_teams +from .users import manage_users from .workspace import get_organization from .workspace import get_secret from .workspace import get_stage diff --git a/singlestoredb/management/audit_logs.py b/singlestoredb/management/audit_logs.py new file mode 100644 index 000000000..916c54412 --- /dev/null +++ b/singlestoredb/management/audit_logs.py @@ -0,0 +1,497 @@ +#!/usr/bin/env python +"""SingleStoreDB Audit Logs Management.""" +import datetime +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +from .manager import Manager +from .utils import camel_to_snake_dict +from .utils import to_datetime +from .utils import vars_to_str + + +class AuditLog(object): + """ + SingleStoreDB audit log entry definition. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`AuditLogsManager`. Audit logs are retrieved using + :meth:`AuditLogsManager.list_audit_logs`. + + Represents an audit log entry from the SingleStore Management API. + Contains information about user actions in the Control Plane that can be + used to track user activity, including Portal activities, workspace operations, + team management, authentication events, and more. + + See Also + -------- + :meth:`AuditLogsManager.list_audit_logs` + :attr:`AuditLogsManager.audit_logs` + + """ + + def __init__( + self, + audit_id: str, + created_at: Union[str, datetime.datetime], + user_id: Optional[str] = None, + user_email: Optional[str] = None, + type: Optional[str] = None, + reason: Optional[str] = None, + source: Optional[str] = None, + user_type: Optional[str] = None, + org_id: Optional[str] = None, + project_id: Optional[str] = None, + workspace_id: Optional[str] = None, + cluster_id: Optional[str] = None, + team_id: Optional[str] = None, + session_id: Optional[str] = None, + labels: Optional[List[str]] = None, + attributes: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + first_name: Optional[str] = None, + last_name: Optional[str] = None, + ): + #: Unique ID of the audit log entry + self.id = audit_id + + #: Timestamp of when the audit log entry was created (RFC3339Nano format) + self.created_at = to_datetime(created_at) + + #: ID of the user who performed the action + self.user_id = user_id + + #: Email of the user who performed the action + self.user_email = user_email + + #: The audit log entry type + self.type = type + + #: A human-readable description of what happened + self.reason = reason + + #: The audit log entry source (Portal, Admin, SystemJob) + self.source = source + + #: The type of user that triggered the audit log entry + self.user_type = user_type + + #: Organization ID tied to this event + self.organization_id = org_id + + #: Project ID tied to this event + self.project_id = project_id + + #: Workspace ID tied to this event + self.workspace_id = workspace_id + + #: Database cluster ID tied to this event + self.cluster_id = cluster_id + + #: Team ID tied to this event + self.team_id = team_id + + #: Authorization session ID tied to this event + self.session_id = session_id + + #: A list of audit keywords + self.labels = labels or [] + + #: Additional keys and values that are specific to the audit log type + self.attributes = camel_to_snake_dict(attributes) if attributes else None + + #: Text error message, if any relating to this entry + self.error = error + + #: The first name of a redacted user + self.first_name = first_name + + #: The last name of a redacted user + self.last_name = last_name + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'AuditLog': + """ + Construct an AuditLog from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`AuditLog` + + """ + return cls( + audit_id=obj['auditID'], + created_at=obj['createdAt'], + user_id=obj.get('userID'), + user_email=obj.get('userEmail'), + type=obj.get('type'), + reason=obj.get('reason'), + source=obj.get('source'), + user_type=obj.get('userType'), + org_id=obj.get('orgID'), + project_id=obj.get('projectID'), + workspace_id=obj.get('workspaceID'), + cluster_id=obj.get('clusterID'), + team_id=obj.get('teamID'), + session_id=obj.get('sessionID'), + labels=obj.get('labels'), + attributes=obj.get('attributes'), + error=obj.get('error'), + first_name=obj.get('firstName'), + last_name=obj.get('lastName'), + ) + + +class AuditLogsManager(Manager): + """ + SingleStoreDB audit logs manager. + + This class should be instantiated using :func:`singlestoredb.manage_audit_logs` + or accessed via :attr:`WorkspaceManager.audit_logs`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + + """ + + #: Object type + obj_type = 'audit_log' + + def list_audit_logs( + self, + start_date: Optional[datetime.datetime] = None, + end_date: Optional[datetime.datetime] = None, + log_type: Optional[str] = None, + source: Optional[str] = None, + limit: Optional[int] = None, + next_token: Optional[str] = None, + first_name: Optional[str] = None, + last_name: Optional[str] = None, + ) -> List[AuditLog]: + """ + List audit log entries for the organization. + + Parameters + ---------- + start_date : datetime.datetime, optional + Start date (inclusive) for filtering audit logs in RFC3339 format + end_date : datetime.datetime, optional + End date (inclusive) for filtering audit logs in RFC3339 format + log_type : str, optional + Filter by audit log entry type + source : str, optional + Filter by source (Portal, Admin, SystemJob) + limit : int, optional + Maximum number of entries to return + next_token : str, optional + Token from previous query for pagination + first_name : str, optional + Filter by first name (for user redaction) + last_name : str, optional + Filter by last name (for user redaction) + + Returns + ------- + List[AuditLog] + List of audit log entries + + Examples + -------- + >>> audit_mgr = singlestoredb.manage_audit_logs() + >>> logs = audit_mgr.list_audit_logs( + ... log_type="Login", + ... limit=100 + ... ) + >>> for log in logs: + ... print(f"{log.created_at}: {log.type} by {log.user_email}") + >>> + >>> # Filter by time range + >>> import datetime + >>> start = datetime.datetime.now() - datetime.timedelta(days=7) + >>> recent_logs = audit_mgr.list_audit_logs(start_date=start) + + """ + params = {} + + if start_date: + params['startDate'] = start_date.isoformat() + if end_date: + params['endDate'] = end_date.isoformat() + if log_type: + params['type'] = log_type + if source: + params['source'] = source + if limit: + params['limit'] = str(limit) + if next_token: + params['nextToken'] = next_token + if first_name: + params['firstName'] = first_name + if last_name: + params['lastName'] = last_name + + res = self._get('auditLogs', params=params if params else None) + return [AuditLog.from_dict(item) for item in res.json()['auditLogs']] + + @property + def audit_logs(self) -> List[AuditLog]: + """Return a list of recent audit logs.""" + return self.list_audit_logs(limit=100) + + def get_audit_logs_for_user( + self, + user_email: str, + start_date: Optional[datetime.datetime] = None, + end_date: Optional[datetime.datetime] = None, + limit: Optional[int] = None, + ) -> List[AuditLog]: + """ + Get audit logs for a specific user by email. + + Note: The API doesn't support filtering by user_id directly, so this method + retrieves all logs and filters them client-side by user_email. + + Parameters + ---------- + user_email : str + Email address of the user + start_date : datetime.datetime, optional + Start date for filtering audit logs + end_date : datetime.datetime, optional + End date for filtering audit logs + limit : int, optional + Maximum number of entries to return + + Returns + ------- + List[AuditLog] + List of audit log entries for the user + + Examples + -------- + >>> audit_mgr = singlestoredb.manage_audit_logs() + >>> user_logs = audit_mgr.get_audit_logs_for_user("user@example.com") + >>> print(f"Found {len(user_logs)} log entries for user") + + """ + # Get all logs and filter client-side since API doesn't support user_id filter + all_logs = self.list_audit_logs( + start_date=start_date, + end_date=end_date, + limit=limit, + ) + # Filter for logs that match the user email + return [log for log in all_logs if log.user_email == user_email] + + def get_audit_logs_for_resource( + self, + resource_type: str, + resource_id: str, + start_date: Optional[datetime.datetime] = None, + end_date: Optional[datetime.datetime] = None, + limit: Optional[int] = None, + ) -> List[AuditLog]: + """ + Get audit logs for a specific resource. + + Note: The API doesn't support filtering by resource IDs directly, so this method + retrieves all logs and filters them client-side by checking the attributes. + + Parameters + ---------- + resource_type : str + Type of the resource (workspace, cluster, team, project, organization) + resource_id : str + ID of the resource + start_date : datetime.datetime, optional + Start date for filtering audit logs + end_date : datetime.datetime, optional + End date for filtering audit logs + limit : int, optional + Maximum number of entries to return + + Returns + ------- + List[AuditLog] + List of audit log entries for the resource + + Examples + -------- + >>> audit_mgr = singlestoredb.manage_audit_logs() + >>> workspace_logs = audit_mgr.get_audit_logs_for_resource( + ... "workspace", "ws-123" + ... ) + >>> print(f"Found {len(workspace_logs)} log entries for workspace") + + """ + # Get all logs and filter client-side since API doesn't support resource filters + all_logs = self.list_audit_logs( + start_date=start_date, + end_date=end_date, + limit=limit, + ) + + # Filter for logs that match the resource + filtered_logs = [] + for log in all_logs: + if resource_type.lower() == 'workspace' and log.workspace_id == resource_id: + filtered_logs.append(log) + elif resource_type.lower() == 'cluster' and log.cluster_id == resource_id: + filtered_logs.append(log) + elif resource_type.lower() == 'team' and log.team_id == resource_id: + filtered_logs.append(log) + elif resource_type.lower() == 'project' and log.project_id == resource_id: + filtered_logs.append(log) + elif ( + resource_type.lower() == 'organization' and + log.organization_id == resource_id + ): + filtered_logs.append(log) + + return filtered_logs + + def get_failed_actions( + self, + start_date: Optional[datetime.datetime] = None, + end_date: Optional[datetime.datetime] = None, + limit: Optional[int] = None, + ) -> List[AuditLog]: + """ + Get audit logs that contain error messages. + + Note: This method filters for logs that have error messages, as the + audit log schema doesn't have a simple success/failure boolean field. + + Parameters + ---------- + start_date : datetime.datetime, optional + Start date for filtering audit logs + end_date : datetime.datetime, optional + End date for filtering audit logs + limit : int, optional + Maximum number of entries to return + + Returns + ------- + List[AuditLog] + List of audit log entries that contain error messages + + Examples + -------- + >>> audit_mgr = singlestoredb.manage_audit_logs() + >>> failed_logs = audit_mgr.get_failed_actions(limit=50) + >>> for log in failed_logs: + ... print(f"{log.created_at}: {log.type} - {log.error}") + + """ + # Get all logs and filter for those with error messages + all_logs = self.list_audit_logs( + start_date=start_date, + end_date=end_date, + limit=limit, + ) + # Filter for logs that have error messages + return [log for log in all_logs if log.error] + + def get_actions_by_type( + self, + log_type: str, + start_date: Optional[datetime.datetime] = None, + end_date: Optional[datetime.datetime] = None, + limit: Optional[int] = None, + ) -> List[AuditLog]: + """ + Get audit logs for a specific log type. + + Parameters + ---------- + log_type : str + Type of audit log entry to filter by (e.g., "Login", "Logout", etc.) + start_date : datetime.datetime, optional + Start date for filtering audit logs + end_date : datetime.datetime, optional + End date for filtering audit logs + limit : int, optional + Maximum number of entries to return + + Returns + ------- + List[AuditLog] + List of audit log entries for the log type + + Examples + -------- + >>> audit_mgr = singlestoredb.manage_audit_logs() + >>> login_logs = audit_mgr.get_actions_by_type("Login") + >>> print(f"Found {len(login_logs)} login events") + + """ + return self.list_audit_logs( + log_type=log_type, + start_date=start_date, + end_date=end_date, + limit=limit, + ) + + +def manage_audit_logs( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> AuditLogsManager: + """ + Retrieve a SingleStoreDB audit logs manager. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + organization_id : str, optional + ID of organization, if using a JWT for authentication + + Returns + ------- + :class:`AuditLogsManager` + + Examples + -------- + >>> import singlestoredb as s2 + >>> audit_mgr = s2.manage_audit_logs() + >>> logs = audit_mgr.audit_logs + >>> print(f"Found {len(logs)} recent audit log entries") + + """ + return AuditLogsManager( + access_token=access_token, + base_url=base_url, + version=version, + organization_id=organization_id, + ) diff --git a/singlestoredb/management/billing_usage.py b/singlestoredb/management/billing_usage.py index 24c8683dc..ded596f42 100644 --- a/singlestoredb/management/billing_usage.py +++ b/singlestoredb/management/billing_usage.py @@ -82,7 +82,7 @@ def from_dict( owner_id=obj['ownerId'], resource_id=obj['resourceId'], resource_name=obj['resourceName'], - resource_type=obj['resource_type'], + resource_type=obj['resourceType'], value=obj['value'], ) out._manager = manager @@ -142,7 +142,7 @@ def from_dict( out = cls( description=obj['description'], metric=str(camel_to_snake(obj['metric'])), - usage=[UsageItem.from_dict(x, manager) for x in obj['Usage']], + usage=[UsageItem.from_dict(x, manager) for x in obj['usage']], ) out._manager = manager return out diff --git a/singlestoredb/management/manager.py b/singlestoredb/management/manager.py index 9474360a8..85667f602 100644 --- a/singlestoredb/management/manager.py +++ b/singlestoredb/management/manager.py @@ -20,7 +20,8 @@ def set_organization(kwargs: Dict[str, Any]) -> None: """Set the organization ID in the dictionary.""" - if kwargs.get('params', {}).get('organizationID', None): + params = kwargs.get('params') or {} + if params.get('organizationID', None): return org = os.environ.get('SINGLESTOREDB_ORGANIZATION') diff --git a/singlestoredb/management/metrics.py b/singlestoredb/management/metrics.py new file mode 100644 index 000000000..bd66d8b86 --- /dev/null +++ b/singlestoredb/management/metrics.py @@ -0,0 +1,382 @@ +#!/usr/bin/env python +"""SingleStoreDB Metrics Management.""" +import re +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from .manager import Manager + + +class MetricDataPoint(object): + """ + A single metric data point from OpenMetrics format. + + This object represents a single measurement value with labels. + + """ + + def __init__( + self, + metric_name: str, + value: float, + labels: Optional[Dict[str, str]] = None, + ): + #: Name of the metric + self.metric_name = metric_name + + #: Value of the measurement + self.value = value + + #: Labels associated with this metric + self.labels = labels or {} + + def __str__(self) -> str: + """Return string representation.""" + labels_str = ','.join(f'{k}="{v}"' for k, v in self.labels.items()) + return f'{self.metric_name}{{{labels_str}}} {self.value}' + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + +class WorkspaceGroupMetrics(object): + """ + Workspace group metrics definition. + + This object represents metrics for a workspace group, containing + parsed OpenMetrics data. + + """ + + def __init__( + self, + workspace_group_id: str, + raw_metrics: str, + data_points: Optional[List[MetricDataPoint]] = None, + ): + #: Workspace group ID these metrics belong to + self.workspace_group_id = workspace_group_id + + #: Raw OpenMetrics text response + self.raw_metrics = raw_metrics + + #: Parsed metric data points + self.data_points = data_points or [] + + def __str__(self) -> str: + """Return string representation.""" + return ( + f'WorkspaceGroupMetrics(workspace_group_id={self.workspace_group_id}, ' + f'data_points={len(self.data_points)})' + ) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_openmetrics_text( + cls, + workspace_group_id: str, + metrics_text: str, + ) -> 'WorkspaceGroupMetrics': + """ + Parse OpenMetrics text format into structured data. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + metrics_text : str + Raw OpenMetrics text response + + Returns + ------- + :class:`WorkspaceGroupMetrics` + + """ + data_points = [] + + # Parse OpenMetrics format + # Example: singlestoredb_cloud_threads_running{extractor="...",node="..."} 1 + pattern = r'([a-zA-Z_:][a-zA-Z0-9_:]*)\{([^}]*)\}\s+([0-9.-]+)' + + for line in metrics_text.split('\n'): + line = line.strip() + if line.startswith('#') or not line: + continue + + match = re.match(pattern, line) + if match: + metric_name = match.group(1) + labels_str = match.group(2) + value = float(match.group(3)) + + # Parse labels + labels = {} + if labels_str: + # Parse label=value pairs + label_pattern = r'([^=,]+)="([^"]*)"' + for label_match in re.finditer(label_pattern, labels_str): + key = label_match.group(1).strip() + val = label_match.group(2) + labels[key] = val + + data_points.append( + MetricDataPoint( + metric_name=metric_name, + value=value, + labels=labels, + ), + ) + + return cls( + workspace_group_id=workspace_group_id, + raw_metrics=metrics_text, + data_points=data_points, + ) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'WorkspaceGroupMetrics': + """ + Construct a WorkspaceGroupMetrics from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values containing metric data + + Returns + ------- + :class:`WorkspaceGroupMetrics` + + """ + workspace_group_id = obj.get('workspaceGroupId', '') + metric_name = obj.get('metricName', '') + + # Convert dict data to data points + data_points = [] + if 'dataPoints' in obj: + for dp in obj['dataPoints']: + data_points.append( + MetricDataPoint( + metric_name=metric_name, + value=float(dp.get('value', 0)), + labels=dp.get('labels', {}), + ), + ) + elif 'value' in obj: + # Single data point + data_points.append( + MetricDataPoint( + metric_name=metric_name, + value=float(obj['value']), + labels=obj.get('labels', {}), + ), + ) + + return cls( + workspace_group_id=workspace_group_id, + raw_metrics='', # No raw metrics for JSON data + data_points=data_points, + ) + + def get_metrics_by_name(self, metric_name: str) -> List[MetricDataPoint]: + """ + Get all data points for a specific metric name. + + Parameters + ---------- + metric_name : str + Name of the metric to filter by + + Returns + ------- + List[MetricDataPoint] + List of data points matching the metric name + + Examples + -------- + >>> metrics = workspace_group.get_metrics() + >>> cpu_metrics = metrics.get_metrics_by_name( + ... "singlestoredb_cloud_cpu_usage" + ... ) + >>> for point in cpu_metrics: + ... print(f"Node {point.labels.get('node')}: {point.value}%") + + """ + return [ + dp for dp in self.data_points if dp.metric_name == metric_name + ] + + def get_metrics_by_label( + self, label_key: str, label_value: str, + ) -> List[MetricDataPoint]: + """ + Get all data points that have a specific label value. + + Parameters + ---------- + label_key : str + Label key to filter by + label_value : str + Label value to filter by + + Returns + ------- + List[MetricDataPoint] + List of data points matching the label + + Examples + -------- + >>> metrics = workspace_group.get_metrics() + >>> node_metrics = metrics.get_metrics_by_label("node", "aggregator-0") + >>> for point in node_metrics: + ... print(f"{point.metric_name}: {point.value}") + + """ + return [ + dp for dp in self.data_points + if dp.labels.get(label_key) == label_value + ] + + @property + def metric_names(self) -> List[str]: + """ + Get list of all unique metric names. + + Returns + ------- + List[str] + List of unique metric names + + """ + return list(set(dp.metric_name for dp in self.data_points)) + + @property + def metric_name(self) -> str: + """ + Get the primary metric name. + + Returns the first metric name if there are multiple metrics, + or empty string if no metrics. + + Returns + ------- + str + Primary metric name + + """ + names = self.metric_names + return names[0] if names else '' + + +class MetricsManager(Manager): + """ + SingleStoreDB metrics manager. + + This class should be instantiated using + :func:`singlestoredb.manage_metrics` or accessed via + :attr:`WorkspaceManager.metrics`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + + """ + + #: Object type + obj_type = 'metrics' + + def get_workspace_group_metrics( + self, + organization_id: str, + workspace_group_id: str, + ) -> WorkspaceGroupMetrics: + """ + Get metrics for a workspace group in OpenMetrics format. + + Parameters + ---------- + organization_id : str + ID of the organization + workspace_group_id : str + ID of the workspace group + + Returns + ------- + :class:`WorkspaceGroupMetrics` + Parsed metrics data + + Examples + -------- + >>> metrics_mgr = singlestoredb.manage_metrics() + >>> metrics = metrics_mgr.get_workspace_group_metrics("org-123", "wg-456") + >>> cpu_metrics = metrics.get_metrics_by_name( + ... "singlestoredb_cloud_cpu_usage" + ... ) + >>> print(f"Found {len(cpu_metrics)} CPU data points") + + """ + url = ( + f'organizations/{organization_id}/' + f'workspaceGroups/{workspace_group_id}/metrics' + ) + res = self._get(url) + + # The API returns text/plain OpenMetrics format + metrics_text = res.text + + return WorkspaceGroupMetrics.from_openmetrics_text( + workspace_group_id=workspace_group_id, + metrics_text=metrics_text, + ) + + +def manage_metrics( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> MetricsManager: + """ + Retrieve a SingleStoreDB metrics manager. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + organization_id : str, optional + ID of organization, if using a JWT for authentication + + Returns + ------- + :class:`MetricsManager` + + Examples + -------- + >>> import singlestoredb as s2 + >>> metrics_mgr = s2.manage_metrics() + >>> metrics = metrics_mgr.get_workspace_group_metrics("org-123", "wg-456") + >>> print(f"Found {len(metrics.data_points)} metric data points") + + """ + return MetricsManager( + access_token=access_token, + base_url=base_url, + version=version, + organization_id=organization_id, + ) diff --git a/singlestoredb/management/private_connections.py b/singlestoredb/management/private_connections.py new file mode 100644 index 000000000..876de993a --- /dev/null +++ b/singlestoredb/management/private_connections.py @@ -0,0 +1,515 @@ +#!/usr/bin/env python +"""SingleStoreDB Private Connections Management.""" +import datetime +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +from ..exceptions import ManagementError +from .manager import Manager +from .utils import NamedList +from .utils import to_datetime +from .utils import vars_to_str + + +class PrivateConnection(object): + """ + SingleStoreDB private connection definition. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`PrivateConnectionsManager`. Private connections are + created using :meth:`PrivateConnectionsManager.create_private_connection`, or + existing private connections are accessed by either + :attr:`PrivateConnectionsManager.private_connections` or by calling + :meth:`PrivateConnectionsManager.get_private_connection`. + + See Also + -------- + :meth:`PrivateConnectionsManager.create_private_connection` + :meth:`PrivateConnectionsManager.get_private_connection` + :attr:`PrivateConnectionsManager.private_connections` + + """ + + def __init__( + self, + private_connection_id: str, + workspace_group_id: str, + service_name: Optional[str] = None, + connection_type: Optional[str] = None, + status: Optional[str] = None, + allow_list: Optional[str] = None, + outbound_allow_list: Optional[str] = None, + allowed_private_link_ids: Optional[List[str]] = None, + kai_endpoint_id: Optional[str] = None, + sql_port: Optional[int] = None, + websockets_port: Optional[int] = None, + endpoint: Optional[str] = None, + workspace_id: Optional[str] = None, + created_at: Optional[Union[str, datetime.datetime]] = None, + updated_at: Optional[Union[str, datetime.datetime]] = None, + active_at: Optional[Union[str, datetime.datetime]] = None, + deleted_at: Optional[Union[str, datetime.datetime]] = None, + ): + #: Unique ID of the private connection + self.id = private_connection_id + + #: ID of the workspace group containing the private connection + self.workspace_group_id = workspace_group_id + + #: Name of the private connection service + self.service_name = service_name + + #: The private connection type (INBOUND, OUTBOUND) + self.type = connection_type + + #: Status of the private connection (PENDING, ACTIVE, DELETED) + self.status = status + + #: The private connection allow list (account ID for AWS, + #: subscription ID for Azure, project name for GCP) + self.allow_list = allow_list + + #: The account ID allowed for outbound connections + self.outbound_allow_list = outbound_allow_list + + #: List of allowed Private Link IDs + self.allowed_private_link_ids = allowed_private_link_ids or [] + + #: VPC Endpoint ID for AWS + self.kai_endpoint_id = kai_endpoint_id + + #: The SQL port + self.sql_port = sql_port + + #: The websockets port + self.websockets_port = websockets_port + + #: The service endpoint + self.endpoint = endpoint + + #: ID of the workspace to connect with + self.workspace_id = workspace_id + + #: Timestamp of when the private connection was created + self.created_at = to_datetime(created_at) + + #: Timestamp of when the private connection was last updated + self.updated_at = to_datetime(updated_at) + + #: Timestamp of when the private connection became active + self.active_at = to_datetime(active_at) + + #: Timestamp of when the private connection was deleted + self.deleted_at = to_datetime(deleted_at) + + self._manager: Optional['PrivateConnectionsManager'] = None + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict( + cls, obj: Dict[str, Any], + manager: 'PrivateConnectionsManager', + ) -> 'PrivateConnection': + """ + Construct a PrivateConnection from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : PrivateConnectionsManager + The PrivateConnectionsManager the PrivateConnection belongs to + + Returns + ------- + :class:`PrivateConnection` + + """ + out = cls( + private_connection_id=obj['privateConnectionID'], + workspace_group_id=obj['workspaceGroupID'], + service_name=obj.get('serviceName'), + connection_type=obj.get('type'), + status=obj.get('status'), + allow_list=obj.get('allowList'), + outbound_allow_list=obj.get('outboundAllowList'), + allowed_private_link_ids=obj.get('allowedPrivateLinkIDs', []), + kai_endpoint_id=obj.get('kaiEndpointID'), + sql_port=obj.get('sqlPort'), + websockets_port=obj.get('websocketsPort'), + endpoint=obj.get('endpoint'), + workspace_id=obj.get('workspaceID'), + created_at=obj.get('createdAt'), + updated_at=obj.get('updatedAt'), + active_at=obj.get('activeAt'), + deleted_at=obj.get('deletedAt'), + ) + out._manager = manager + return out + + def update( + self, + allow_list: Optional[str] = None, + ) -> None: + """ + Update the private connection definition. + + Parameters + ---------- + allow_list : str, optional + The private connection allow list + + """ + if self._manager is None: + raise ManagementError( + msg='No private connections manager is associated with this object.', + ) + + data = {} + if allow_list is not None: + data['allowList'] = allow_list + + if not data: + return + + self._manager._patch(f'privateConnections/{self.id}', json=data) + self.refresh() + + def delete(self) -> None: + """Delete the private connection.""" + if self._manager is None: + raise ManagementError( + msg='No private connections manager is associated with this object.', + ) + self._manager._delete(f'privateConnections/{self.id}') + + def refresh(self) -> 'PrivateConnection': + """Update the object to the current state.""" + if self._manager is None: + raise ManagementError( + msg='No private connections manager is associated with this object.', + ) + new_obj = self._manager.get_private_connection(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + +class PrivateConnectionKaiInfo(object): + """ + SingleStore Kai private connection information. + + This object contains information needed to create a private connection + to SingleStore Kai for a workspace. + + """ + + def __init__( + self, + service_name: str, + ): + #: VPC Endpoint Service Name for AWS + self.service_name = service_name + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'PrivateConnectionKaiInfo': + """ + Construct a PrivateConnectionKaiInfo from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`PrivateConnectionKaiInfo` + """ + return cls( + service_name=obj['serviceName'], + ) + + +class PrivateConnectionOutboundAllowList(object): + """ + Outbound allow list for a workspace. + + """ + + def __init__( + self, + outbound_allow_list: str, + ): + #: The account ID allowed for outbound connections + self.outbound_allow_list = outbound_allow_list + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'PrivateConnectionOutboundAllowList': + """ + Construct a PrivateConnectionOutboundAllowList from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`PrivateConnectionOutboundAllowList` + + """ + return cls( + outbound_allow_list=obj['outboundAllowList'], + ) + + +class PrivateConnectionsManager(Manager): + """ + SingleStoreDB private connections manager. + + This class should be instantiated using + :func:`singlestoredb.manage_private_connections` or accessed via + :attr:`WorkspaceManager.private_connections`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + + """ + + #: Object type + obj_type = 'private_connection' + + def create_private_connection( + self, + workspace_group_id: str, + service_name: Optional[str] = None, + connection_type: Optional[str] = None, + kai_endpoint_id: Optional[str] = None, + allow_list: Optional[str] = None, + sql_port: Optional[int] = None, + websockets_port: Optional[int] = None, + workspace_id: Optional[str] = None, + ) -> PrivateConnection: + """ + Create a new private connection. + + Parameters + ---------- + workspace_group_id : str + The ID of the workspace group containing the private connection + service_name : str, optional + The name of the private connection service + connection_type : str, optional + The private connection type ('INBOUND', 'OUTBOUND') + kai_endpoint_id : str, optional + VPC Endpoint ID for AWS + allow_list : str, optional + The private connection allow list + sql_port : int, optional + The SQL port + websockets_port : int, optional + The websockets port + workspace_id : str, optional + The ID of the workspace to connect with + + Returns + ------- + :class:`PrivateConnection` + + Examples + -------- + >>> pc_mgr = singlestoredb.manage_private_connections() + >>> connection = pc_mgr.create_private_connection( + ... workspace_group_id="wg-123", + ... service_name="My PrivateLink", + ... connection_type="INBOUND", + ... kai_endpoint_id="vpce-123456789abcdef01" + ... ) + + """ + data = { + k: v for k, v in dict( + workspaceGroupID=workspace_group_id, + serviceName=service_name, + type=connection_type, + kaiEndpointID=kai_endpoint_id, + allowList=allow_list, + sqlPort=sql_port, + websocketsPort=websockets_port, + workspaceID=workspace_id, + ).items() if v is not None + } + + res = self._post('privateConnections', json=data) + return self.get_private_connection(res.json()['privateConnectionID']) + + def get_private_connection(self, connection_id: str) -> PrivateConnection: + """ + Retrieve a private connection definition. + + Parameters + ---------- + connection_id : str + ID of the private connection + + Returns + ------- + :class:`PrivateConnection` + + Examples + -------- + >>> pc_mgr = singlestoredb.manage_private_connections() + >>> connection = pc_mgr.get_private_connection("conn-123") + + """ + res = self._get(f'privateConnections/{connection_id}') + return PrivateConnection.from_dict(res.json(), manager=self) + + @property + def private_connections(self) -> NamedList[PrivateConnection]: + """ + List all private connections. + + Returns + ------- + NamedList[PrivateConnection] + List of private connections + + Examples + -------- + >>> pc_mgr = singlestoredb.manage_private_connections() + >>> connections = pc_mgr.private_connections + >>> for conn in connections: + ... print(f"{conn.service_name}: {conn.type}") + + """ + res = self._get('privateConnections') + return NamedList([PrivateConnection.from_dict(item, self) for item in res.json()]) + + def delete_private_connection(self, connection_id: str) -> None: + """ + Delete a private connection. + + Parameters + ---------- + connection_id : str + ID of the private connection to delete + + Examples + -------- + >>> pc_mgr = singlestoredb.manage_private_connections() + >>> pc_mgr.delete_private_connection("conn-123") + """ + self._delete(f'privateConnections/{connection_id}') + + def update_private_connection( + self, + connection_id: str, + allow_list: Optional[str] = None, + ) -> PrivateConnection: + """ + Update a private connection. + + Parameters + ---------- + connection_id : str + ID of the private connection to update + allow_list : str, optional + The private connection allow list + + Returns + ------- + :class:`PrivateConnection` + Updated private connection object + + Examples + -------- + >>> pc_mgr = singlestoredb.manage_private_connections() + >>> connection = pc_mgr.update_private_connection( + ... "conn-123", + ... allow_list="my-allow-list" + ... ) + + """ + data = {} + if allow_list is not None: + data['allowList'] = allow_list + + if not data: + return self.get_private_connection(connection_id) + + self._patch(f'privateConnections/{connection_id}', json=data) + return self.get_private_connection(connection_id) + + +def manage_private_connections( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> PrivateConnectionsManager: + """ + Retrieve a SingleStoreDB private connections manager. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + organization_id : str, optional + ID of organization, if using a JWT for authentication + + Returns + ------- + :class:`PrivateConnectionsManager` + + Examples + -------- + >>> import singlestoredb as s2 + >>> pc_mgr = s2.manage_private_connections() + >>> connections = pc_mgr.private_connections + >>> print(f"Found {len(connections)} private connections") + + """ + return PrivateConnectionsManager( + access_token=access_token, + base_url=base_url, + version=version, + organization_id=organization_id, + ) diff --git a/singlestoredb/management/storage_dr.py b/singlestoredb/management/storage_dr.py new file mode 100644 index 000000000..8cdde7802 --- /dev/null +++ b/singlestoredb/management/storage_dr.py @@ -0,0 +1,468 @@ +#!/usr/bin/env python +"""SingleStoreDB Storage Disaster Recovery Management.""" +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from .manager import Manager +from .utils import vars_to_str + + +class ReplicatedDatabase(object): + """ + Replicated database configuration for Storage DR. + + Represents information related to a database's replication status. + """ + + def __init__( + self, + database_name: str, + region: str, + duplication_state: str, + ): + #: Name of the database + self.database_name = database_name + + #: Name of the region + self.region = region + + #: Duplication state of the database (Pending, Active, Inactive, Error) + self.duplication_state = duplication_state + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary representation.""" + return { + 'databaseName': self.database_name, + 'region': self.region, + 'duplicationState': self.duplication_state, + } + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'ReplicatedDatabase': + """ + Construct a ReplicatedDatabase from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`ReplicatedDatabase` + + """ + return cls( + database_name=obj['databaseName'], + region=obj['region'], + duplication_state=obj['duplicationState'], + ) + + +class StorageDRCompute(object): + """ + Storage DR compute operation information. + + Represents information related to a workspace group's latest storage DR operation. + """ + + def __init__( + self, + storage_dr_type: str, + storage_dr_state: str, + total_workspaces: int, + total_attachments: int, + completed_workspaces: int, + completed_attachments: int, + ): + #: Name of Storage DR operation (Failover, Failback, + #: PreProvisionStart, PreProvisionStop) + self.storage_dr_type = storage_dr_type + + #: Status of Storage DR operation (Active, Completed, Failed, Expired, Canceled) + self.storage_dr_state = storage_dr_state + + #: The total number of workspaces to setup + self.total_workspaces = total_workspaces + + #: The total number of database attachments to setup + self.total_attachments = total_attachments + + #: The number of workspaces that have been setup + self.completed_workspaces = completed_workspaces + + #: The number of database attachments that have been setup + self.completed_attachments = completed_attachments + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'StorageDRCompute': + """ + Construct a StorageDRCompute from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`StorageDRCompute` + + """ + return cls( + storage_dr_type=obj['storageDRType'], + storage_dr_state=obj['storageDRState'], + total_workspaces=obj['totalWorkspaces'], + total_attachments=obj['totalAttachments'], + completed_workspaces=obj['completedWorkspaces'], + completed_attachments=obj['completedAttachments'], + ) + + +class StorageDRStatus(object): + """ + Storage disaster recovery status information. + + Represents Storage DR status information for a workspace group. + """ + + def __init__( + self, + compute: StorageDRCompute, + storage: List[ReplicatedDatabase], + ): + #: Compute operation information + self.compute = compute + + #: List of replicated databases + self.storage = storage + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'StorageDRStatus': + """ + Construct a StorageDRStatus from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`StorageDRStatus` + + """ + compute = StorageDRCompute.from_dict(obj['compute']) + storage = [ReplicatedDatabase.from_dict(db) for db in obj['storage']] + + return cls( + compute=compute, + storage=storage, + ) + + +class StorageDRRegion(object): + """Available region for Storage DR setup.""" + + def __init__( + self, + region_id: str, + region_name: str, + provider: str, + ): + #: Region ID + self.region_id = region_id + + #: Region name + self.region_name = region_name + + #: Cloud provider + self.provider = provider + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'StorageDRRegion': + """ + Construct a StorageDRRegion from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`StorageDRRegion` + + """ + return cls( + region_id=obj['regionID'], + region_name=obj['regionName'], + provider=obj['provider'], + ) + + +class StorageDRManager(Manager): + """ + SingleStoreDB Storage DR manager. + + This class should be instantiated using + :func:`singlestoredb.manage_storage_dr` or accessed via + :attr:`WorkspaceGroupManager.storage_dr`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + + """ + + #: Object type + obj_type = 'storage_dr' + + def get_status(self, workspace_group_id: str) -> StorageDRStatus: + """ + Get Storage DR status for a workspace group. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + + Returns + ------- + :class:`StorageDRStatus` + Storage DR status information + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> status = dr_mgr.get_status("wg-123") + >>> print(f"DR State: {status.compute.storage_dr_state}") + + """ + res = self._get(f'workspaceGroups/{workspace_group_id}/storage/DR/status') + return StorageDRStatus.from_dict(res.json()) + + def get_available_regions(self, workspace_group_id: str) -> List[StorageDRRegion]: + """ + Get available regions for Storage DR setup. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + + Returns + ------- + List[StorageDRRegion] + List of available regions for DR + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> regions = dr_mgr.get_available_regions("wg-123") + >>> for region in regions: + ... print(f"{region.provider}: {region.region_name}") + + """ + res = self._get(f'workspaceGroups/{workspace_group_id}/storage/DR/regions') + return [StorageDRRegion.from_dict(region) for region in res.json()] + + def setup_storage_dr( + self, + workspace_group_id: str, + region_id: str, + database_names: List[str], + auto_replication: Optional[bool] = None, + backup_bucket_kms_key_id: Optional[str] = None, + data_bucket_kms_key_id: Optional[str] = None, + ) -> None: + """ + Setup Storage DR for a workspace group. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + region_id : str + Region ID of the secondary region + database_names : List[str] + List of database names (can be empty if setting up auto-replication) + auto_replication : bool, optional + If true, all existing and future databases will be automatically replicated + backup_bucket_kms_key_id : str, optional + KMS key ID for backup bucket encryption (AWS only) + data_bucket_kms_key_id : str, optional + KMS key ID for data bucket encryption (AWS only) + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> dr_mgr.setup_storage_dr( + ... "wg-123", + ... "region-456", + ... ["db1", "db2"], + ... auto_replication=True + ... ) + + """ + data: Dict[str, Any] = { + 'regionID': region_id, + 'databaseNames': database_names, + } + + if auto_replication is not None: + data['autoReplication'] = auto_replication + if backup_bucket_kms_key_id is not None: + data['backupBucketKMSKeyID'] = backup_bucket_kms_key_id + if data_bucket_kms_key_id is not None: + data['dataBucketKMSKeyID'] = data_bucket_kms_key_id + + self._post(f'workspaceGroups/{workspace_group_id}/storage/DR/setup', json=data) + + def start_failover(self, workspace_group_id: str) -> None: + """ + Start failover operation for Storage DR. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> dr_mgr.start_failover("wg-123") + + """ + self._post(f'workspaceGroups/{workspace_group_id}/storage/DR/failover') + + def start_failback(self, workspace_group_id: str) -> None: + """ + Start failback operation for Storage DR. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> dr_mgr.start_failback("wg-123") + + """ + self._post(f'workspaceGroups/{workspace_group_id}/storage/DR/failback') + + def start_pre_provision(self, workspace_group_id: str) -> None: + """ + Start pre-provisioning for Storage DR. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> dr_mgr.start_pre_provision("wg-123") + + """ + self._post(f'workspaceGroups/{workspace_group_id}/storage/DR/startPreProvision') + + def stop_pre_provision(self, workspace_group_id: str) -> None: + """ + Stop pre-provisioning for Storage DR. + + Parameters + ---------- + workspace_group_id : str + ID of the workspace group + + Examples + -------- + >>> dr_mgr = singlestoredb.manage_storage_dr() + >>> dr_mgr.stop_pre_provision("wg-123") + + """ + self._post(f'workspaceGroups/{workspace_group_id}/storage/DR/stopPreProvision') + + +def manage_storage_dr( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> StorageDRManager: + """ + Retrieve a SingleStoreDB Storage DR manager. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + organization_id : str, optional + ID of organization, if using a JWT for authentication + + Returns + ------- + :class:`StorageDRManager` + + Examples + -------- + >>> import singlestoredb as s2 + >>> dr_mgr = s2.manage_storage_dr() + >>> status = dr_mgr.get_status("wg-123") + >>> print(f"DR State: {status.compute.storage_dr_state}") + + """ + return StorageDRManager( + access_token=access_token, + base_url=base_url, + version=version, + organization_id=organization_id, + ) diff --git a/singlestoredb/management/teams.py b/singlestoredb/management/teams.py new file mode 100644 index 000000000..9ba4c1db8 --- /dev/null +++ b/singlestoredb/management/teams.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python +"""SingleStoreDB Teams Management.""" +import datetime +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +from ..exceptions import ManagementError +from .manager import Manager +from .utils import NamedList +from .utils import to_datetime +from .utils import vars_to_str + + +class IdentityRole(object): + """ + Identity role definition. + + This object is not instantiated directly. It is used in results + of API calls on teams and users. + + """ + + def __init__( + self, + role_id: str, + role_name: str, + resource_type: str, + resource_id: str, + granted_at: Union[str, datetime.datetime], + granted_by: str, + ): + #: Role ID + self.role_id = role_id + + #: Role name + self.role_name = role_name + + #: Resource type the role applies to + self.resource_type = resource_type + + #: Resource ID the role applies to + self.resource_id = resource_id + + #: When the role was granted + self.granted_at = to_datetime(granted_at) + + #: Who granted the role + self.granted_by = granted_by + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'IdentityRole': + """ + Construct an IdentityRole from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`IdentityRole` + + """ + return cls( + role_id=obj['roleID'], + role_name=obj['roleName'], + resource_type=obj['resourceType'], + resource_id=obj['resourceID'], + granted_at=obj['grantedAt'], + granted_by=obj['grantedBy'], + ) + + +class Team(object): + """ + SingleStoreDB team definition. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`TeamsManager`. Teams are created using + :meth:`TeamsManager.create_team`, or existing teams are accessed by either + :attr:`TeamsManager.teams` or by calling :meth:`TeamsManager.get_team`. + + See Also + -------- + :meth:`TeamsManager.create_team` + :meth:`TeamsManager.get_team` + :attr:`TeamsManager.teams` + + """ + + def __init__( + self, + team_id: str, + name: str, + description: str, + member_users: Optional[List[Dict[str, Any]]] = None, + member_teams: Optional[List[Dict[str, Any]]] = None, + created_at: Optional[Union[str, datetime.datetime]] = None, + ): + #: Unique ID of the team + self.id = team_id + + #: Name of the team + self.name = name + + #: Description of the team + self.description = description + + #: List of member users with user info + self.member_users = member_users or [] + + #: List of member teams with team info + self.member_teams = member_teams or [] + + #: Timestamp of when the team was created + self.created_at = to_datetime(created_at) + + self._manager: Optional['TeamsManager'] = None + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any], manager: 'TeamsManager') -> 'Team': + """ + Construct a Team from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : TeamsManager + The TeamsManager the Team belongs to + + Returns + ------- + :class:`Team` + + """ + out = cls( + team_id=obj['teamID'], + name=obj['name'], + description=obj['description'], + member_users=obj.get('memberUsers', []), + member_teams=obj.get('memberTeams', []), + created_at=obj.get('createdAt'), + ) + out._manager = manager + return out + + def update( + self, + name: Optional[str] = None, + description: Optional[str] = None, + add_member_user_ids: Optional[List[str]] = None, + add_member_user_emails: Optional[List[str]] = None, + add_member_team_ids: Optional[List[str]] = None, + remove_member_user_ids: Optional[List[str]] = None, + remove_member_user_emails: Optional[List[str]] = None, + remove_member_team_ids: Optional[List[str]] = None, + ) -> None: + """ + Update the team definition. + + Parameters + ---------- + name : str, optional + New name for the team + description : str, optional + New description for the team + add_member_user_ids : List[str], optional + List of user IDs to add as members + add_member_user_emails : List[str], optional + List of user emails to add as members + add_member_team_ids : List[str], optional + List of team IDs to add as members + remove_member_user_ids : List[str], optional + List of user IDs to remove from members + remove_member_user_emails : List[str], optional + List of user emails to remove from members + remove_member_team_ids : List[str], optional + List of team IDs to remove from members + + """ + if self._manager is None: + raise ManagementError( + msg='No teams manager is associated with this object.', + ) + + data = { + k: v for k, v in dict( + name=name, + description=description, + addMemberUserIDs=add_member_user_ids, + addMemberUserEmails=add_member_user_emails, + addMemberTeamIDs=add_member_team_ids, + removeMemberUserIDs=remove_member_user_ids, + removeMemberUserEmails=remove_member_user_emails, + removeMemberTeamIDs=remove_member_team_ids, + ).items() if v is not None + } + + if not data: + return + + self._manager._patch(f'teams/{self.id}', json=data) + self.refresh() + + def delete(self) -> None: + """Delete the team.""" + if self._manager is None: + raise ManagementError( + msg='No teams manager is associated with this object.', + ) + self._manager._delete(f'teams/{self.id}') + + def refresh(self) -> 'Team': + """Update the object to the current state.""" + if self._manager is None: + raise ManagementError( + msg='No teams manager is associated with this object.', + ) + new_obj = self._manager.get_team(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + @property + def identity_roles(self) -> List[IdentityRole]: + """ + Get identity roles granted to this team. + + Returns + ------- + List[IdentityRole] + List of identity roles granted to the team + + """ + if self._manager is None: + raise ManagementError( + msg='No teams manager is associated with this object.', + ) + res = self._manager._get(f'teams/{self.id}/identityRoles') + return [IdentityRole.from_dict(item) for item in res.json()] + + +class TeamsManager(Manager): + """ + SingleStoreDB teams manager. + + This class should be instantiated using :func:`singlestoredb.manage_teams` + or accessed via :attr:`WorkspaceManager.teams`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + + """ + + #: Object type + obj_type = 'team' + + def create_team( + self, + name: str, + description: Optional[str] = None, + ) -> Team: + """ + Create a new team. + + Parameters + ---------- + name : str + Name of the team + description : str, optional + Description of the team + + Returns + ------- + :class:`Team` + + Examples + -------- + >>> teams_mgr = singlestoredb.manage_teams() + >>> team = teams_mgr.create_team( + ... name="Data Science Team", + ... description="Team for data science projects" + ... ) + >>> print(team.name) + Data Science Team + + """ + data = { + 'name': name, + } + if description is not None: + data['description'] = description + + res = self._post('teams', json=data) + return self.get_team(res.json()['teamID']) + + def get_team(self, team_id: str) -> Team: + """ + Retrieve a team definition. + + Parameters + ---------- + team_id : str + ID of the team + + Returns + ------- + :class:`Team` + + Examples + -------- + >>> teams_mgr = singlestoredb.manage_teams() + >>> team = teams_mgr.get_team("team-123") + >>> print(team.name) + My Team + + """ + res = self._get(f'teams/{team_id}') + return Team.from_dict(res.json(), manager=self) + + def list_teams( + self, + name_filter: Optional[str] = None, + description_filter: Optional[str] = None, + ) -> NamedList[Team]: + """ + List all teams for the current organization. + + Parameters + ---------- + name_filter : str, optional + Filter teams by name (substring match) + description_filter : str, optional + Filter teams by description (substring match) + + Returns + ------- + NamedList[Team] + List of teams + + Examples + -------- + >>> teams_mgr = singlestoredb.manage_teams() + >>> teams = teams_mgr.list_teams() + >>> for team in teams: + ... print(f"{team.name}: {team.description}") + + >>> # Filter by name + >>> data_teams = teams_mgr.list_teams(name_filter="data") + + """ + params = { + k: v for k, v in dict( + name=name_filter, + description=description_filter, + ).items() if v is not None + } + + res = self._get('teams', params=params if params else None) + return NamedList([Team.from_dict(item, self) for item in res.json()]) + + @property + def teams(self) -> NamedList[Team]: + """Return a list of available teams.""" + return self.list_teams() + + def delete_team(self, team_id: str) -> None: + """ + Delete a team. + + Parameters + ---------- + team_id : str + ID of the team to delete + + Examples + -------- + >>> teams_mgr = singlestoredb.manage_teams() + >>> teams_mgr.delete_team("team-123") + + """ + self._delete(f'teams/{team_id}') + + def update_team( + self, + team_id: str, + name: Optional[str] = None, + description: Optional[str] = None, + add_member_user_ids: Optional[List[str]] = None, + add_member_user_emails: Optional[List[str]] = None, + add_member_team_ids: Optional[List[str]] = None, + remove_member_user_ids: Optional[List[str]] = None, + remove_member_user_emails: Optional[List[str]] = None, + remove_member_team_ids: Optional[List[str]] = None, + ) -> Team: + """ + Update a team. + + Parameters + ---------- + team_id : str + ID of the team to update + name : str, optional + New name for the team + description : str, optional + New description for the team + add_member_user_ids : List[str], optional + List of user IDs to add as members + add_member_user_emails : List[str], optional + List of user emails to add as members + add_member_team_ids : List[str], optional + List of team IDs to add as members + remove_member_user_ids : List[str], optional + List of user IDs to remove from members + remove_member_user_emails : List[str], optional + List of user emails to remove from members + remove_member_team_ids : List[str], optional + List of team IDs to remove from members + + Returns + ------- + :class:`Team` + Updated team object + + Examples + -------- + >>> teams_mgr = singlestoredb.manage_teams() + >>> team = teams_mgr.update_team( + ... "team-123", + ... name="Updated Team Name", + ... description="Updated description", + ... add_member_user_emails=["user@example.com"] + ... ) + + """ + data = { + k: v for k, v in dict( + name=name, + description=description, + addMemberUserIDs=add_member_user_ids, + addMemberUserEmails=add_member_user_emails, + addMemberTeamIDs=add_member_team_ids, + removeMemberUserIDs=remove_member_user_ids, + removeMemberUserEmails=remove_member_user_emails, + removeMemberTeamIDs=remove_member_team_ids, + ).items() if v is not None + } + + if not data: + return self.get_team(team_id) + + self._patch(f'teams/{team_id}', json=data) + return self.get_team(team_id) + + def get_team_identity_roles(self, team_id: str) -> List[IdentityRole]: + """ + Get identity roles granted to a team. + + Parameters + ---------- + team_id : str + ID of the team + + Returns + ------- + List[IdentityRole] + List of identity roles granted to the team + + Examples + -------- + >>> teams_mgr = singlestoredb.manage_teams() + >>> roles = teams_mgr.get_team_identity_roles("team-123") + >>> for role in roles: + ... print(f"{role.role_name} on {role.resource_type}") + + """ + res = self._get(f'teams/{team_id}/identityRoles') + return [IdentityRole.from_dict(item) for item in res.json()] + + +def manage_teams( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> TeamsManager: + """ + Retrieve a SingleStoreDB teams manager. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + organization_id : str, optional + ID of organization, if using a JWT for authentication + + Returns + ------- + :class:`TeamsManager` + + Examples + -------- + >>> import singlestoredb as s2 + >>> teams_mgr = s2.manage_teams() + >>> teams = teams_mgr.teams + >>> print(f"Found {len(teams)} teams") + + """ + return TeamsManager( + access_token=access_token, + base_url=base_url, + version=version, + organization_id=organization_id, + ) diff --git a/singlestoredb/management/users.py b/singlestoredb/management/users.py new file mode 100644 index 000000000..97ad6ccbb --- /dev/null +++ b/singlestoredb/management/users.py @@ -0,0 +1,486 @@ +#!/usr/bin/env python +"""SingleStoreDB Users Management.""" +import datetime +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +from ..exceptions import ManagementError +from .manager import Manager +from .utils import NamedList +from .utils import to_datetime +from .utils import vars_to_str + + +class IdentityRole(object): + """ + Identity role definition for users. + + This object is not instantiated directly. It is used in results + of API calls on users and teams. + + """ + + def __init__( + self, + role_id: str, + role_name: str, + resource_type: str, + resource_id: str, + granted_at: Union[str, datetime.datetime], + granted_by: str, + ): + #: Role ID + self.role_id = role_id + + #: Role name + self.role_name = role_name + + #: Resource type the role applies to + self.resource_type = resource_type + + #: Resource ID the role applies to + self.resource_id = resource_id + + #: When the role was granted + self.granted_at = to_datetime(granted_at) + + #: Who granted the role + self.granted_by = granted_by + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any]) -> 'IdentityRole': + """ + Construct an IdentityRole from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + + Returns + ------- + :class:`IdentityRole` + + """ + return cls( + role_id=obj['roleID'], + role_name=obj['roleName'], + resource_type=obj['resourceType'], + resource_id=obj['resourceID'], + granted_at=obj['grantedAt'], + granted_by=obj['grantedBy'], + ) + + +class UserInvitation(object): + """ + SingleStoreDB user invitation definition. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`UsersManager`. + + """ + + def __init__( + self, + invitation_id: str, + email: str, + state: str, + created_at: Union[str, datetime.datetime], + acted_at: Optional[Union[str, datetime.datetime]] = None, + message: Optional[str] = None, + team_ids: Optional[List[str]] = None, + ): + #: Unique ID of the invitation + self.id = invitation_id + + #: Email address of the invited user + self.email = email + + #: State of the invitation (Pending, Accepted, Refused, Revoked) + self.state = state + + #: Timestamp of when the invitation was created + self.created_at = to_datetime(created_at) + + #: Timestamp of most recent state change + self.acted_at = to_datetime(acted_at) + + #: Welcome message + self.message = message + + #: List of team IDs the user will be added to + self.team_ids = team_ids or [] + + self._manager: Optional['UsersManager'] = None + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any], manager: 'UsersManager') -> 'UserInvitation': + """ + Construct a UserInvitation from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : UsersManager + The UsersManager the UserInvitation belongs to + + Returns + ------- + :class:`UserInvitation` + + """ + out = cls( + invitation_id=obj['invitationID'], + email=obj['email'], + state=obj['state'], + created_at=obj['createdAt'], + acted_at=obj.get('actedAt'), + message=obj.get('message'), + team_ids=obj.get('teamIDs', []), + ) + out._manager = manager + return out + + +class User(object): + """ + SingleStoreDB user definition. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`UsersManager`. Users are accessed by calling + :meth:`UsersManager.get_user` or :meth:`UsersManager.get_user_identity_roles`. + + See Also + -------- + :meth:`UsersManager.get_user_identity_roles` + + """ + + def __init__( + self, + user_id: str, + email: str, + first_name: str, + last_name: str, + ): + #: Unique ID of the user + self.id = user_id + + #: Email address of the user + self.email = email + + #: First name of the user + self.first_name = first_name + + #: Last name of the user + self.last_name = last_name + + self._manager: Optional['UsersManager'] = None + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any], manager: 'UsersManager') -> 'User': + """ + Construct a User from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : UsersManager + The UsersManager the User belongs to + + Returns + ------- + :class:`User` + """ + out = cls( + user_id=obj['userID'], + email=obj['email'], + first_name=obj['firstName'], + last_name=obj['lastName'], + ) + out._manager = manager + return out + + @property + def identity_roles(self) -> List[IdentityRole]: + """ + Get identity roles granted to this user. + + Returns + ------- + List[IdentityRole] + List of identity roles granted to the user + + Examples + -------- + >>> user = users_mgr.get_user("user-123") + >>> roles = user.identity_roles + >>> for role in roles: + ... print(f"{role.role_name} on {role.resource_type}") + """ + if self._manager is None: + raise ManagementError( + msg='No users manager is associated with this object.', + ) + return self._manager.get_user_identity_roles(self.id) + + +class UsersManager(Manager): + """ + SingleStoreDB users manager. + + This class should be instantiated using :func:`singlestoredb.manage_users` + or accessed via :attr:`WorkspaceManager.users`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + """ + + #: Object type + obj_type = 'user' + + def get_user_identity_roles(self, user_id: str) -> List[IdentityRole]: + """ + Get identity roles granted to a user. + + Parameters + ---------- + user_id : str + ID of the user + + Returns + ------- + List[IdentityRole] + List of identity roles granted to the user + + Examples + -------- + >>> users_mgr = singlestoredb.manage_users() + >>> roles = users_mgr.get_user_identity_roles("user-123") + >>> for role in roles: + ... print(f"{role.role_name} on {role.resource_type} ({role.resource_id})") + ... print(f" Granted by {role.granted_by} at {role.granted_at}") + + """ + res = self._get(f'users/{user_id}/identityRoles') + return [IdentityRole.from_dict(item) for item in res.json()] + + def get_user(self, user_id: str) -> User: + """ + Get basic user information. + + Note: This method creates a User object with the provided user_id. + Full user details may not be available through the current API. + + Parameters + ---------- + user_id : str + ID of the user + + Returns + ------- + :class:`User` + User object + + Examples + -------- + >>> users_mgr = singlestoredb.manage_users() + >>> user = users_mgr.get_user("user-123") + >>> roles = user.identity_roles + + """ + # Note: The API doesn't seem to have a direct GET /users/{userID} endpoint + # based on the documentation provided. We create a basic User object + # that can be used to get identity roles. + user = User( + user_id=user_id, + email='', # Will be populated if/when user details endpoint is available + first_name='', + last_name='', + ) + user._manager = self + return user + + def create_user_invitation( + self, + email: str, + team_ids: Optional[List[str]] = None, + ) -> UserInvitation: + """ + Create a user invitation. + + Parameters + ---------- + email : str + Email address of the user to invite + team_ids : List[str], optional + List of team IDs to add the user to upon acceptance + + Returns + ------- + :class:`UserInvitation` + Created user invitation + + Examples + -------- + >>> users_mgr = singlestoredb.manage_users() + >>> invitation = users_mgr.create_user_invitation( + ... email="user@example.com", + ... team_ids=["team-123"] + ... ) + >>> print(invitation.state) + Pending + + """ + data: Dict[str, Any] = { + 'email': email, + } + if team_ids is not None: + data['teamIDs'] = team_ids + + res = self._post('userInvitations', json=data) + return self.get_user_invitation(res.json()['invitationID']) + + def get_user_invitation(self, invitation_id: str) -> UserInvitation: + """ + Get a user invitation. + + Parameters + ---------- + invitation_id : str + ID of the invitation + + Returns + ------- + :class:`UserInvitation` + User invitation object + + Examples + -------- + >>> users_mgr = singlestoredb.manage_users() + >>> invitation = users_mgr.get_user_invitation("invitation-123") + >>> print(f"Invitation for {invitation.email} is {invitation.state}") + + """ + res = self._get(f'userInvitations/{invitation_id}') + return UserInvitation.from_dict(res.json(), manager=self) + + def list_user_invitations(self) -> NamedList[UserInvitation]: + """ + List all user invitations for the current organization. + + Returns + ------- + NamedList[UserInvitation] + List of user invitations + + Examples + -------- + >>> users_mgr = singlestoredb.manage_users() + >>> invitations = users_mgr.list_user_invitations() + >>> for invitation in invitations: + ... print(f"{invitation.email}: {invitation.state}") + + """ + res = self._get('userInvitations') + return NamedList([UserInvitation.from_dict(item, self) for item in res.json()]) + + def delete_user_invitation(self, invitation_id: str) -> None: + """ + Delete (revoke) a user invitation. + + Parameters + ---------- + invitation_id : str + ID of the invitation to delete + + Examples + -------- + >>> users_mgr = singlestoredb.manage_users() + >>> users_mgr.delete_user_invitation("invitation-123") + + """ + self._delete(f'userInvitations/{invitation_id}') + + @property + def user_invitations(self) -> NamedList[UserInvitation]: + """Return a list of user invitations.""" + return self.list_user_invitations() + + +def manage_users( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> UsersManager: + """ + Retrieve a SingleStoreDB users manager. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the management API + organization_id : str, optional + ID of organization, if using a JWT for authentication + + Returns + ------- + :class:`UsersManager` + + Examples + -------- + >>> import singlestoredb as s2 + >>> users_mgr = s2.manage_users() + >>> # Get roles for a specific user + >>> roles = users_mgr.get_user_identity_roles("user-123") + >>> print(f"User has {len(roles)} identity roles") + + """ + return UsersManager( + access_token=access_token, + base_url=base_url, + version=version, + organization_id=organization_id, + ) diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index 8ba179c30..56e66ec19 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -13,11 +13,20 @@ from typing import Dict from typing import List from typing import Optional +from typing import TYPE_CHECKING from typing import Union +if TYPE_CHECKING: + from .storage_dr import StorageDRRegion, StorageDRStatus + from .private_connections import PrivateConnection + from .private_connections import PrivateConnectionKaiInfo + from .private_connections import PrivateConnectionOutboundAllowList + from .. import config from .. import connection from ..exceptions import ManagementError +from .metrics import WorkspaceGroupMetrics +from .storage_dr import ReplicatedDatabase from .billing_usage import BillingUsageItem from .files import FileLocation from .files import FilesObject @@ -676,6 +685,9 @@ class Workspace(object): resume_attachments: Optional[List[Dict[str, Any]]] scaling_progress: Optional[int] last_resumed_at: Optional[datetime.datetime] + auto_scale: Optional[Dict[str, Any]] + kai_enabled: Optional[bool] + scale_factor: Optional[int] def __init__( self, @@ -693,6 +705,9 @@ def __init__( resume_attachments: Optional[List[Dict[str, Any]]] = None, scaling_progress: Optional[int] = None, last_resumed_at: Optional[Union[str, datetime.datetime]] = None, + auto_scale: Optional[Dict[str, Any]] = None, + kai_enabled: Optional[bool] = None, + scale_factor: Optional[int] = None, ): #: Name of the workspace self.name = name @@ -744,6 +759,15 @@ def __init__( #: Timestamp when workspace was last resumed self.last_resumed_at = to_datetime(last_resumed_at) + #: Autoscaling configuration + self.auto_scale = camel_to_snake_dict(auto_scale) + + #: Whether SingleStore Kai is enabled + self.kai_enabled = kai_enabled + + #: Current scale factor + self.scale_factor = scale_factor + self._manager: Optional[WorkspaceManager] = None def __str__(self) -> str: @@ -786,6 +810,9 @@ def from_dict(cls, obj: Dict[str, Any], manager: 'WorkspaceManager') -> 'Workspa last_resumed_at=obj.get('lastResumedAt'), resume_attachments=obj.get('resumeAttachments'), scaling_progress=obj.get('scalingProgress'), + auto_scale=obj.get('autoScale'), + kai_enabled=obj.get('kaiEnabled'), + scale_factor=obj.get('scaleFactor'), ) out._manager = manager return out @@ -984,6 +1011,183 @@ def resume( ) self.refresh() + @property + def private_connections(self) -> List['PrivateConnection']: + """Return a list of private connections for this workspace.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + res = self._manager._get(f'workspaces/{self.id}/privateConnections') + from .private_connections import PrivateConnection + return [ + PrivateConnection.from_dict(item, None) # type: ignore + for item in res.json() + ] + + @property + def kai_info(self) -> 'PrivateConnectionKaiInfo': + """Get information to create private connection to SingleStore Kai.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + res = self._manager._get(f'workspaces/{self.id}/privateConnections/kai') + return PrivateConnectionKaiInfo.from_dict(res.json()) + + @property + def outbound_allowlist(self) -> 'PrivateConnectionOutboundAllowList': + """Get the outbound allow list for this workspace.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + res = self._manager._get( + f'workspaces/{self.id}/privateConnections/outboundAllowList', + ) + return PrivateConnectionOutboundAllowList.from_dict(res.json()) + + def get_cpu_metrics( + self, + start_time: Optional[datetime.datetime] = None, + end_time: Optional[datetime.datetime] = None, + ) -> Optional['WorkspaceGroupMetrics']: + """ + Get CPU usage metrics for this workspace. + + Parameters + ---------- + start_time : datetime.datetime, optional + Start time for metrics data + end_time : datetime.datetime, optional + End time for metrics data + + Returns + ------- + WorkspaceGroupMetrics or None + CPU usage metric, or None if not available + + Examples + -------- + >>> workspace = manager.get_workspace('workspace-id') + >>> cpu_metric = workspace.get_cpu_metrics() + >>> if cpu_metric: + ... print(f"Current CPU usage: {cpu_metric.get_latest_value()}%") + """ + # Get the workspace group to access metrics + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + workspace_group = self._manager.get_workspace_group(self.group_id) + metrics = workspace_group.get_metrics( + start_time=start_time, + end_time=end_time, + metric_names=['cpu_usage', 'cpu_utilization'], + workspace_id=self.id, + ) + + # Try common CPU metric names + for name in ['cpu_usage', 'cpu_utilization', 'cpu']: + if name in metrics: + return metrics[name] + + return None + + def get_memory_metrics( + self, + start_time: Optional[datetime.datetime] = None, + end_time: Optional[datetime.datetime] = None, + ) -> Optional['WorkspaceGroupMetrics']: + """ + Get memory usage metrics for this workspace. + + Parameters + ---------- + start_time : datetime.datetime, optional + Start time for metrics data + end_time : datetime.datetime, optional + End time for metrics data + + Returns + ------- + WorkspaceGroupMetrics or None + Memory usage metric, or None if not available + + Examples + -------- + >>> workspace = manager.get_workspace('workspace-id') + >>> memory_metric = workspace.get_memory_metrics() + >>> if memory_metric: + ... print(f"Current memory usage: {memory_metric.get_latest_value()} MB") + """ + # Get the workspace group to access metrics + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + workspace_group = self._manager.get_workspace_group(self.group_id) + metrics = workspace_group.get_metrics( + start_time=start_time, + end_time=end_time, + metric_names=['memory_usage', 'memory_utilization'], + workspace_id=self.id, + ) + + # Try common memory metric names + for name in ['memory_usage', 'memory_utilization', 'memory']: + if name in metrics: + return metrics[name] + + return None + + def get_storage_metrics( + self, + start_time: Optional[datetime.datetime] = None, + end_time: Optional[datetime.datetime] = None, + ) -> Optional['WorkspaceGroupMetrics']: + """ + Get storage usage metrics for this workspace. + + Parameters + ---------- + start_time : datetime.datetime, optional + Start time for metrics data + end_time : datetime.datetime, optional + End time for metrics data + + Returns + ------- + WorkspaceGroupMetrics or None + Storage usage metric, or None if not available + + Examples + -------- + >>> workspace = manager.get_workspace('workspace-id') + >>> storage_metric = workspace.get_storage_metrics() + >>> if storage_metric: + ... print(f"Current storage usage: {storage_metric.get_latest_value()} GB") + """ + # Get the workspace group to access metrics + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + workspace_group = self._manager.get_workspace_group(self.group_id) + metrics = workspace_group.get_metrics( + start_time=start_time, + end_time=end_time, + metric_names=['storage_usage', 'disk_usage'], + workspace_id=self.id, + ) + + # Try common storage metric names + for name in ['storage_usage', 'disk_usage', 'storage']: + if name in metrics: + return metrics[name] + + return None + class WorkspaceGroup(object): """ @@ -1288,6 +1492,420 @@ def workspaces(self) -> NamedList[Workspace]: [Workspace.from_dict(item, self._manager) for item in res.json()], ) + @property + def private_connections(self) -> List['PrivateConnection']: + """Return a list of private connections for this workspace group.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + res = self._manager._get(f'workspaceGroups/{self.id}/privateConnections') + from .private_connections import PrivateConnection + return [ + PrivateConnection.from_dict(item, None) # type: ignore + for item in res.json() + ] + + def get_metrics( + self, + start_time: Optional[datetime.datetime] = None, + end_time: Optional[datetime.datetime] = None, + metric_names: Optional[List[str]] = None, + workspace_id: Optional[Union[str, 'Workspace']] = None, + aggregation_type: Optional[str] = None, + resolution: Optional[str] = None, + ) -> Dict[str, 'WorkspaceGroupMetrics']: + """ + Get metrics for this workspace group. + + Parameters + ---------- + start_time : datetime.datetime, optional + Start time for metrics data + end_time : datetime.datetime, optional + End time for metrics data + metric_names : List[str], optional + List of specific metric names to retrieve + workspace_id : str or Workspace, optional + ID of specific workspace to get metrics for, or a Workspace instance + aggregation_type : str, optional + Type of aggregation ('avg', 'sum', 'max', 'min') + resolution : str, optional + Time resolution for data points ('1m', '5m', '1h', '1d') + + Returns + ------- + Dict[str, WorkspaceGroupMetrics] + Dictionary mapping metric names to metric objects + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + params = {} + if start_time: + params['startTime'] = start_time.isoformat() + if end_time: + params['endTime'] = end_time.isoformat() + if metric_names: + params['metricNames'] = ','.join(metric_names) + if workspace_id: + # Handle both string IDs and Workspace instances + if hasattr(workspace_id, 'id'): + params['workspaceID'] = workspace_id.id + else: + params['workspaceID'] = workspace_id + if aggregation_type: + params['aggregationType'] = aggregation_type + if resolution: + params['resolution'] = resolution + + path = ( + f'organizations/{self._manager.organization.id}/workspaceGroups/' + f'{self.id}/metrics' + ) + res = self._manager._get(path, params=params if params else None) + + metrics_data = res.json() + metrics_dict = {} + + # Handle different possible response structures + if isinstance(metrics_data, list): + for metric_obj in metrics_data: + metric = WorkspaceGroupMetrics.from_dict(metric_obj) + metrics_dict[metric.metric_name] = metric + elif isinstance(metrics_data, dict): + if 'metrics' in metrics_data: + for metric_obj in metrics_data['metrics']: + metric = WorkspaceGroupMetrics.from_dict(metric_obj) + metrics_dict[metric.metric_name] = metric + else: + # Assume the dict itself contains metric data + for name, data in metrics_data.items(): + if isinstance(data, dict): + data['metricName'] = name + metric = WorkspaceGroupMetrics.from_dict(data) + metrics_dict[name] = metric + + return metrics_dict + + @property + def storage_dr_status(self) -> Optional['StorageDRStatus']: + """ + Get Storage DR status for this workspace group. + + Returns + ------- + StorageDRStatus or None + Storage DR status information, or None if no manager is associated + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> status = wg.storage_dr_status + >>> if status: + ... print(f"DR enabled: {status.dr_enabled}") + ... print(f"Primary region: {status.primary_region}") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + try: + path = f'workspaceGroups/{self.id}/storage/DR/status' + res = self._manager._get(path) + from .storage_dr import StorageDRStatus + return StorageDRStatus.from_dict(res.json()) + except Exception: + return None + + @property + def available_dr_regions(self) -> List['StorageDRRegion']: + """ + Get available regions for Storage DR setup for this workspace group. + + Returns + ------- + List[StorageDRRegion] + List of available DR regions + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> regions = wg.available_dr_regions + >>> for region in regions: + ... print(f"{region.region_name} ({region.provider})") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + path = f'workspaceGroups/{self.id}/storage/DR/regions' + res = self._manager._get(path) + from .storage_dr import StorageDRRegion + return [StorageDRRegion.from_dict(item) for item in res.json()] + + def setup_storage_dr( + self, + backup_region: str, + replicated_databases: List[Union[str, 'ReplicatedDatabase']], + ) -> Optional['StorageDRStatus']: + """ + Set up Storage DR for this workspace group. + + Parameters + ---------- + backup_region : str + ID of the backup region + replicated_databases : List[str or ReplicatedDatabase] + List of database names or ReplicatedDatabase objects to replicate + + Returns + ------- + StorageDRStatus or None + Updated Storage DR status + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> status = wg.setup_storage_dr( + ... backup_region="us-west-2", + ... replicated_databases=["production_db", "analytics_db"] + ... ) + >>> if status: + ... print(f"DR setup status: {status.status}") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + # Convert database names/objects to config dictionaries + db_configs = [] + for db in replicated_databases: + if isinstance(db, str): + # For string database names, just pass the name + db_configs.append({'databaseName': db}) + else: + # For ReplicatedDatabase objects, use their to_dict method + db_configs.append(db.to_dict()) + + data = { + 'backupRegion': backup_region, + 'replicatedDatabases': db_configs, + } + + path = f'workspaceGroups/{self.id}/storage/DR/setup' + self._manager._post(path, json=data) + + # Return updated status + return self.storage_dr_status + + def start_failover(self) -> Optional['StorageDRStatus']: + """ + Start failover to the secondary region for this workspace group. + + Returns + ------- + StorageDRStatus or None + Updated Storage DR status + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> status = wg.start_failover() + >>> if status: + ... print(f"Failover status: {status.failover_status}") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + path = f'workspaceGroups/{self.id}/storage/DR/failover' + self._manager._patch(path) + return self.storage_dr_status + + def start_failback(self) -> Optional['StorageDRStatus']: + """ + Start failback to the primary region for this workspace group. + + Returns + ------- + StorageDRStatus or None + Updated Storage DR status + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> status = wg.start_failback() + >>> if status: + ... print(f"Failback status: {status.status}") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + path = f'workspaceGroups/{self.id}/storage/DR/failback' + self._manager._patch(path) + return self.storage_dr_status + + def start_pre_provision(self) -> Optional['StorageDRStatus']: + """ + Start pre-provisioning from primary region for this workspace group. + + Returns + ------- + StorageDRStatus or None + Updated Storage DR status + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> status = wg.start_pre_provision() + >>> if status: + ... print(f"Pre-provision status: {status.pre_provision_status}") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + path = f'workspaceGroups/{self.id}/storage/DR/startPreProvision' + self._manager._patch(path) + return self.storage_dr_status + + def stop_pre_provision(self) -> Optional['StorageDRStatus']: + """ + Stop pre-provisioning from primary region for this workspace group. + + Returns + ------- + StorageDRStatus or None + Updated Storage DR status + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> status = wg.stop_pre_provision() + >>> if status: + ... print(f"Pre-provision status: {status.pre_provision_status}") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + path = f'workspaceGroups/{self.id}/storage/DR/stopPreProvision' + self._manager._patch(path) + return self.storage_dr_status + + def update_retention_period(self, retention_days: int) -> None: + """ + Update the retention period for continuous backups for this workspace group. + + Parameters + ---------- + retention_days : int + Number of days to retain backups + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> wg.update_retention_period(retention_days=30) + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + data = { + 'retentionDays': retention_days, + } + + path = f'workspaceGroups/{self.id}/storage/retentionPeriod' + self._manager._patch(path, json=data) + + def wait_for_dr_operation( + self, + operation_type: str, + target_status: str, + interval: int = 30, + timeout: int = 3600, + ) -> Optional['StorageDRStatus']: + """ + Wait for a DR operation to complete for this workspace group. + + Parameters + ---------- + operation_type : str + Type of operation ('failover', 'failback', 'pre_provision') + target_status : str + Target status to wait for + interval : int, optional + Polling interval in seconds + timeout : int, optional + Maximum time to wait in seconds + + Returns + ------- + StorageDRStatus or None + Final Storage DR status + + Raises + ------ + ManagementError + If timeout is reached + + Examples + -------- + >>> wg = workspace_mgr.get_workspace_group("wg-123") + >>> wg.start_failover() + >>> final_status = wg.wait_for_dr_operation("failover", "completed") + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + import time + + elapsed = 0 + while elapsed < timeout: + status = self.storage_dr_status + if status is None: + raise ManagementError(msg='Unable to get storage DR status') + + if ( + operation_type == 'failover' and + status.compute.storage_dr_state == target_status + ): + return status + elif ( + operation_type == 'failback' and + status.compute.storage_dr_state == target_status + ): + return status + elif ( + operation_type == 'pre_provision' and + status.compute.storage_dr_state == target_status + ): + return status + + time.sleep(interval) + elapsed += interval + + raise ManagementError( + msg=( + f'Timeout waiting for {operation_type} operation to ' + f'reach {target_status}' + ), + ) + class StarterWorkspace(object): """ @@ -1557,7 +2175,7 @@ def usage( metric=snake_to_camel(metric), startTime=from_datetime(start_time), endTime=from_datetime(end_time), - aggregate_by=aggregate_by.lower() if aggregate_by else None, + aggregateBy=aggregate_by.lower() if aggregate_by else None, ).items() if v is not None }, ) @@ -1871,6 +2489,7 @@ def create_starter_workspace( Returns ------- :class:`StarterWorkspace` + """ payload = { diff --git a/singlestoredb/tests/test_management.py b/singlestoredb/tests/test_management.py index 4b2af1bd3..6abf059d2 100755 --- a/singlestoredb/tests/test_management.py +++ b/singlestoredb/tests/test_management.py @@ -1562,3 +1562,401 @@ def test_str_repr(self): # Test __repr__ assert repr(region) == str(region) + + +@pytest.mark.management +class TestTeams(unittest.TestCase): + """Test cases for teams management.""" + + manager = None + team = None + + @classmethod + def setUpClass(cls): + """Set up the test environment.""" + cls.manager = s2.manage_teams() + + # Create a test team + name = clean_name(f'test-team-{secrets.token_urlsafe(10)}') + cls.team = cls.manager.create_team( + name=name, + description='Test team for unit tests', + ) + + @classmethod + def tearDownClass(cls): + """Clean up the test environment.""" + if cls.team is not None: + try: + cls.team.delete() + except Exception: + pass + cls.manager = None + cls.team = None + + def test_create_team(self): + """Test creating a team.""" + assert self.team is not None + assert self.team.name.startswith('test-team-') + assert self.team.description == 'Test team for unit tests' + assert isinstance(self.team.member_users, list) + assert isinstance(self.team.member_teams, list) + + def test_get_team(self): + """Test getting a team by ID.""" + team = self.manager.get_team(self.team.id) + assert team.id == self.team.id + assert team.name == self.team.name + + def test_list_teams(self): + """Test listing teams.""" + teams = self.manager.list_teams() + team_ids = [t.id for t in teams] + assert self.team.id in team_ids + + def test_update_team(self): + """Test updating a team.""" + new_description = 'Updated test team description' + self.team.update(description=new_description) + + # Verify update + updated_team = self.manager.get_team(self.team.id) + assert updated_team.description == new_description + + def test_str_repr(self): + """Test string representation of team.""" + s = str(self.team) + assert self.team.name in s + assert repr(self.team) == str(self.team) + + def test_no_manager_error(self): + """Test error when no manager is associated.""" + team = self.manager.get_team(self.team.id) + team._manager = None + + with self.assertRaises(s2.ManagementError) as cm: + team.update() + assert 'No teams manager' in cm.exception.msg + + +@pytest.mark.management +class TestPrivateConnections(unittest.TestCase): + """Test cases for private connections management.""" + + manager = None + + @classmethod + def setUpClass(cls): + """Set up the test environment.""" + cls.manager = s2.manage_private_connections() + + @classmethod + def tearDownClass(cls): + """Clean up the test environment.""" + cls.manager = None + + +@pytest.mark.management +class TestAuditLogs(unittest.TestCase): + """Test cases for audit logs management.""" + + manager = None + + @classmethod + def setUpClass(cls): + """Set up the test environment.""" + cls.manager = s2.manage_audit_logs() + + @classmethod + def tearDownClass(cls): + """Clean up the test environment.""" + cls.manager = None + + def test_list_audit_logs(self): + """Test listing audit logs.""" + logs = self.manager.list_audit_logs(limit=10) + # Should return a list (may be empty) + assert isinstance(logs, list) + + +@pytest.mark.management +class TestUsers(unittest.TestCase): + """Test cases for users management.""" + + manager = None + + @classmethod + def setUpClass(cls): + """Set up the test environment.""" + cls.manager = s2.manage_users() + + @classmethod + def tearDownClass(cls): + """Clean up the test environment.""" + cls.manager = None + + def test_get_user(self): + """Test getting a user object.""" + # Create a basic user object (no actual API call since user ID is arbitrary) + user = self.manager.get_user('test-user-123') + assert user.id == 'test-user-123' + assert user._manager is not None + assert user.email == '' # Empty since no actual API call + assert user.first_name == '' + assert user.last_name == '' + + def test_user_from_dict(self): + """Test User.from_dict conversion.""" + from singlestoredb.management.users import User + + data = { + 'userID': 'user-123', + 'email': 'test@example.com', + 'firstName': 'Test', + 'lastName': 'User', + } + + user = User.from_dict(data, self.manager) + assert user.id == 'user-123' + assert user.email == 'test@example.com' + assert user.first_name == 'Test' + assert user.last_name == 'User' + assert user._manager is self.manager + + def test_user_invitation_from_dict(self): + """Test UserInvitation.from_dict conversion.""" + from singlestoredb.management.users import UserInvitation + + data = { + 'invitationID': 'invite-123', + 'email': 'invitee@example.com', + 'state': 'Pending', + 'createdAt': '2023-01-01T00:00:00Z', + 'actedAt': '2023-01-02T00:00:00Z', + 'message': 'Welcome to our team!', + 'teamIDs': ['team-1', 'team-2'], + } + + invitation = UserInvitation.from_dict(data, self.manager) + assert invitation.id == 'invite-123' + assert invitation.email == 'invitee@example.com' + assert invitation.state == 'Pending' + assert invitation.message == 'Welcome to our team!' + assert invitation.team_ids == ['team-1', 'team-2'] + assert invitation._manager is self.manager + + +@pytest.mark.management +class TestWorkspaceManagerIntegration(unittest.TestCase): + """Test cases for workspace manager integration with new modules.""" + + manager = None + workspace_group = None + password = None + + @classmethod + def setUpClass(cls): + """Set up the test environment.""" + cls.manager = s2.manage_workspaces() + + us_regions = [x for x in cls.manager.regions if 'US' in x.name] + cls.password = secrets.token_urlsafe(20) + '-x&$' + + name = clean_name(secrets.token_urlsafe(20)[:20]) + + cls.workspace_group = cls.manager.create_workspace_group( + f'wg-integration-test-{name}', + region=random.choice(us_regions).id, + admin_password=cls.password, + firewall_ranges=['0.0.0.0/0'], + ) + + @classmethod + def tearDownClass(cls): + """Clean up the test environment.""" + if cls.workspace_group is not None: + cls.workspace_group.terminate(force=True) + cls.workspace_group = None + cls.manager = None + cls.password = None + + +@pytest.mark.management +class TestDataClasses(unittest.TestCase): + """Test cases for data classes and object conversion.""" + + def test_team_from_dict(self): + """Test Team.from_dict conversion.""" + from singlestoredb.management.teams import Team, TeamsManager + + manager = TeamsManager() + data = { + 'teamID': 'team-123', + 'name': 'Test Team', + 'description': 'Test Description', + 'memberUsers': [ + { + 'userID': 'user-1', + 'email': 'user1@example.com', + 'firstName': 'User', + 'lastName': 'One', + }, + { + 'userID': 'user-2', + 'email': 'user2@example.com', + 'firstName': 'User', + 'lastName': 'Two', + }, + ], + 'memberTeams': [ + { + 'teamID': 'team-1', + 'name': 'Subteam One', + 'description': 'Sub team description', + }, + ], + 'createdAt': '2023-01-01T00:00:00Z', + } + + team = Team.from_dict(data, manager) + assert team.id == 'team-123' + assert team.name == 'Test Team' + assert team.description == 'Test Description' + assert len(team.member_users) == 2 + assert len(team.member_teams) == 1 + assert team.member_users[0]['userID'] == 'user-1' + assert team.member_teams[0]['teamID'] == 'team-1' + assert team._manager is manager + + def test_private_connection_from_dict(self): + """Test PrivateConnection.from_dict conversion.""" + from singlestoredb.management.private_connections import ( + PrivateConnection, + PrivateConnectionsManager, + ) + + manager = PrivateConnectionsManager() + data = { + 'privateConnectionID': 'conn-123', + 'workspaceGroupID': 'wg-456', + 'serviceName': 'Test Connection', + 'type': 'INBOUND', + 'status': 'ACTIVE', + 'allowList': 'my-allow-list', + 'sqlPort': 3306, + 'websocketsPort': 443, + 'createdAt': '2023-01-01T00:00:00Z', + 'updatedAt': '2023-01-02T00:00:00Z', + } + + conn = PrivateConnection.from_dict(data, manager) + assert conn.id == 'conn-123' + assert conn.workspace_group_id == 'wg-456' + assert conn.service_name == 'Test Connection' + assert conn.type == 'INBOUND' + assert conn.status == 'ACTIVE' + assert conn.allow_list == 'my-allow-list' + assert conn.sql_port == 3306 + assert conn.websockets_port == 443 + assert conn._manager is manager + + def test_audit_log_from_dict(self): + """Test AuditLog.from_dict conversion.""" + from singlestoredb.management.audit_logs import AuditLog + + data = { + 'auditID': 'log-123', + 'createdAt': '2023-01-01T00:00:00Z', + 'userID': 'user-123', + 'userEmail': 'test@example.com', + 'type': 'CREATE_WORKSPACE', + 'reason': 'User created a new workspace', + 'source': 'Portal', + 'userType': 'Customer', + 'orgID': 'org-456', + 'projectID': 'proj-789', + 'workspaceID': 'ws-101', + } + + log = AuditLog.from_dict(data) + assert log.id == 'log-123' + assert log.user_id == 'user-123' + assert log.user_email == 'test@example.com' + assert log.type == 'CREATE_WORKSPACE' + assert log.reason == 'User created a new workspace' + assert log.source == 'Portal' + assert log.user_type == 'Customer' + assert log.organization_id == 'org-456' + assert log.project_id == 'proj-789' + assert log.workspace_id == 'ws-101' + + def test_workspace_group_metrics_from_openmetrics(self): + """Test WorkspaceGroupMetrics.from_openmetrics_text parsing.""" + from singlestoredb.management.metrics import WorkspaceGroupMetrics + + openmetrics_text = ( + '# TYPE singlestoredb_cloud_threads_running gauge' + 'singlestoredb_cloud_threads_running{' + "extractor=\"monitoring-customer-prd/memsql-exporter\"," + "node=\"node-3337afc7-443e-4126-b784-413903527186-aggregator-0\"," + "role=\"CA\"," + "workspace_group_id=\"3337afc7-443e-4126-b784-413903527186\"," + "workspace_name=\"singlestore-central\"} 1" + 'singlestoredb_cloud_cpu_usage{' + "node=\"aggregator-0\",workspace_group_id=\"wg-123\"} 75.5" + ) + + metrics = WorkspaceGroupMetrics.from_openmetrics_text('wg-123', openmetrics_text) + + assert metrics.workspace_group_id == 'wg-123' + assert len(metrics.data_points) == 2 + + # Test first metric + threads_metrics = metrics.get_metrics_by_name( + 'singlestoredb_cloud_threads_running', + ) + assert len(threads_metrics) == 1 + assert threads_metrics[0].value == 1.0 + assert threads_metrics[0].labels['role'] == 'CA' + + # Test second metric + cpu_metrics = metrics.get_metrics_by_name('singlestoredb_cloud_cpu_usage') + assert len(cpu_metrics) == 1 + assert cpu_metrics[0].value == 75.5 + + def test_storage_dr_status_from_dict(self): + """Test StorageDRStatus.from_dict conversion.""" + from singlestoredb.management.storage_dr import StorageDRStatus + + data = { + 'compute': { + 'storageDRType': 'Failover', + 'storageDRState': 'Active', + 'totalWorkspaces': 2, + 'totalAttachments': 5, + 'completedWorkspaces': 1, + 'completedAttachments': 3, + }, + 'storage': [ + { + 'databaseName': 'test_db', + 'region': 'us-east-1', + 'duplicationState': 'Active', + }, + { + 'databaseName': 'prod_db', + 'region': 'us-west-2', + 'duplicationState': 'Pending', + }, + ], + } + + status = StorageDRStatus.from_dict(data) + assert status.compute.storage_dr_type == 'Failover' + assert status.compute.storage_dr_state == 'Active' + assert status.compute.total_workspaces == 2 + assert status.compute.completed_workspaces == 1 + assert len(status.storage) == 2 + assert status.storage[0].database_name == 'test_db' + assert status.storage[0].duplication_state == 'Active' + assert status.storage[1].duplication_state == 'Pending'