diff --git a/.flake8 b/.flake8 index d2f090777..44749c7df 100644 --- a/.flake8 +++ b/.flake8 @@ -12,4 +12,13 @@ per-file-ignores = singlestoredb/fusion/grammar.py:E501 singlestoredb/http/__init__.py:F401 singlestoredb/management/__init__.py:F401 + singlestoredb/management/cluster.py:F401 + singlestoredb/management/export.py:F401 + singlestoredb/management/project.py:F401 + singlestoredb/management/workspace.py:F401 + # The v1/ and v2/ modules are version namespaces: they re-export the + # shared implementations under the names the manage_* factories look up, + # so unused-import is expected there. + singlestoredb/management/v1/*.py:F401 + singlestoredb/management/v2/*.py:F401 singlestoredb/mysql/__init__.py:F401 diff --git a/.github/workflows/code-check.yml b/.github/workflows/code-check.yml index 526dc61d4..5ef487bc7 100644 --- a/.github/workflows/code-check.yml +++ b/.github/workflows/code-check.yml @@ -171,8 +171,13 @@ jobs: SINGLESTOREDB_FUSION_ENABLE_HIDDEN: "1" - name: Run HTTP protocol tests + # -n 0 overrides the -n 3 in pyproject.toml's addopts: the HTTP/Data API + # run must be serial. Setup goes over SINGLESTOREDB_INIT_DB_URL (MySQL), + # so load_sql takes its `SET GLOBAL HTTP_PROXY_PORT` + `RESTART PROXY` + # branch (singlestoredb/tests/utils.py:227) once per worker, and a proxy + # restart drops any HTTP request another worker has in flight. run: | - pytest -v -m 'not management' --cov=singlestoredb --pyargs singlestoredb.tests + pytest -v -n 0 -m 'not management' --cov=singlestoredb --pyargs singlestoredb.tests env: COVERAGE_FILE: "coverage-http.cov" SINGLESTOREDB_URL: "http://root:root@127.0.0.1:9081" diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 9b388cf62..6c9546fa9 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -58,8 +58,13 @@ jobs: SINGLESTOREDB_FUSION_ENABLE_HIDDEN: "1" - name: Run HTTP protocol tests + # -n 0 overrides the -n 3 in pyproject.toml's addopts: the HTTP/Data API + # run must be serial. Setup goes over SINGLESTOREDB_INIT_DB_URL (MySQL), + # so load_sql takes its `SET GLOBAL HTTP_PROXY_PORT` + `RESTART PROXY` + # branch (singlestoredb/tests/utils.py:227) once per worker, and a proxy + # restart drops any HTTP request another worker has in flight. run: | - pytest -v -m 'not management' --cov=singlestoredb --pyargs singlestoredb.tests + pytest -v -n 0 -m 'not management' --cov=singlestoredb --pyargs singlestoredb.tests env: COVERAGE_FILE: "coverage-http.cov" SINGLESTOREDB_URL: "http://root:root@127.0.0.1:9081" diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 5712c3ef5..d3a669c1e 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -114,7 +114,10 @@ jobs: CIBW_BUILD: "cp39-*" CIBW_SKIP: "pp* *-musllinux* *-manylinux_i686" CIBW_TEST_COMMAND: "pytest -v --pyargs singlestoredb.tests.test_basics" - CIBW_TEST_REQUIRES: "pytest" + # xdist because pyproject.toml's addopts sets -n/--dist, and PYTHONPATH + # points --pyargs at the workspace, so that pyproject is the inifile + # here. Without the plugin pytest exits on the unknown arguments. + CIBW_TEST_REQUIRES: "pytest pytest-xdist" CIBW_ENVIRONMENT: "SINGLESTOREDB_URL='mysql://${{ secrets.CLUSTER_USER }}:${{ secrets.CLUSTER_PASSWORD }}@${{ needs.setup-database.outputs.cluster-host }}:3306/${{ needs.setup-database.outputs.cluster-database }}?pure_python=0'" PYTHONPATH: ${{ github.workspace }} diff --git a/.github/workflows/smoke-test.yml b/.github/workflows/smoke-test.yml index 27c2a9282..688a2dc18 100644 --- a/.github/workflows/smoke-test.yml +++ b/.github/workflows/smoke-test.yml @@ -122,7 +122,13 @@ jobs: - name: Run tests if: ${{ matrix.driver == 'https' }} - run: pytest -v --pyargs singlestoredb.tests.test_basics + # -n 0 overrides the -n 3 in pyproject.toml's addopts: the Data API is + # not run in parallel. This job avoids the `RESTART PROXY` hazard the + # code-check/coverage HTTP steps hit -- no SINGLESTOREDB_INIT_DB_URL + # here, so load_sql's setup connection is itself HTTP and skips that + # branch -- but the driver is the same one, so it gets the same + # treatment rather than being the lone parallel Data API run. + run: pytest -v -n 0 --pyargs singlestoredb.tests.test_basics env: PYTHONPATH: ${{ github.workspace }} SINGLESTOREDB_URL: "${{ matrix.driver }}://${{ secrets.CLUSTER_USER }}:${{ secrets.CLUSTER_PASSWORD }}@${{ needs.setup-database.outputs.cluster-host }}:443/${{ needs.setup-database.outputs.cluster-database }}?pure_python=${{ matrix.pure-python }}&buffered=${{ matrix.buffered }}" diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 80ebe2469..8720daa7a 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -26,7 +26,7 @@ the components fit together, their responsibilities, and their interactions. The SingleStoreDB Python SDK provides: - **DB-API 2.0 compliant interface** to SingleStore databases -- **Cloud management API** for workspace and cluster lifecycle management +- **Cloud management API** for cluster lifecycle management - **Fusion SQL** for client-side SQL command extension - **External Functions** (UDFs) for deploying Python functions to SingleStore - **AI integrations** for chat and embeddings @@ -70,20 +70,26 @@ singlestoredb/ │ ├── management/ # Cloud management API │ ├── manager.py # Base REST client -│ ├── workspace.py # Workspace/WorkspaceGroup/Stage +│ ├── cluster.py # Cluster/StarterCluster +│ ├── project.py # Project definitions +│ ├── stage.py # Stage file storage │ ├── organization.py # Organization management │ ├── region.py # Region definitions │ ├── job.py # Job management │ ├── files.py # File operations │ ├── billing_usage.py # Billing and usage tracking -│ └── export.py # Data export operations +│ ├── export.py # Data export operations +│ ├── workspace.py # v1 re-exports (deprecated) +│ ├── v1/ # Version 1 routes (deprecated) +│ └── v2/ # Version 2 routes (default) │ ├── fusion/ # Client-side SQL extensions │ ├── handler.py # SQLHandler base class │ ├── registry.py # Handler registration │ ├── result.py # FusionSQLResult │ └── handlers/ # Built-in handlers -│ ├── workspace.py # Workspace commands +│ ├── cluster.py # Cluster commands +│ ├── workspace.py # Workspace commands (deprecated) │ ├── stage.py # Stage commands │ ├── job.py # Job commands │ ├── files.py # File commands @@ -128,7 +134,7 @@ layer provides a unified interface with protocol-specific implementations. ### Connection Architecture -The entry point is `singlestoredb.connect()` in `singlestoredb/connection.py:1312`: +The entry point is `singlestoredb.connect()` in `singlestoredb/connection.py:1354`: ```python import singlestoredb as s2 @@ -252,7 +258,7 @@ conn.show.plan(plan_id) # SHOW PLAN - execution plan details **Fusion SQL Integration:** - Client-side interception of extended SQL commands -- Workspace management via SQL syntax +- Cluster management via SQL syntax - Stage (file storage) operations via SQL **Multiple Result Formats:** @@ -431,6 +437,18 @@ export SINGLESTOREDB_FUSION_ENABLED=1 The management API (`singlestoredb/management/`) provides programmatic access to SingleStore's cloud management features. +The API is versioned. Version-neutral code lives in the top-level modules, whose +base classes implement version 2 — the default, set once in +`singlestoredb/_management_version.py`. `management/v1/` holds the version 1 +overrides and `management/v2/` pins the version 2 routes. Version is selected by +the `management.version` option (`SINGLESTOREDB_MANAGEMENT_VERSION`) or a +`version=` argument to any `manage_*` function. + +Version 2 replaced version 1's workspace groups and workspaces with a single +flat `Cluster` resource, so `WorkspaceManager`, `WorkspaceGroup` and `Workspace` +are deprecated in favor of `ClusterManager` and `Cluster`. Grouping is expressed +by `Project`, an organizational unit rather than a deployment parent. + ### Architecture ``` @@ -448,12 +466,12 @@ SingleStore's cloud management features. │ ▼ ┌─────────────────────────────────────────────────────────────────────┐ -│ WorkspaceManager │ -│ singlestoredb/management/workspace.py │ +│ ClusterManager │ +│ singlestoredb/management/cluster.py │ ├─────────────────────────────────────────────────────────────────────┤ -│ workspace_groups() regions() organizations() │ -│ create_workspace() get_workspace() billing() │ -│ starter_workspaces() create_workspace_group() │ +│ clusters regions organizations │ +│ create_cluster() get_cluster() billing │ +│ starter_clusters projects get_project() │ └─────────────────────────────────────────────────────────────────────┘ ``` @@ -463,24 +481,24 @@ SingleStore's cloud management features. import singlestoredb as s2 # Initialize manager with API token -mgr = s2.manage_workspaces() +mgr = s2.manage_clusters() -# List workspace groups -for wg in mgr.workspace_groups(): - print(wg.name, wg.id) +# List clusters +for c in mgr.clusters: + print(c.name, c.id) -# Create a workspace -ws = mgr.create_workspace( - name='my-workspace', - workspace_group=wg, +# Create a cluster +c = mgr.create_cluster( + name='my-cluster', + region='US West 2 (Oregon)', size='S-00', ) -# Connect to workspace -conn = ws.connect() +# Connect to the cluster +conn = c.connect() # Stage operations (file storage) -stage = wg.stage +stage = c.stage stage.upload_file('local.csv', '/data/uploaded.csv') stage.download_file('/data/uploaded.csv', 'downloaded.csv') stage.listdir('/data') @@ -491,13 +509,17 @@ stage.listdir('/data') | Class | File | Purpose | |-------|------|---------| | `Manager` | `manager.py` | Base REST client with auth | -| `WorkspaceManager` | `workspace.py` | Main management interface | -| `WorkspaceGroup` | `workspace.py` | Group of workspaces | -| `Workspace` | `workspace.py` | Database instance | -| `Stage` | `workspace.py` | File storage operations | -| `StarterWorkspace` | `workspace.py` | Free tier workspace | +| `ClusterManager` | `cluster.py` | Main management interface | +| `Cluster` | `cluster.py` | Database deployment | +| `StarterCluster` | `cluster.py` | Shared-tier deployment | +| `Project` | `project.py` | Grouping for an org's clusters | +| `Stage` | `stage.py` | File storage operations | | `Organization` | `organization.py` | Organization management | -| `Billing` | `workspace.py` | Usage and billing | +| `Billing` | `billing.py` | Usage and billing | +| `WorkspaceManager` | `v1/workspace.py` | v1 interface (deprecated) | +| `WorkspaceGroup` | `v1/workspace.py` | v1 group of workspaces (deprecated) | +| `Workspace` | `v1/workspace.py` | v1 database instance (deprecated) | +| `StarterWorkspace` | `v1/workspace.py` | v1 shared tier (deprecated) | --- @@ -598,7 +620,8 @@ Located in `singlestoredb/fusion/handlers/`: | Handler | Commands | |---------|----------| -| `workspace.py` | `SHOW WORKSPACE GROUPS`, `CREATE WORKSPACE`, etc. | +| `cluster.py` | `SHOW CLUSTERS`, `CREATE CLUSTER`, `SHOW PROJECTS`, etc. | +| `workspace.py` | `SHOW WORKSPACE GROUPS`, `CREATE WORKSPACE`, etc. (deprecated) | | `stage.py` | `UPLOAD`, `DOWNLOAD`, `CREATE STAGE FOLDER` | | `job.py` | `SHOW JOBS`, `CREATE JOB`, `DROP JOB` | | `files.py` | File management commands | @@ -1131,12 +1154,13 @@ Feature options: | Purpose | Primary File | |---------|-------------| | Entry point | `singlestoredb/__init__.py` | -| Connect function | `singlestoredb/connection.py:1312` | +| Connect function | `singlestoredb/connection.py:1354` | | MySQL connection | `singlestoredb/mysql/connection.py` | | Cursor types | `singlestoredb/mysql/cursors.py` | | HTTP connection | `singlestoredb/http/connection.py` | | Configuration | `singlestoredb/config.py` | -| Management API | `singlestoredb/management/workspace.py` | +| Management API | `singlestoredb/management/cluster.py` | +| Management API version default | `singlestoredb/_management_version.py` | | Fusion handlers | `singlestoredb/fusion/handler.py` | | UDF decorator | `singlestoredb/functions/decorator.py` | | Plugin UDF server | `singlestoredb/functions/ext/plugin/server.py` | diff --git a/README.md b/README.md index eacdd6b02..c2b617867 100644 --- a/README.md +++ b/README.md @@ -13,8 +13,8 @@ analytics and vector search. (port 9000) using the same interface - **Flexible Result Formats**: Return query results as tuples, dictionaries, named tuples, NumPy arrays, Pandas DataFrames, Polars DataFrames, or PyArrow Tables -- **Workspace Management**: Full API for managing SingleStore Cloud workspaces, - clusters, regions, and files programmatically +- **Deployment Management**: Full API for managing SingleStore Cloud clusters, + projects, regions, and files programmatically - **Vector Store**: Pinecone-compatible vector database API for similarity search applications with built-in connection pooling - **User-Defined Functions**: Deploy Python functions as SingleStore UDFs with @@ -173,29 +173,37 @@ df = cur.fetchone() ## Management API -The SDK provides a workspace management API for managing SingleStore deployments -programmatically. This includes creating and managing workspaces, clusters, +The SDK provides a management API for managing SingleStore deployments +programmatically. This includes creating and managing clusters, projects, regions, and files. ```python import singlestoredb as s2 -# Get a workspace manager (uses SINGLESTOREDB_MANAGEMENT_TOKEN env var by default) -manager = s2.manage_workspaces() +# Get a cluster manager (uses SINGLESTOREDB_MANAGEMENT_TOKEN env var by default) +manager = s2.manage_clusters() -# List all workspaces -for ws in manager.workspaces: - print(ws.name, ws.state) +# List all clusters +for c in manager.clusters: + print(c.name, c.state) -# Create a new workspace -ws = manager.workspaces.create( - name='my-workspace', - workspace_group_id='', +# Create a new cluster +c = manager.create_cluster( + name='my-cluster', + region='US West 2 (Oregon)', + size='S-00', ) ``` +The API is versioned, and version 2 — the flat `Cluster` resource shown above — +is the default. Version 1, which called a deployment a `Workspace` inside a +`WorkspaceGroup` and is reached through `s2.manage_workspaces()`, still works +but is deprecated in its entirety. Select a version with the +`management.version` option (`SINGLESTOREDB_MANAGEMENT_VERSION`) or by passing +`version=` to any `manage_*` function. + See the [API documentation](https://singlestore-labs.github.io/singlestoredb-python) -for full details on workspace, cluster, region, and file management. +for full details on cluster, project, region, and file management. ## Vector Store @@ -225,7 +233,7 @@ Pinecone-compatible operations for vector similarity search. Fusion SQL extends the SQL commands handled by the client with custom handlers. These commands are processed locally rather than sent to the database server. Built-in handlers provide SQL-like commands for managing -workspaces, running notebook jobs, and more. +clusters, running notebook jobs, and more. ```python import os @@ -235,23 +243,24 @@ import singlestoredb as s2 conn = s2.connect() # Show available cloud regions -conn.execute('SHOW REGIONS') - -# List workspace groups -conn.execute('SHOW WORKSPACE GROUPS') +conn.execute('SHOW CLUSTER REGIONS') -# List workspaces in a specific group -conn.execute("SHOW WORKSPACES IN GROUP 'my-group' EXTENDED") +# List clusters +conn.execute('SHOW CLUSTERS EXTENDED') -# Create a new workspace group +# Create a new cluster conn.execute(""" - CREATE WORKSPACE GROUP 'analytics-team' + CREATE CLUSTER 'analytics-team' IN REGION 'US West 2 (Oregon)' - WITH PASSWORD 'my-password' + WITH SIZE 'S-00' WITH FIREWALL RANGES '10.0.0.0/8' """) ``` +The `WORKSPACE` and `WORKSPACE GROUP` commands, and the version-less +`SHOW REGIONS`, still work but are deprecated along with the rest of management +API v1. + See [singlestoredb/fusion/README.md](singlestoredb/fusion/README.md) for details on writing custom Fusion SQL handlers. diff --git a/docs/adr/0001-versioned-management-api-wrappers.md b/docs/adr/0001-versioned-management-api-wrappers.md new file mode 100644 index 000000000..b64882e5d --- /dev/null +++ b/docs/adr/0001-versioned-management-api-wrappers.md @@ -0,0 +1,134 @@ +# ADR 0001: Versioned Management API Wrappers + +## Status + +Accepted. Amended — the original decision included a cross-version bridge +(`VersionedMixin`, per-entity `_response` storage, manager clones) that has +since been removed. See [Revisions](#revisions) for what changed and why. + +## Context + +The Management API has multiple versions (v1, v2, etc.) with differing endpoints and response shapes. A `Manager` instance is locked to one version via its `_base_url`, and all entities created through that manager use that version. + +v2 is not an additive revision of v1. Workspace groups and workspaces were replaced by a single flat `Cluster` resource, Stage moved from a top-level resource to one nested under the cluster, regions lost their IDs, and the `inferenceapis/` routes disappeared entirely. So "v2 is v1 plus overrides" is not true in general, and v1 is expected to be abandoned outright rather than maintained alongside v2. + +We needed a way to: +- Serve both versions from one package while they overlap +- Let each version differ in behavior without duplicating what they share +- Keep backward compatibility with existing import paths and usage patterns +- Make retiring v1 a deletion rather than an excavation + +## Decision + +### Folder structure + +Version-specific modules live in `management/v1/`, `management/v2/`, etc. Version-neutral implementations live in the top-level `management/` modules (`manager.py`, `stage.py`, `job.py`, `organization.py`, `region.py`, `files.py`, `utils.py`). + +A module belongs at the top level when both versions share the implementation and the only difference is the URL the shared code is pointed at. Anything else — a different resource model, a different request or response shape, a route that exists at one version only — belongs in a version folder. + +Each version folder is a **complete set** for the resources that version has: every class reachable at that version must be importable from its folder, whether as a real subclass or a re-export of the shared implementation. There is no cross-version fallback; requesting a class from a version where it doesn't exist raises an error. + +**Rule 1: no cross-version imports, in either direction.** `management/v1/` must not import from `management/v2/` and vice versa. Shared code moves up to `management/`; it never travels sideways. This is what makes retiring a version an `rm -rf` of its folder plus removal of the back-compat shims, and it is enforced by `TestVersionPackagesAreIndependent` in `singlestoredb/tests/test_management_versioning.py` — an AST walk over each folder's imports plus a `sys.meta_path` blocker that imports every module of one version with the other forbidden. + +Top-level modules also serve as thin re-export shims for stable import paths (`from singlestoredb.management.workspace import Workspace` still resolves to the v1 class). Version routing happens in the top-level functions only — duplicating one into a version folder both invites the copies to drift and makes the folder un-deletable. + +**Rule 2: everything exported from `singlestoredb.management` is version-neutral.** A caller who does not name a version gets the version the `management.version` option names; an explicit `version=` argument always wins. That applies to the `manage_*()` factories and equally to the module-level helpers (`get_organization`, `get_secret`, `get_stage`), which were the v1 implementations under a neutral name until they were routed through `_versioned_attr()`. A resolved version that lacks the resource raises and names the replacement — `manage_clusters()` at v1 points at workspaces, `manage_workspaces()` at v2 points at clusters — rather than silently answering from the version that happens to have it. `manage_workspaces()` is the one exception to the option's reach, for the reason given under [API version in URL](#api-version-in-url). + +Two consequences worth stating: + +- The one place the resolution rule lives is `_resolve_version()` in `management/_version_import.py`. Nothing else reads the option. +- Callers that are v1-only *by design* rather than by default — Fusion, the UDF `stage://` handling, the AI inference helpers — go through the private `_manage_workspaces_v1()`, which ignores the option. They are asking for a workspace manager specifically, so an org-wide preference for another version has nothing to say to them. The same reasoning applies to test suites: each version's suite pins its own `version=`, so the ambient option cannot change what is under test. + +Because the module that implements a helper differs by version — v1 hangs them off workspaces, v2 off clusters — the neutral layer looks them up by name in the resolved version *package* (`_versioned_attr('get_stage', ver)`) rather than hard-coding a module per version. Each version package re-exports its own from `__init__.py`, so adding a version is an export list, not a branch in the dispatcher. + +### Inheritance model + +The shared base class carries the **newest** version's behavior. Older versions subclass it and override backward. So: + +```python +# management/stage.py -- the shared base is level-set to v2 +class Stage(FileLocation): + def _fs_path(self, path=''): + return f'clusters/{self._deployment_id}/stage/fs/{path}' + +# v1/stage.py -- the backward override +class Stage(_Stage): + def _fs_path(self, path=''): + return f'stage/{self._deployment_id}/fs/{path}' +``` + +and v2 uses the shared class unchanged — there is no `v2/stage.py`; `v2`'s +`Stage` is re-exported from `v2/cluster.py`, which imports it from +`management/stage.py`. The direction matters: with v2 as the subclass, deleting `v1/` would strand the base class it inherits from. With v1 as the subclass, deleting `v1/` leaves the current behavior standing on its own. + +Version differences are expressed as **class attributes on the shared class**, repointed by the version subclass, rather than as runtime `if version == ...` branches: + +- `JobsManager._deployment_target_type`, `_starter_target_type` — the `targetType` strings each version uses +- `Organization._jobs_manager_class`, `_inference_api_manager_class` +- `Organizations._organization_class` — so a v1 manager hands out a v1-configured organization +- `Stage._fs_path` — the one thing that differs about Stage + +A resource that exists at one version only lives in that version's folder, and the shared base raises a `ManagementError` explaining the absence if the operation has no equivalent. `inference_api.py` is v1-only for this reason, and `Organization.inference_apis` raises from the shared base for every version past v1. + +The inverse mistake is just as easy to make: an operation that looks version-specific but is not. `RegionManager.list_shared_tier_regions` was written as a v1-only override on the strength of a `GET /v2/regions/sharedtier` 404 that turned out not to happen — the route answers identically at both versions, so it now lives in the shared base and `v1/region.py` is a pure re-export. Confirm the absence against the live API before encoding it, because the OpenAPI dump does not describe v2. + +### Convention-based module lookup + +`_import_versioned_module(version, module_name)` in `management/_version_import.py` imports `singlestoredb.management.{version}.{module_name}`, distinguishing "this version is unsupported" from "this version has no such module" in its error message. The `manage_*()` factories are its only callers. Its companion `_versioned_attr(name, version)` looks a name up in the version *package* instead, for the helpers whose implementing module differs by version. No registry or registration is needed — the folder structure is the registry. + +### API version in URL + +Each manager class has a `default_version` class attribute, and the URL is built as `urljoin(base_url_root, version or type(self).default_version) + '/'`. A class that implements one specific version's routes names that version as a literal — `v1/workspace.py` pins `'v1'`, `v2/cluster.py` pins `'v2'` — so it keeps addressing its own routes when a newer version becomes the default. A version-neutral class takes `DEFAULT_VERSION` from `_version_import`, which is `singlestoredb._management_version.DEFAULT_MANAGEMENT_VERSION`. That constant is the single place the *current* version is named: it supplies the registered default of the `management.version` option as well, so the option and the classes cannot drift apart, and retargeting the SDK at a new version is a one-line change. `_management_version.py` imports nothing, which is what lets `config.py` and `management/` both read it — `config.py` is imported before `management`, and `management.manager` imports `config`, so the constant cannot live in either. + +Version numbers still appear as literals where the point *is* a specific version rather than the current one: the `default_version` of a version-specific class, and the guards in `manage_workspaces()` (v1-only resource) and `manage_clusters()` (absent at v1). Those are facts about v1 and v2 that a v3 must not silently change. + +`default_version` is **not** resolved from `config.get_option('management.version')`: doing that let a v1-only class declare itself to be v2 whenever the option was set. `config.get_default()` is no better — `Option.__init__` folds the environment variable into the registered default, so a class reading it would take a v1 URL from `SINGLESTOREDB_MANAGEMENT_VERSION=v1`. Both holes are guarded by tests in `test_management_versioning.py`, the second from a subprocess, since it only appears at import with the variable already set. + +The `management.version` option is consulted by the version-neutral entry points, never by the manager classes. `manage_clusters()` consults it even though clusters exist at a single version, resolving first and then raising when the resolved version has no clusters — so an option naming v1 is answered as the deliberate request it is. `manage_workspaces()` is the exception, along with the private `_manage_workspaces_v1()` behind it: it is pinned to v1 and never reads the option, because workspaces exist only there, and resolving would turn a bare call into an exception once the default moved past v1. See the rule 2 note under [Alternatives](#v2-subclasses-v1). + +### Deprecation of the v1 grammar + +`manage_workspaces()` and the workspace-group vocabulary are deprecated in favor of `manage_clusters()`. The deprecation warning lives in `manage_workspaces()`; the un-warned body is `_manage_workspaces_v1()`. Internal callers that are v1-only by design — Fusion handlers, the UDF `stage://` handling, the AI helpers — call the private form, so they do not emit a warning the caller can do nothing about. + +## Alternatives Considered + +### Single manager with version parameter per method call + +Rejected: would pollute every method signature and make it unclear which version's response schema applies to the returned entity. + +### v2 subclasses v1 + +Rejected: it inverts the dependency relative to the lifecycle. v1 is the version that goes away, so it must be the leaf. It also does not describe v2 honestly — a `Cluster` is not a `Workspace` with overrides. + +### Separate, unrelated manager classes per version + +Rejected: the versions genuinely share most of their surface (files, jobs, secrets, billing, the HTTP plumbing), and duplicating it would let the copies drift. Sharing a level-set base with backward overrides in `v1/` keeps one implementation of the common part without making either version depend on the other. + +### Runtime `if version == 'v1'` branches in shared code + +Rejected: it spreads version knowledge across every method that has any, and the branches survive the deletion of `v1/` as dead code that still reads as live. A class attribute puts the difference in one declaration, at the version that owns it. + +### Fallback to v1 if a class doesn't exist in v2 + +Rejected: silent fallback hides bugs. If you ask for a v2 class and it doesn't exist, that's an error worth surfacing. + +## Consequences + +- Adding a new API version means creating a folder, moving the newest behavior into the shared base, and leaving a backward override in the now-older folder +- Retiring a version means deleting its folder and the shims that re-export from it; nothing else refers to it +- Import paths are stable — existing code using `from singlestoredb.management.workspace import Workspace` continues to work unchanged +- Version differences are declarations rather than control flow, so "what differs at v1?" is answerable by reading `v1/` +- Entities carry no stored API response, so an object cannot be re-interpreted as another version after the fact; getting a different version's view means asking that version's manager + +## Revisions + +The accepted decision originally included a cross-version bridge, removed in +full on the `versioned-management-api` branch: + +- **`VersionedMixin` and `.v1`/`.v2` attribute switching.** A `__getattr__` intercepting `v\d+` let any manager or entity hop versions in place, returning a cached clone or a re-parsed entity. Removed: callers reach a version through the factory they call, and the bridge required exactly the cross-version coupling that rule 1 forbids. Nothing consumed it outside its own tests. +- **`_response` storage on every entity.** Entities stashed their raw API response so another version's `from_dict` could re-read it. Removed with the bridge — with v1 and v2 modeling different resources, re-parsing one version's payload as another was not meaningful anyway. +- **Clone-support state on `Manager`** (`_access_token`, `_base_url_root`, `_organization_id`) and the v1↔v2 field translators (`v1/_translate.py`, `v1/cluster.py`) existed only to feed the bridge, and went with it. +- **Inheritance direction inverted** from "v2 subclasses v1" to "shared base level-set to the newest version, `v1/` holds backward overrides", for the reasons in the alternatives above. +- **`default_version` resolved from the config option.** The original text described it as resolved from `config.get_option('management.version')`; making it dynamic was the bug that let a v1 class report itself as v2. A version-specific class now pins a literal, and a version-neutral one takes the shared `DEFAULT_MANAGEMENT_VERSION` constant, which no runtime setting can move. +- **`management/versioned.py` renamed to `_version_import.py`**, since all that remains of it is the version-module importer. +- **Rule 2 added.** The original text only described version routing in the `manage_*()` factories, which left `singlestoredb.management.get_organization`/`get_secret`/`get_stage` re-exported straight from `v1/`: neutral names that ignored the option and would vanish with the v1 package. They now dispatch on the resolved version. `manage_workspaces()`, however, stays **pinned to v1** along with its private `_manage_workspaces_v1()`: workspaces exist only at v1, so the option has nothing to select between, and letting it resolve would mean a bare `manage_workspaces()` raises once the option defaults to v2 — v1 ceasing to work rather than v1 being deprecated. It emits a `DeprecationWarning` pointing at `manage_clusters()` and returns a working v1 manager. An explicit `version='v2'` still raises. diff --git a/docs/management-api-audit.md b/docs/management-api-audit.md index c265debbb..3b36030f4 100644 --- a/docs/management-api-audit.md +++ b/docs/management-api-audit.md @@ -438,6 +438,290 @@ These are not in the scope of this audit pass but are worth noting: 3. **`fields=` query param** on every GET — intentionally skipped per scope. 4. **DR / identity / privateConnections / delegatedEntities sub-resources** on workspace groups — intentionally skipped per scope. +5. **`GET /projects` was missing from the spec dump, and `projectID` is required + on `POST /v2/clusters`.** Both confirmed live against + `https://api.singlestore.com` (2026-08-21): + + - `GET /v1/projects` and `GET /v2/projects` both return + `[{projectID, name, edition, createdAt}]`, with `edition` one of + `SHARED | STANDARD | ENTERPRISE`. Neither route appeared anywhere in the + 1.1.124 `dev-docs/management_api.openapi` snapshot this audit was written + against — one instance of that dump not being authoritative. The refreshed + dump (1.2.171, 2026-08-25) does document `/v2/projects`; `/v1/projects` is + still unpublished. + - `POST /v2/clusters` fails with `400 projectID is required` for any body + without it, including one that is otherwise complete. `POST + /v1/workspaceGroups` assigns a project implicitly: every group in the test + organization sits in `Standard Project` without the SDK ever sending an ID. + Handled by `ClusterManager._resolve_project_id`, which takes the caller's + `project`, then the project of the deployment the code is running in, then + the organization's only project, and otherwise raises naming the + candidates. The argument accepts a project *name* as well as an ID: + `_project_id_for` treats a UUID as an ID and anything else as a name to + look up, which is safe because the route answers `400 uuid: incorrect UUID + length` for a non-UUID ID. The API does not promise names are unique, so an + ambiguous name raises rather than resolving to the first match. + + **⚠ Correction (established while testing the notebooks).** Priority two + was originally `SINGLESTOREDB_PROJECT`. That variable names a project of + the *inference* API, not of this one — a notebook reports an ID there that + `GET /v2/projects/{id}` answers `404 project not found` for — so reading it + broke `CREATE CLUSTER` in every notebook. Reading the project off the + current deployment replaces it and is a better default anyway: a new + cluster lands beside the one it was created from. + - `POST /v2/sharedtier/virtualClusters` does **not** require `projectID` — + validation runs through to `databaseName` without it — so + `create_starter_cluster` resolves `project` only when one is given. + - Field-validation order on `POST /v2/clusters` is `region` → `projectID` → + `firewallRanges` (which must be present, `[]` to disallow all inbound + traffic) → `sizeConfig`. + **Correction (2026-08-28): the size field is `sizeConfig`, not `size`.** + It was `size` when this was recorded. The rename shipped on 2026-08-26, + was backed out the next morning (verified live 2026-08-27: `sizeConfig` + drew `400 ... unknown field "sizeConfig"`), and landed again by + 2026-08-28, when `size` began drawing `400 request body contains an + unknown field "size"`. Only the outer key changed; the object inside is + still `{size, scaleFactor}`, and `PATCH /v2/clusters/{id}` moved with it. +6. **`POST /v2/sharedtier/virtualClusters` accepts only `AWS` | `AZURE` | `GCP` + verbatim.** Also confirmed live (2026-08-21). Any other capitalization — + including the mixed-case `Azure` that `GET /v2/regions` itself reports — + fails with `500 Unspecified is not a valid CloudServiceProvider`, so a + region's `provider` cannot be passed through as-is. `POST /v2/clusters` is + case-insensitive on the same field (only an unknown provider is rejected, + with `400 invalid provider ...; value must be aws or azure or gcp`). + `create_starter_cluster` upper-cases it; `create_cluster` does not. + The v1 starter route is a different path, and the v1 shared-tier region list + reports only `AWS us-east-1`, so v1 never hit this. + **Correction (verified live 2026-08-24): `GET /v2/regions/sharedtier` is + *not* missing.** An earlier pass of this audit recorded it as a 404 and a + "real gap"; that was wrong. The route returns **200** with + `[{"region": "US East 1 (N. Virginia)", "provider": "AWS", + "regionName": "us-east-1"}]` — the same shape and the same content as at v1. + So shared-tier regions *are* discoverable from v2 and a v2-only client need + not hard-code them. `RegionManager.list_shared_tier_regions` and + `ClusterManager.shared_tier_regions` both implement it, and neither raises. + The 36 entries `GET /v2/regions` returns still carry only `region`, + `provider`, `regionName` with nothing marking shared-tier capability, so the + `sharedtier` route remains the only way to tell: sending a region absent + from it fails at create time with `500 error creating virtual workspace + (): no shared tier region found for provider AWS and region + us-east-2`. +7. **v2 deployment name format.** `POST /v2/clusters` requires the name to match + `[a-z0-9]([a-z0-9-]*[a-z0-9])?` at 1-32 characters: an uppercase letter, an + underscore, a dot, a space, or a leading/trailing hyphen draws `400 name: + must be in a valid format`, and anything longer draws `400 name: the length + must be between 1 and 32`. Repeated hyphens are accepted. + `POST /v2/sharedtier/virtualClusters` applies none of this — it took + `STARTER_cl_test_abc-` unchanged. Neither rule is in the spec dump. + Full validation order on `POST /v2/clusters`: `region` presence → + `projectID` → `firewallRanges` → `sizeConfig` → `name` → region existence. +8. **`POST /v2/clusters` ignores `adminPassword` and generates its own.** + Confirmed live (2026-08-21) with two throwaway clusters, both since + terminated. Whatever password is sent, the created cluster's `admin` user + gets a server-generated one, returned as `adminPassword` in the *create + response only* — `GET /v2/clusters/{id}` has no such field. Losing that + value means losing `admin` access to the cluster. v1 honored the password it + was given, so nothing in the v1 wrapper had to keep it. `create_cluster` + therefore carries it onto the returned object as + `Cluster.admin_password` (backed by a private attribute so it stays out of + `str()`/`repr()`), and the `admin_password` parameter's docstring carries a + warning that v2 discards it. Not in the spec dump. + + **Re-confirmed 2026-08-25 with a connection attempt**, which the original + probe had not made — the earlier evidence was only that the create response + echoed a value different from the one posted, which on its own does not + establish which of the two authenticates. One throwaway `S-00` + (`probe-adminpw-1787666027`, since terminated) created with + `adminPassword: 'Probe-Sent-Pw-2026a!'` returned the generated + `'{:D}TK*[F3Ll}Ups2pNv'`. Connecting as `admin` over the MySQL protocol with + the value we *sent* was refused with `1045: Access denied for user 'admin'`; + the same connection with the *returned* value succeeded. The generated + password is therefore the real credential, and the sent one is discarded + rather than merely unreported. The reference documents the field + (`docs.singlestore.com/cloud/reference/management-api/reference/`), so this + is a server-behavior bug and not a missing parameter — worth raising with + the API team on that basis. + + **The documented invalid-password fallback does not explain it.** The + reference states the password must be ≥14 characters with an uppercase, a + lowercase, a numeric and a special character, at most two consecutive + sequential characters and at most three consecutive identical characters, and + that "if a password is not specified **or if an invalid password is + provided**, a valid password is generated and returned in the response + object." That fallback is silent and therefore indistinguishable from the + field being ignored, so the probe was re-run 2026-08-25 on two password + shapes at once, on two throwaway clusters (`probe-pw-random-1787666958`, + `probe-pw-original-1787666958`, both since terminated): + + | | `Xq7#vTm2$pLw9Kz@` | `Probe-Sent-Pw-2026a!` | + |---|---|---| + | POST echoed what we sent | no | no | + | sent value authenticates | **no** (1045) | **no** (1045) | + | returned value authenticates | yes | yes | + + The first is random, 16 characters, all four character classes, no sequential + run, no repeated character and no dictionary word — it trips no documented + rule, nor the guessable undocumented ones (a dictionary check on + `Probe`/`Sent`, or `-` not counting as special). It was discarded identically. + The generated replacements came back 25-26 characters. So the field is inert + on this route rather than rejecting non-compliant input. Not ruled out: that + the generate-always behavior is specific to this organization, its tier, or + `aws/us-east-1` — every probe has run there. + + Incidental findings from these probes, none in the spec dump: + `POST /v2/clusters` rejects a null `firewallRanges` with + `400 firewallRanges cannot be null (indicate empty list [] to disallow all + inbound traffic)` **even when `allowAllTraffic: true` is sent** — the field + is unconditionally required, so `allow_all_traffic` alone is not a usable + way to open a new cluster. And the create error text calls the resource a + workspace (`error creating workspace (): ...`) despite the v2 cluster + vocabulary, including in the on-demand quota rejection. +9. **`PATCH /v2/clusters/{id}` accepts `name` and silently ignores it.** + Confirmed live (2026-08-21) with a throwaway cluster, since terminated. + `name` is a *known* field on the route — an unknown field draws + `400 request body contains an unknown field "bogusField"` and `name` does + not — and the PATCH returns success and even cycles the cluster through + PENDING, but the name never changes in `GET /v2/clusters/{id}` or + `GET /v2/clusters` (polled for two minutes). v1 workspace groups *could* be + renamed, so this is a v2 regression rather than a wrapper bug; the accepted + field makes it undetectable from the client. Worth raising with the API + team. `Cluster.update()` still passes the field through — there is nothing + better for it to do — and `TestCluster::test_update` pins the behavior. + Separately, the same route applies changes **asynchronously**: after a + `PATCH` with new `firewallRanges`, the immediately following + `GET /v2/clusters/{id}` still reports the old ranges while the cluster is + PENDING, so the trailing `refresh()` inside `Cluster.update()` does not + reflect the change. + **Wrapper status:** `Cluster.update()` now takes + `wait_on_active`/`wait_interval`/`wait_timeout`, defaulting to `False` for + backward compatibility. With it, `update()` waits for ACTIVE and then — if + `firewall_ranges` or `allow_all_traffic` was passed — for the new ranges to + be reported, before the trailing `refresh()`. On this path the firewall wait + compares against the ranges that were requested rather than merely checking + for non-empty, because the pre-PATCH ranges are already non-empty; see + `ClusterManager._wait_on_firewall`. The asynchrony itself is still an API + bug: a caller who does not opt in still gets stale values back, and the SDK + is only papering over it. Worth raising with the API team alongside the + silently-ignored `name`. +10. **Stage path normalization is inconsistent for directories.** Pre-existing + behavior in the shared `management/stage.py`, not v2-specific. + `mkdir()`/`rmdir()` append the trailing slash themselves + (`re.sub(r'/*$', '', path) + '/'`), but `info()` — and therefore + `exists()`, `is_dir()`, `is_file()` — pass the path through unchanged, so + `mkdir('d')` followed by `is_dir('d')` returns `False` while + `is_dir('d/')` returns `True`. The v1 suite happens to pass the slash + everywhere, which is why this never surfaced. Candidate fix: normalize in + `info()` too, or in `_fs_path()`. +11. **`Stage.open(path, 'r')` on a missing object raises `ManagementError`, + not `FileNotFoundError`.** Also pre-existing shared behavior. The rest of + `open()`'s builtin-open emulation raises `OSError` subclasses + (`FileExistsError` for `'x'` on an existing path, `IsADirectoryError` from + `_download_file`), so the bare 404 coming through is an inconsistency + rather than a deliberate contract. Left as-is because changing it is + visible to v1 callers too; pinned by `TestStage::test_open`. +12. **`POST /v2/clusters` applies `firewallRanges` asynchronously, outside the + state machine.** Confirmed live (2026-08-21): a cluster created with + `firewallRanges: ['0.0.0.0/0']` reaches ACTIVE with a resolvable endpoint + while `GET /v2/clusters/{id}` still reports `firewallRanges: []` — which + is deny-all, so connection attempts in that window time out at the TCP + level rather than failing authentication. `create_cluster()`'s + `wait_on_active` and its endpoint wait both completed before the firewall + landed, so the documented "wait until usable" contract was not actually + met: an SDK caller could get back an ACTIVE cluster whose endpoint refused + every connection. How long the gap lasts varies — one run had the firewall + in place by the time the tests ran, the next did not, which is what made + `TestCluster::test_connect` flaky. + **Wrapper status: fixed.** `ClusterManager._wait_on_firewall()` polls + `GET /v2/clusters/{id}` until the cluster admits inbound traffic, and + `create_cluster()` calls it under `wait_on_active`, after `_wait_on_state` + and `_wait_on_endpoint`, whenever `firewall_ranges` or `allow_all_traffic` + was requested. "Admits traffic" means non-empty `firewallRanges` **or** + `allowAllTraffic` — see item 13, the API stores a requested `0.0.0.0/0` as + the latter — rather than equality with the requested ranges, which the + server is free to normalize. The wait + is skipped for `firewall_ranges=[]`, which is a legitimate deny-all request + (see item 5) and would otherwise hang for the full timeout. The helper + lives in `v2/cluster.py` rather than the shared `manager.py` because this + is a v2 quirk and the v1 workspace path must not be affected. The live + suite's `_wait_for_firewall()` workaround is gone; `TestCluster.setUpClass` + now just asserts the SDK delivered a firewall that admits something. + The API behavior is still a bug — a caller passing `wait_on_active=False`, + or using `GET` directly, still sees the deny-all window — and is worth + raising with the API team. + **The wrapper fix is not airtight.** Observed 2026-08-25 on + `probe-pw-random-1787666958`: `create_cluster(wait_on_active=True, + firewall_ranges=['0.0.0.0/0'], allow_all_traffic=True)` returned, and the + *first* connection attempt still failed with `2003 ... (timed out)` — a TCP + timeout, the deny-all signature — while a second attempt seconds later + authenticated fine against the same endpoint. So `_wait_on_firewall()` + observed a cluster the `GET` already described as admitting traffic before + the data plane actually did. Polling the control plane cannot close this; + only a connect-retry loop would. Left as-is because it is the API's race to + fix, but any test that connects immediately after `create_cluster` should + retry rather than trust the first attempt. +13. **`POST /v2/clusters` stores `firewallRanges: ['0.0.0.0/0']` as + `allowAllTraffic: True` with `firewallRanges: []`.** Confirmed live + (2026-08-21) on a cluster since terminated: after the create settled, + `GET /v2/clusters/{id}` reported `firewallRanges: []` and + `allowAllTraffic: True`, and the endpoint accepted connections — port 3306 + open — so the empty list there does *not* mean deny-all in that + combination. The round trip is lossy: what was asked for as a range comes + back as a boolean, so a client cannot compare a create request against the + resulting cluster field by field. This is not consistent between runs + either — an earlier run of the same test had `firewallRanges: + ['0.0.0.0/0']` stored verbatim (that is what item 9's "still reports the + old ranges" observation was made against), which suggests either + region-dependent handling or an ordering effect in how the two fields are + written. Worth raising with the API team. + **Wrapper status:** `_wait_on_firewall()` treats either representation as + "reachable", and the update path additionally accepts `allowAllTraffic` as + satisfying a requested `0.0.0.0/0`. `Cluster.allow_all_traffic` was already + parsed, so nothing else changed. +14. **Four v1 workspace-group capabilities have no v2 equivalent.** Recorded + here so that `CREATE CLUSTER` can be read against a list rather than + against `create_workspace_group`'s signature. `create_workspace_group` + posts `adminPassword`, `backupBucketKMSKeyID`, `dataBucketKMSKeyID` and + `smartDR`; `POST /v2/clusters` has none of the last three, and ignores the + first (item 8). Only `highAvailabilityTwoZones` survives the move, renamed + to `multiAZ` (`v2/cluster.py:250`). + + Consequence for the Fusion grammar: `CREATE CLUSTER` deliberately offers no + `WITH PASSWORD`, KMS-key or `SMART DR` clause. A clause for any of them + would parse, be sent, and be dropped without comment — worse than not + offering it, because the statement would read as though it had taken + effect. The three KMS/DR fields are flag-only in the spec dump (see + cross-cutting item 2), so their absence at v2 is not independently + confirmable from the dump either. + + Because the password is generated and reported only in the create response, + `CREATE CLUSTER` returns a one-row result carrying `Name`, `ID`, `Endpoint` + and `AdminPassword`. `CREATE WORKSPACE GROUP` returns no row, and the + divergence is deliberate: at v1 the caller already knew the password + because it chose it, and at v2 a cluster created without capturing the + response has no reachable `admin` user. + + **Both previously open questions were settled on 2026-08-25** by one + throwaway `S-00` (`probe-adminpw-1787666027`, since terminated); see the + re-confirmation paragraph in item 8 for the POST half. + + - **`PATCH /v2/clusters/{id}` does not honour `adminPassword` either.** The + PATCH was accepted and the cluster reported ACTIVE, but connecting as + `admin` with the patched value was refused with `1045: Access denied`, + while the password generated by the original create *continued to work*. + This is the same accept-and-silently-ignore shape item 9 records for + `name`. `WITH PASSWORD` is therefore not implementable as + create-then-PATCH, and `CREATE CLUSTER` keeps returning the generated + password as its `AdminPassword` column instead. + + The first run of this probe allowed only a 30-second settle window and + never saw the cluster leave ACTIVE, so a slow-but-honored PATCH would have + looked ignored. Re-run 2026-08-25 with a 180-second window on two clusters + and two password shapes (see item 8): in all cases the patched value was + refused with `1045: Access denied` while the password generated by the + original create still authenticated. The timing caveat is closed. + - Re-confirmation of item 8 against the current API: **done**, and it holds. + The re-check mattered because finding 6's `GET /v2/regions/sharedtier` + claim had since been corrected from a live probe. --- @@ -470,8 +754,24 @@ test churn: - `WorkspaceManager.create_starter_workspace`: `project_id`. 4. **v2/ override layer** — v2 classes subclass v1 (per ADR 0001) and only override what differs. Adding fields in v1 propagates to v2 automatically - through inheritance — **no v2 changes are needed for any work in this - audit**. + through inheritance, so the *field* additions in this audit need no v2 + counterpart. + + > **Correction.** An earlier revision of this section concluded that "no v2 + > changes are needed for any work in this audit." That was drawn from the + > stale, partial `dev-docs/management_api.openapi` snapshot and is wrong. + > v2 is not a field-compatible overlay on v1: it replaces the two-level + > `workspaceGroups` → `workspaces` hierarchy with a single flat `clusters` + > resource, and it moves the Stage, shared-tier, egress, and metrics paths. + > Inheritance alone therefore leaves v2 sending v1 paths to `/v2/`, which + > 404s. See `docs/adr/0001-versioned-management-api-wrappers.md`. + > + > `dev-docs/management_api.openapi` is **not authoritative**. It was + > refreshed from `https://api.singlestore.com/spec` on 2026-08-25 (1.2.171), + > which fixed the v2 coverage but dropped all but nine v1 routes; the + > `egress` family is still missing at both versions. Confirm endpoint + > existence by probing the live API (see the header comment in that file), + > not by reading the spec. --- diff --git a/docs/shared-deployment-pool-plan.md b/docs/shared-deployment-pool-plan.md new file mode 100644 index 000000000..012050a95 --- /dev/null +++ b/docs/shared-deployment-pool-plan.md @@ -0,0 +1,197 @@ +# Sharing deployments across the management test suites + +## Goal + +Cut the serial wall time of the management suite by deploying fewer clusters, +not by polling faster. Measured target: **~1300s off an 8915s run (~15%)**, with +no change to what is asserted. + +## Why this is the lever + +A traced run (`SINGLESTOREDB_MANAGEMENT_TRACE=1`) spends 8874s of its 8915s +elapsed inside the management API -- 3729s in requests, 5145s asleep in +`wait_on_*` loops. There is no local work to optimize. Waiting for an S-00 +cluster to reach ACTIVE costs ~460s and is irreducible, so the only serial +lever is **deploying fewer of them**. + +Per-class fixture cost from that run: + +| fixture | cost | deploys | +| --- | --- | --- | +| `test_fusion.py::TestStageFusion` | 891s | 2 clusters | +| `test_fusion.py::TestJobsFusion` | 686s | 1 cluster | +| `test_management_v2.py::TestJob` | 366s | 1 cluster | +| `test_management_v2.py::TestStage` | 247s | 1 cluster | +| **total** | **2190s** | **5 clusters** | + +All four need is *a* live cluster. `TestStageFusion` needs two, because it +exercises `IN ''`; two therefore covers all four classes. + +Pool cost is one 2-cluster deployment, ~890s (which is what `TestStageFusion` +measures today for exactly that). **2190s -> ~890s.** + +## Audit: why these four are safe to share + +Established by reading every assertion in each class: + +* Every Stage assertion is scoped to one deployment's filesystem, and every + path is namespaced with the class's `cls.id`, so two classes on one cluster + cannot see each other's files. +* Job listings are filtered by job id (`show jobs {job_id} like ...`), never by + a bare listing of the deployment's jobs. +* None of the four asserts a row count over an org-wide listing. + +### Must NOT join the pool + +* `test_management_v2.py::TestCluster` (125s) and + `test_management_v1.py::TestWorkspace` (243s) -- these test the deployment + objects themselves. `TestCluster::test_update` PATCHes the cluster and cycles + it back through PENDING; a pooled cluster would break every other class. +* `TestClusterFusionCreateDrop`, `TestClusterFusionSuspendResume` -- both mutate + or destroy their subject by definition. +* `TestWorkspaceFusion` -- its three workspace groups are the subject of its + `SHOW WORKSPACE GROUPS` assertions, and it deploys them without waiting + (39.7s total), so there is nothing to save. + +## The one real implementation hazard + +`conftest.py::pytest_runtest_setup` sets a per-class owner and sweeps the +previous owner's deployments as soon as the run moves to the next class: + +```python +utils.set_owner(owner) +if previous: + _sweep_live_deployments(previous) # -> cleanup_tracked(previous) +``` + +`install_deployment_tracking()` (conftest.py:78) patches the creation methods, +so a pooled cluster built inside a `setUpClass` is tracked under *that class* +and **terminated the moment the run leaves it**. Naively hoisting the fixture +therefore produces a pool that dies after its first consumer. + +`cleanup_tracked` matches on `x[0] == owner` (utils.py:493), and +`pytest_unconfigure` calls it with `owner=None`, which matches everything. So +an entry tracked under the empty owner survives every per-class sweep and is +terminated exactly once, at session end. Create the pool with the owner +temporarily cleared: + +```python +prev = utils.get_owner() +utils.set_owner('') # session-owned: no per-class sweep matches '' +try: + cluster = mgr.create_cluster(...) +finally: + utils.set_owner(prev) +``` + +No change to `conftest.py` or the sweep is needed -- this uses the existing +mechanism as designed. + +## Steps + +1. **Add the pool helper** in `singlestoredb/tests/utils.py`: a lazily built, + process-wide pool of N v2 clusters with the owner-clearing block above. + Cache on a module global; return the same objects on every call. Have it + `raise unittest.SkipTest` for the same reasons the current fixtures do (no + US regions, no STANDARD project), so skip behaviour is unchanged. + *Verify:* a unit test that calls it twice and asserts the same cluster ids + come back, and that the tracked entry's owner is `''`. + +2. **Move `TestStageFusion` onto the pool** (`cls.cluster`, `cls.cluster_2`). + It already needs exactly two. Drop the cluster creation and the + `terminate(force=True)` calls from its `tearDownClass`; keep the env-var + save/restore and the `load_sql`/`drop_database` calls, which are per-class + and cheap (local server, not the pool cluster). + *Verify:* `pytest -v singlestoredb/tests/test_fusion.py::TestStageFusion` + passes and the trace shows no `POST clusters`. + +3. **Move `TestJobsFusion`, `test_management_v2.py::TestStage` and + `test_management_v2.py::TestJob` onto pool cluster 0.** Same edit shape. + *Verify:* each class passes standalone, then all four pooled classes pass in + one run -- that ordering is what proves the sweep does not eat the pool. + +4. **Re-run with `SINGLESTOREDB_MANAGEMENT_TRACE=1`** and confirm + `POST clusters` drops by 3 and the four fixtures total ~890s instead of + 2190s. + +## Optional follow-ups, in value order + +* **Deploy the two pool clusters concurrently** (two threads in the pool + builder, joined before returning). Pool cost ~890s -> ~460s, another ~430s. + Self-contained: two threads inside one fixture, not test-level parallelism. +* **Check whether the v1 suites can share the pool too.** `Cluster` carries a + `group` attribute (`v2/cluster.py:163`), so a v2 cluster may be addressable + as a v1 workspace group -- v1 Stage is keyed `stage/{id}/fs` where v2 is + `clusters/{id}/stage/fs`. If `cluster.group` works as that id, then + `test_management_v1.py::TestStage` and `::TestJob` (340s, deploys a group + *and* a workspace) could join, worth another ~350s. **Verify before + designing for it** -- this is a hypothesis, not a known fact. +* **`test_management_v1.py::TestStage` sharing `TestJob`'s workspace group.** + Only ~15-25s, since it creates a group without waiting on it. Low priority. + +## Out of scope + +Test-level parallelism (pytest-xdist / concurrent class execution). The pool is +a prerequisite for it but independent of it: every number above is a serial +saving. Note that a process-wide pool and xdist interact -- under xdist each +worker builds its own pool, so the saving is per worker, not per session. + +### Follow-up: what parallelism needs from the pool + +Since taken up. The pool is process-wide, so which worker a borrowing class +lands on decides how many pools get built. Four borrowers spread over four +workers is four pools -- ~890s apiece, and the saving above is gone. + +The borrowers therefore carry `xdist_group` marks +(`utils.SHARED_CLUSTER_STAGE_GROUP`, `utils.SHARED_CLUSTER_JOBS_GROUP`) and the +suite runs under `--dist loadgroup`. Two groups rather than one, split by what +they borrow: + +| group | classes | pool | +| --- | --- | --- | +| `shared-cluster-stage` | `TestStageFusion`, v2 `TestStage` | 2 clusters | +| `shared-cluster-jobs` | `TestJobsFusion`, v2 `TestJob` | 1 cluster | + +One group would serialise all four classes behind a single pool build. Two run +concurrently on separate workers, so the extra pool costs one cluster and no +wall time -- the builds overlap -- and halves the chain. The marks are inert +without `-n`: one process, one pool of two, exactly the serial behaviour above. + +`--dist loadgroup` is load-bearing beyond the groups. xdist's default `--dist +load` distributes individual tests, so a unittest class is split across workers +and each one runs `setUpClass` itself: one cluster fixture becomes N +deployments. That is slower and more expensive than running serially. Because +forgetting it costs money, it is not left to the invocation: `addopts` in +`pyproject.toml` sets `-n 3 --dist loadgroup` for every run. The command line is +applied after `addopts`, so `-n 0` still gives a serial run and an explicit +`--dist` still wins. 3 rather than `auto` because the ceiling is the API's +tolerance for concurrent provisioning, not the host's CPUs. + +One thing parallelism does not fix, and three it breaks: + +* `TestClusterFusionCreateDrop::test_create_drop_cluster` is a single test of + most of twenty minutes. One test cannot be split, so ~1200s is the floor on + wall time whatever `-n` is. +* Peak concurrent deployments rises even though total cluster-hours does not -- + the two pools, `TestClusterFusion`'s three, `CreateDrop`, `TestCluster`, v1's + group plus workspace, and both starter deployments can all be in flight at + once. The org's cluster quota and the shared-tier starter limit are what cap + `-n`, not anything in the tests. + + The API misbehaves under that burst. A cross-process cap on in-flight + creations was tried and removed: `utils.deployment_slot()` held one of + `SINGLESTOREDB_TEST_DEPLOY_CONCURRENCY` (default 3) `flock`ed slot files for + the whole `create_*` call. It did not fix the failures it was aimed at, and it + added wall time to every parallel run, so it is gone. Nothing bounds + concurrent provisioning now -- `-n` is capped by the org's cluster quota and + the shared-tier starter limit, and by whatever the API tolerates. **Open.** +* UNVERIFIED: `USE_DATA_API=1` with `-n` may not be safe on a shared + container. `load_sql` ends with `SET GLOBAL HTTP_PROXY_PORT` and + `RESTART PROXY` (`utils.py:227`), which every worker runs, and a restart + while another worker has an HTTP request in flight would drop it. The MySQL + path does not reach that branch. Not hit yet -- the parallel runs so far have + been over the MySQL protocol. +* `SINGLESTOREDB_MANAGEMENT_TRACE`'s terminal summary is lost: `conftest.py` + accumulates the traces in module globals filled in the workers, and + `pytest_terminal_summary` runs in the controller, which sees none of them. + The per-event stderr log still works. Measure with a serial run. diff --git a/docs/src/api.rst b/docs/src/api.rst index 9e0a872c8..46fc2de23 100644 --- a/docs/src/api.rst +++ b/docs/src/api.rst @@ -228,10 +228,152 @@ Management API -------------- The management objects allow you to create, destroy, and interact with -workspaces in the SingleStoreDB Cloud. +deployments in the SingleStoreDB Cloud. + +The API is versioned. Version 2 is the default, and calls a deployment a +:class:`Cluster`; version 1 called it a :class:`Workspace` inside a +:class:`WorkspaceGroup`. Which version you get is controlled by the +``management.version`` option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` +environment variable), or by passing ``version=`` to the ``manage_*`` +functions. + +The :func:`manage_clusters` function will return a :class:`ClusterManager` +object that can be used to interact with the Management API. The v1 entry +point, :func:`manage_workspaces`, is documented under `Workspaces (v1)`_. + +.. currentmodule:: singlestoredb + +.. autosummary:: + :toctree: generated/ + + manage_clusters + + +ClusterManager +.............. + +ClusterManager objects are returned by the :func:`manage_clusters` function. +They allow you to retrieve information about the clusters in your account, or +create new ones. Clusters are flat: unlike v1 workspaces they are not contained +in a group, so there is nothing to resolve before reaching one. + +.. currentmodule:: singlestoredb.management.cluster + +.. autosummary:: + :toctree: generated/ + + ClusterManager + ClusterManager.organization + ClusterManager.organizations + ClusterManager.billing + ClusterManager.clusters + ClusterManager.starter_clusters + ClusterManager.projects + ClusterManager.regions + ClusterManager.shared_tier_regions + ClusterManager.create_cluster + ClusterManager.create_starter_cluster + ClusterManager.get_cluster + ClusterManager.get_starter_cluster + ClusterManager.get_project + + +Cluster +....... + +Cluster objects are retrieved from :meth:`ClusterManager.get_cluster` or by +retrieving an element from :attr:`ClusterManager.clusters`. They are created +with :meth:`ClusterManager.create_cluster`, which replaces v1's two-step +workspace group plus workspace creation. + +.. autosummary:: + :toctree: generated/ + + Cluster + Cluster.organization + Cluster.stage + Cluster.admin_password + Cluster.connect + Cluster.refresh + Cluster.update + Cluster.suspend + Cluster.resume + Cluster.terminate + + +StarterCluster +.............. + +Starter clusters are the shared-tier deployment. They are created with +:meth:`ClusterManager.create_starter_cluster` and retrieved from +:meth:`ClusterManager.get_starter_cluster` or +:attr:`ClusterManager.starter_clusters`. They support a smaller surface than a +:class:`Cluster` — there is no update, suspend or resume. + +.. autosummary:: + :toctree: generated/ + + StarterCluster + StarterCluster.organization + StarterCluster.stage + StarterCluster.connect + StarterCluster.refresh + StarterCluster.create_user + StarterCluster.terminate + + +Project +....... + +Projects group the clusters in an organization. Project objects are retrieved +from :meth:`ClusterManager.get_project` or by retrieving an element from +:attr:`ClusterManager.projects`, and a project is what +:meth:`ClusterManager.create_cluster` places a new cluster in. An organization +with exactly one project does not need to name it. + +.. autosummary:: + :toctree: generated/ + + Project + + +Workspaces (v1) +............... + +.. deprecated:: Management API v1 as a whole is deprecated, not just the + workspace vocabulary below. ``management.version`` now defaults to ``'v2'``, + and every entry point that resolves to v1 raises a + :class:`DeprecationWarning` -- whether v1 was named with ``version='v1'`` or + inherited from the ``management.version`` option + (``SINGLESTOREDB_MANAGEMENT_VERSION``). + + **v1 still works.** Deprecated here means warned about, not removed: every + function and class below still operates against the live v1 endpoints, and + :func:`manage_workspaces` still returns a working + :class:`WorkspaceManager` without being asked for a version. Nothing raises + because the default moved. When you are ready to move off v1: + + ============================== ============================== + v1 v2 + ============================== ============================== + :func:`manage_workspaces` :func:`manage_clusters` + :class:`WorkspaceManager` :class:`ClusterManager` + :class:`WorkspaceGroup` :class:`Cluster` + :class:`Workspace` :class:`Cluster` + :class:`StarterWorkspace` :class:`StarterCluster` + ============================== ============================== + + Note that :class:`WorkspaceGroup` and :class:`Workspace` both collapse onto + :class:`Cluster`: v2 is flat, so there is no container resource and no + two-step create. Grouping is expressed by :class:`Project`, which is an + organizational unit rather than a deployment parent. + + :func:`manage_files` and :func:`manage_regions` need no migration -- their + routes are identical at both versions, so simply stop passing + ``version='v1'``. The :func:`manage_workspaces` function will return a :class:`WorkspaceManager` -object that can be used to interact with the Management API. +object that can be used to interact with version 1 of the Management API. .. currentmodule:: singlestoredb @@ -305,7 +447,8 @@ Workspaces are created within WorkspaceGroups. They can be created using either Region ...... -Region objects are accessed from the :attr:`WorkspaceManager.regions` attribute. +Region objects are accessed from the :attr:`ClusterManager.regions` attribute, +or from :attr:`WorkspaceManager.regions` at v1. .. currentmodule:: singlestoredb.management.region @@ -318,8 +461,9 @@ Region objects are accessed from the :attr:`WorkspaceManager.regions` attribute. Organization ............ -Organization objects are retrieved from :attr:`WorkspaceManager.organization`. -They provide access to organization-level resources and operations. +Organization objects are retrieved from :attr:`ClusterManager.organization`, or +from :attr:`WorkspaceManager.organization` at v1. They provide access to +organization-level resources and operations. .. currentmodule:: singlestoredb.management.organization @@ -329,7 +473,6 @@ They provide access to organization-level resources and operations. Organization Organization.get_secret Organization.jobs - Organization.inference_apis Secret @@ -413,12 +556,11 @@ The following classes are used as parameters and return values in the jobs API. Stage Files ........... -To interact with files in your Stage, use the -:attr:`WorkspaceGroup.stage` attribute. -It will return a :class:`Stage` object which defines the following -methods and attributes. +To interact with files in your Stage, use the :attr:`Cluster.stage` attribute +(:attr:`WorkspaceGroup.stage` at v1). It will return a :class:`Stage` object +which defines the following methods and attributes. -.. currentmodule:: singlestoredb.management.workspace +.. currentmodule:: singlestoredb.management.stage .. autosummary:: :toctree: generated/ @@ -499,11 +641,11 @@ personal, shared, or model files. FilesObject ........... -:class:`FilesObject`s are returned by the :meth:`StageObject.upload_file` -:meth:`FilesObject.upload_folder`, :meth:`FilesObject.mkdir`, -:meth:`FilesObject.rename`, and :meth:`FilesObject.info` methods. +:class:`FilesObject`s are returned by the :meth:`FileSpace.upload_file`, +:meth:`FileSpace.upload_folder`, :meth:`FileSpace.mkdir`, +:meth:`FileSpace.rename`, and :meth:`FileSpace.info` methods. -.. currentmodule:: singlestoredb.management.workspace +.. currentmodule:: singlestoredb.management.files .. autosummary:: :toctree: generated/ diff --git a/docs/src/conf.py b/docs/src/conf.py index b127c3bfa..aef853cf4 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -46,11 +46,11 @@ autoclass_content = 'class' intersphinx_mapping = { - 'python': ('https://docs.python.org/', None), - 'pandas': ('http://pandas.pydata.org/pandas-docs/version/0.19.2/', None), - 'numpy': ('http://docs.scipy.org/doc/numpy/', None), - 'scipy': ('http://docs.scipy.org/doc/scipy/reference/', None), - 'matplotlib': ('http://matplotlib.sourceforge.net/', None), + 'python': ('https://docs.python.org/3/', None), + 'pandas': ('https://pandas.pydata.org/docs/', None), + 'numpy': ('https://numpy.org/doc/stable/', None), + 'scipy': ('https://docs.scipy.org/doc/scipy/', None), + 'matplotlib': ('https://matplotlib.org/stable/', None), } # Add any paths that contain templates here, relative to this directory. diff --git a/docs/stage-upload-round-trips-plan.md b/docs/stage-upload-round-trips-plan.md new file mode 100644 index 000000000..c36b0b694 --- /dev/null +++ b/docs/stage-upload-round-trips-plan.md @@ -0,0 +1,288 @@ +# Cutting the round trips in the Stage and Files write paths + +`UPLOAD FILE TO STAGE 'stats.csv' IN '' FROM 'mydata.csv'` cost **six HTTP +requests** to upload one file when this was written. Only one of them transfers +anything. This plan is staged so each part lands and ships on its own. Stages 1 +and 2 have landed and took the count to **three**, or four with `OVERWRITE`. +Stage 3 and Stage 4 are open. + +## Measured baseline + +Against a live organization (`S2DB Eng - Launchpad`, one ACTIVE cluster), for a +60-byte CSV: + +| # | Call | Origin | Typical | +|---|------|--------|---------| +| 1 | `GET /v2/clusters` | `get_deployment` resolving `IN ''` | ~2000-3600 ms | +| 2 | `GET /v2/projects` | `Cluster.from_dict` → `_project_from_id`, per manager | ~280 ms | +| 3 | `GET .../stage/fs/?metadata=1` | `Stage.upload_file`'s `exists()` | ~1000 ms | +| 4 | `GET .../stage/fs/?metadata=1` | `Stage._upload`'s `exists()` — **the same check** | ~1000 ms | +| 5 | `PUT .../stage/fs/` | the upload | 1500-31000 ms | +| 6 | `GET .../stage/fs/?metadata=1` | `_upload` returns `info()`, which Fusion **discards** | ~1000 ms | + +With `OVERWRITE` against a file that already exists, `remove()` inserts an +`is_dir()` metadata `GET` plus a `DELETE`, making it eight. + +Rows 4 and 6 are gone as of `400d5b71` — Stage 1. Row 2 is gone with Stage 2, +and Stage 1c collapsed the `OVERWRITE` path's second metadata `GET`. The live +count is now rows 1, 3 and 5: **three**, or four with `OVERWRITE`. + +One call the table missed: a cluster payload that carries a `region` made +`Cluster.from_dict` resolve it against `ClusterManager.regions`, which is a +`GET /v2/regions` on the first cluster a manager built — the same eager-resolve +shape Stage 2 removed from the project. So the live figure for a fresh manager +was one more than the count above. `Cluster.region` got the same lazy treatment +as `Cluster.project` at the same time, so the three-call figure now holds for a +realistic listing too; +`test_a_region_on_the_payload_costs_nothing` pins it. + +### The part we cannot fix + +Stage route latency is erratic and payload-independent. Repeated 8-byte uploads +to the same cluster measured `PUT` at 30665 ms, 1534 ms and 2522 ms; a `DELETE` +of an 8-byte file took 27214 ms while another took 1376 ms. A 100 KB `PUT` took +31286 ms — the same range as 8 bytes, so this is not transfer time. + +Two consequences for this plan. First, no client change makes an upload reliably +fast; the ceiling is the route. Second, **reducing the call count is still the +right work**, because every call is an independent chance to draw a 30-second +stall. Going from six calls to three halves the exposure. Do not expect the +stopwatch to prove it on any single run — the variance swamps the difference. +Stage 4 exists to get the route itself looked at. + +### Two facts established while measuring + +* There is **no server-side filter by name.** A cluster can be fetched by ID — + `ClusterManager.get_cluster(id)` (`v2/cluster.py:1455`) is a path lookup, + `GET /v2/clusters/` — but a name can only be resolved by listing every + cluster and filtering client-side. This is why `IN ''` already costs one + call, through `_deployment_by_id` (`fusion/handlers/utils.py:566-575`), and only + `IN ''` pays the full listing. Nothing in this plan changes that; the + listing is the floor for the name spelling. +* The listing cost is fixed route overhead, not payload: ~2245 ms for a + one-cluster org. Trimming what comes back would not help even if it were + possible. +* The `PUT` response body is only `{"name": ..., "path": ...}`. It carries none + of `size`, `type`, `format`, `mimetype` or `writable`, so + `FilesObject.from_dict` cannot be fed from it. `_upload` cannot skip its + trailing `info()` by reusing the write response, which is why Stage 1b has to + be done at the caller instead. + +## Stage 1 — remove the two redundant calls in the upload path — **landed (`400d5b71`)** + +Independent of every other stage. Two unambiguous defects, no design question. + +### 1a. The duplicated `exists()` / `remove()` + +`Stage.upload_file` (`singlestoredb/management/stage.py:188-197`) checks +`exists()`, raises or `remove()`s, then delegates to `_upload`, which does +exactly the same thing again (`stage.py:293-296`). `FileSpace.upload_file` and +`FileSpace._upload` carry the identical pair (`management/files.py:694-704` and +`files.py:803-806`). + +Delete the check from both `upload_file` methods and let `_upload` own it. The +messages are already identical within each class — `'stage path already +exists'`, `'file path already exists'` — so nothing observable changes. + +One wrinkle to handle rather than inherit: `upload_file` currently opens the +local file *after* its `exists()` check, so removing the check means the handle +is opened before `_upload` raises `OSError` on a non-`overwrite` conflict, +leaking it until GC. Wrap the `open()` in a `with` block. + +`upload_folder` calls `upload_file` per file (`stage.py:267`), so a folder upload +saves one call per file. + +**Verify:** existing `test_upload_file` coverage in `test_management_v2.py:1500` +and `test_management_v1.py:380` already asserts the conflict `OSError` and the +`overwrite=True` path; both must still pass. Add a unit test that counts +requests through a mocked manager and pins the count, so the redundancy cannot +come back — that harness is the one thing this plan needs that does not exist +yet. + +### 1b. The discarded `info()` + +`_upload` ends `return self.info(stage_path)`, and +`UploadStageFileHandler.run` (`fusion/handlers/stage.py:199-204`) throws the +result away. Same in `fusion/handlers/files.py:201` and +`fusion/handlers/models.py:154` — all three Fusion upload handlers return +`None`. + +The `PUT` body cannot supply the `FilesObject` (see above), and `upload_file`'s +public contract returns one, so the `info()` cannot simply go. Give the handlers +a path that does not ask for it. Preferred shape: a private +`_upload(..., fetch_info: bool = True)` returning `Optional[FilesObject]`, with +the three Fusion handlers calling `_upload(..., fetch_info=False)` through a thin +`upload_file`-shaped helper so they keep the `IsADirectoryError` check on the +local path. + +**Verify:** the request-count test from 1a covers this too. Assert that the +public `upload_file` still returns a populated `FilesObject`. + +**Stage 1 payoff:** six calls to four, or eight to six with `OVERWRITE`. + +### 1c. The `OVERWRITE` path fetches the same metadata twice — **landed** + +Left over from 1a rather than introduced by it. `_upload` (`stage.py:287-291`) +calls `exists()`, which is `info()` behind a `try` (`stage.py:395`), and then +`remove()`, which opens with `is_dir()` — `info()` again (`stage.py:732`), on the +same path, with nothing in between that could have changed it. + +Fetch the metadata once in `_upload` and branch on the object: absent → `PUT`; +present and not `overwrite` → `OSError`; present and a directory → +`IsADirectoryError`; otherwise `DELETE` and `PUT`. No caching, no new state — +the second call is reading a value the frame already holds. `remove()` keeps its +own `is_dir()` for its other callers. + +`FileSpace._upload` (`files.py:803-806`) carries the same pair. + +**Verify:** the request-count harness pins the `OVERWRITE` count at four. The +`IsADirectoryError` that `remove()` currently raises through `_upload` must +still be raised, with the same message. + +**1c payoff:** six calls to five with `OVERWRITE`; nothing on the plain path. + +**Landed as:** the shared `FileLocation._info_or_none` (`management/files.py`), +which `Stage._upload` and `FileSpace._upload` branch on. `remove()` keeps its +own `is_dir()`, as planned. Pinned by +`test_an_overwrite_costs_one_check_and_one_delete`, +`test_an_overwrite_of_a_folder_raises_on_the_one_check` and the two file-space +equivalents in `test_management_utils.py`. + +## Stage 2 — resolve the deployment in one call, not two — **landed** + +Depends on nothing in Stage 1. No caching, no new state, no open decision. + +`get_deployment` resolves `IN ''` by filtering `manager.clusters` +(`fusion/handlers/utils.py:491`), which costs `GET /v2/clusters`. It then costs a +second call it never uses: `Cluster.from_dict` resolves `_project_from_id` for +every cluster in the listing (`v2/cluster.py:409` and `:813`), and that reads +`ClusterManager.projects`, so name resolution drags `GET /v2/projects` (~280 ms) +along behind it. + +### Why there is no caching here + +An earlier draft of this stage proposed memoizing name → ID, on the theory that a +notebook looping four `CREATE STAGE FOLDER ... IN ''` statements pays name +resolution four times. It does, but that is **cross-statement** state, and the +staleness it buys is not worth it: a renamed or replaced cluster keeps resolving +to the old ID until the entry expires, and `DROP CLUSTER` / `CREATE CLUSTER` +would each need to invalidate it. Dropped. + +A **statement-scoped** cache — the narrow, obviously-safe version — was checked +and is dead weight. Inside one statement there is nothing to hit twice: + +* every stage handler calls `get_deployment` exactly **once** per `run` + (`fusion/handlers/stage.py:92,199,293,370,435,499`); +* `ClusterManager.projects` is *already* a one-hour `ttl_property` + (`v2/cluster.py:1015`), so `_project_from_id` costs one `GET /v2/projects` per + manager no matter how many clusters the listing holds. + +A per-statement memo would have a 0% hit rate on the upload path. The fix is not +to cache the second call, it is to not make it. + +### The change + +* **Make `Cluster.project` lazy.** This is the whole payoff. Stop calling + `_project_from_id` in `Cluster.from_dict`; keep the `projectID` and resolve + `Project` on first access to `.project`. There are exactly two readers — + `fusion/handlers/cluster.py:74` (`SHOW CLUSTERS EXTENDED`) and + `v2/cluster.py:1207` — and both must keep working, including the + `Project(id=..., name='')` fallback for an ID that matches no project. + `StarterCluster` (`v2/cluster.py:813`) gets the same treatment. Note that the + `ttl_property` on `projects` stays useful: `SHOW CLUSTERS EXTENDED` reads + `.project` per row, and one manager must still serve all of them from one + fetch. + +An earlier draft paired this with server-side filtering +(`GET /v2/clusters?name=`) in place of the list-everything-then-filter in +`get_deployment`. Dropped: **the API has no name filter.** Only ID lookup is +server-side, and that path is already taken by `_deployment_by_id`. The +client-side filter and its ambiguity check stay exactly as they are. + +### What is left afterwards + +One call, `GET /v2/clusters`, at a ~2000 ms floor that is fixed route overhead. +That floor is then the entire cost of name resolution and there is nothing +further the client can do about it — it is Stage 4's reporting job. + +**Verify:** the request-count harness from Stage 1, extended to count a whole +Fusion statement rather than a `Stage` call. Pin the count for +`UPLOAD FILE TO STAGE ... IN ''` at three, and assert no +`GET /v2/projects` is issued. Separately assert `SHOW CLUSTERS EXTENDED` still +reports the project name and issues `GET /v2/projects` exactly once regardless of +cluster count. + +**Landed as:** `Cluster.project` and `StarterCluster.project` are properties +over a stored `_project_id`, resolved by `_lazy_project` on first read +(`v2/cluster.py`). `_project_from_id` and its `` fallback are unchanged +and still what does the resolving; a `Project` passed to the constructor is +still kept as it stands. + +`Cluster.region` was given the same treatment in the same pass — `_region_args` +stores the reported name, `_lazy_region` matches it against +`ClusterManager.regions` on first read and falls back to a `Region` built from +what the cluster itself reported. That is the "one call the table missed" above, +and it is why the three-call figure holds for a payload carrying a region. + +Two consequences worth knowing: + +* Neither lazy value is in `vars(cluster)`, so `vars_to_str` would drop both + from `str(cluster)` — and resolving them in `__repr__` would make printing a + cluster issue two requests. `vars_to_str` grew an `extra=` argument for + exactly this: `Cluster.__str__` passes the resolved object if something has + already read the property and the reported ID / name otherwise, so printing + stays free. `test_printing_a_cluster_costs_nothing` and + `test_printing_a_cluster_shows_what_is_resolved` pin both halves. +* `SHOW CLUSTERS EXTENDED` reported `ProjectID`, not the project name the plan + said. Renamed to `ProjectName`, along with `SHOW STARTER CLUSTERS EXTENDED`'s + column, since `_project_from_id`'s `` fallback means the name is + always readable. Column-name assertions in `test_fusion.py` moved with it. + +**The harness:** `CountingClusterManager`, `counting_cluster_manager` and +`run_fusion_statement` in `singlestoredb/tests/utils.py`, next to the +`CountingManager` Stage 1 introduced. It serves the management routes a +statement resolves a deployment through from fixture payloads, delegates the +Stage filesystem routes to a `CountingManager` sharing its `calls` list, and +raises on any route nobody accounted for. Before Stages 2 and 1c landed it +reproduced the counts in the table above exactly: four for a plain upload, six +with `OVERWRITE`. Stage 3 should extend it rather than write another one. + +## Stage 3 — the folder and listing paths + +Depends on nothing. Lower priority; same class of defect, different methods. + +* `mkdir` (`stage.py:322-334`, `files.py:815+`) does `exists()`, then possibly + `info()`, then `PUT`, then `info()` again — up to four calls to create one + folder. `CREATE STAGE FOLDER` discards the return, exactly like 1b. The + `exists()` and the `info()` are the same `GET` on the same path back to back, + so the same one-fetch-and-branch rewrite as 1c applies. +* `remove` calls `is_dir()`, which is a full `info()`, before its `DELETE`. 1c + stops the upload path paying for it; `remove` keeps it for its other callers. +* `SHOW STAGE FILES ... EXTENDED` calls `stage.info(x)` **per entry** on top of + the `listdir` (`fusion/handlers/stage.py:105-116`). A 20-file listing is 21 + calls. Check whether `listdir` can be asked for metadata in one request; if + not, this one is inherent and should be documented as such rather than + "fixed". + +## Stage 4 — get the route latency looked at + +Not an SDK change. The numbers in this document — 30-second stalls on 8-byte +writes, a 2-second floor on `GET /v2/clusters` — belong with whoever owns those +routes. Two things worth doing here: + +* Extend `SINGLESTOREDB_MANAGEMENT_TRACE` coverage so Stage calls show up in the + per-route breakdown the way management calls already do, giving anyone + reporting this a reproduction rather than an anecdote. +* File the `GET /v2/clusters` floor separately from the Stage `PUT`/`DELETE` + variance. They are different routes and probably different causes. + +## Order and independence + +Stages 1, 2 and 3 touch disjoint code and can land in any order or in parallel. +Stages 1 and 2 have landed. Stage 3 is what is left of the SDK-side work, and +Stage 4 is reporting and can proceed alongside. + +The request-count harness is in `singlestoredb/tests/utils.py` +(`CountingManager` for a `Stage` or `FileSpace` call, `CountingClusterManager` +plus `run_fusion_statement` for a whole statement) and is what makes Stage 3 +checkable. diff --git a/pyproject.toml b/pyproject.toml index c8d624154..ff445f2bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ test = [ "pydantic", "pytest", "pytest-cov", + "pytest-xdist", "singlestore-vectorstore>=0.1.2", "uvicorn", ] @@ -90,8 +91,27 @@ exclude = ["docs*", "resources*", "examples*", "licenses*"] "*" = ["*.typed", "*.sql", "*.csv", "*.ipynb"] [tool.pytest.ini_options] +# Parallel by default. Both options are overridable on the command line, which +# comes after addopts: -n 1 or -n 0 for a serial run (0 takes xdist out of the +# picture entirely, which is what --pdb needs), --dist load to change grouping. +# +# loadgroup rather than xdist's default load: the default splits a unittest +# class across workers, and each one then runs setUpClass itself, so one shared +# cluster fixture becomes N deployments. loadgroup keeps a class on one worker +# and honours the xdist_group marks below. It is here rather than in the +# invocation because forgetting it costs real money. +# +# 3 workers, not `auto`: the ceiling is the management API's tolerance for +# concurrent provisioning and the org's cluster quota, not this host's CPUs. +# +# Note that xdist must be installed for pytest to start at all with these set +# (`pip install -e ".[test]"`), that SINGLESTOREDB_MANAGEMENT_TRACE's terminal +# summary needs -n 0, and that USE_DATA_API=1 in parallel is unverified. +addopts = ["-n", "3", "--dist", "loadgroup"] markers = [ "management", + "management_v1: exercises the v1 management API, which v2 has replaced. Deselect with -m 'not management_v1'; the v1 endpoints only need a nightly gate now that v2 is the default.", + "xdist_group: pytest-xdist's own marker, declared here so a run without the plugin installed does not warn on it. Applied to the classes that borrow from the shared cluster pool; see singlestoredb/tests/utils.py.", ] [tool.mypy] diff --git a/resources/create_test_cluster.py b/resources/create_test_cluster.py index 48c22e221..186fadfa8 100755 --- a/resources/create_test_cluster.py +++ b/resources/create_test_cluster.py @@ -71,8 +71,11 @@ sys.exit(1) -# Connect to workspace -wm = s2.manage_workspaces(options.token or None) +# Connect to workspace. This is still the deprecated v1 workspace-group +# grammar because the v1 test suite it sets up needs workspace groups; +# it gets ported to manage_clusters() when that suite goes. Pinned to v1 +# because manage_workspaces() otherwise follows the management.version option. +wm = s2.manage_workspaces(options.token or None, version='v1') # Find matching region if '::' in options.region: diff --git a/resources/drop_test_cluster.py b/resources/drop_test_cluster.py index 30725afd5..16ed7539d 100755 --- a/resources/drop_test_cluster.py +++ b/resources/drop_test_cluster.py @@ -23,8 +23,11 @@ sys.exit(1) -# Connect to workspace -wm = s2.manage_workspaces(options.token or None) +# Connect to workspace. This is still the deprecated v1 workspace-group +# grammar because the v1 test suite it sets up needs workspace groups; +# it gets ported to manage_clusters() when that suite goes. Pinned to v1 +# because manage_workspaces() otherwise follows the management.version option. +wm = s2.manage_workspaces(options.token or None, version='v1') wg_name = 'Python Client Testing' diff --git a/singlestoredb/__init__.py b/singlestoredb/__init__.py index 897163e31..0a7fb5faa 100644 --- a/singlestoredb/__init__.py +++ b/singlestoredb/__init__.py @@ -25,7 +25,7 @@ DataError, ManagementError, ) from .management import ( - manage_cluster, manage_workspaces, manage_files, manage_regions, + manage_clusters, manage_files, manage_regions, manage_workspaces, ) from .types import ( Date, Time, Timestamp, DateFromTicks, TimeFromTicks, TimestampFromTicks, diff --git a/singlestoredb/_management_version.py b/singlestoredb/_management_version.py new file mode 100644 index 000000000..4425ca418 --- /dev/null +++ b/singlestoredb/_management_version.py @@ -0,0 +1,21 @@ +#!/usr/bin/env python +""" +Management API versions this SDK is built against. + +This module imports nothing, so both :mod:`singlestoredb.config` -- which +registers the ``management.version`` option -- and +:mod:`singlestoredb.management._version_import` can read these without an +import cycle. Neither package can host them: ``config`` is imported before +``management``, and ``management.manager`` imports ``config``. +""" +#: Management API version used when nothing else names one. Changing this +#: retargets the ``management.version`` option default, the ``manage_*`` +#: factories, and the ``default_version`` of every version-neutral manager +#: class. Classes that implement one specific version name it literally +#: instead, and do not follow this. +DEFAULT_MANAGEMENT_VERSION = 'v2' + +#: Management API version being wound down. Public entry points that resolve +#: to it raise a :class:`DeprecationWarning`, and everything under +#: ``singlestoredb.management.v1`` goes away with it. +DEPRECATED_MANAGEMENT_VERSION = 'v1' diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py index 6636fe8d5..23d7d9009 100644 --- a/singlestoredb/ai/chat.py +++ b/singlestoredb/ai/chat.py @@ -6,8 +6,8 @@ import httpx -from singlestoredb import manage_workspaces -from singlestoredb.management.inference_api import InferenceAPIInfo +from singlestoredb.management.v1.inference_api import InferenceAPIInfo +from singlestoredb.management.workspace import _manage_workspaces_v1 try: from langchain_openai import ChatOpenAI @@ -49,7 +49,7 @@ def SingleStoreChatFactory( hosting_platform = os.environ.get('SINGLESTOREDB_INFERENCE_API_HOSTING_PLATFORM') if base_url is None or hosting_platform is None: inference_api_manager = ( - manage_workspaces().organizations.current.inference_apis + _manage_workspaces_v1().organizations.current.inference_apis ) info = inference_api_manager.get(model_name=model_name) if not info.internal_connection_url: diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index ac2ced1f5..bd7975829 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -8,8 +8,8 @@ import httpx -from singlestoredb import manage_workspaces -from singlestoredb.management.inference_api import InferenceAPIInfo +from singlestoredb.management.v1.inference_api import InferenceAPIInfo +from singlestoredb.management.workspace import _manage_workspaces_v1 try: from langchain_openai import OpenAIEmbeddings @@ -133,7 +133,7 @@ def SingleStoreEmbeddingsFactory( hosting_platform = os.environ.get('SINGLESTOREDB_INFERENCE_API_HOSTING_PLATFORM') if base_url is None or hosting_platform is None: inference_api_manager = ( - manage_workspaces().organizations.current.inference_apis + _manage_workspaces_v1().organizations.current.inference_apis ) info = inference_api_manager.get(model_name=model_name) if not info.internal_connection_url: diff --git a/singlestoredb/config.py b/singlestoredb/config.py index 594b06cf8..bb9914b69 100644 --- a/singlestoredb/config.py +++ b/singlestoredb/config.py @@ -4,6 +4,7 @@ import os from . import auth +from ._management_version import DEFAULT_MANAGEMENT_VERSION from .utils.config import check_bool # noqa: F401 from .utils.config import check_dict_str_str # noqa: F401 from .utils.config import check_float # noqa: F401 @@ -309,12 +310,23 @@ environ=['SINGLESTOREDB_MANAGEMENT_BASE_URL'], ) +# Set this to 'v1' -- or pass version='v1' to a manage_* factory -- to address +# the v1 endpoints, which remain reachable until management/v1/ is removed. +# The default comes from _management_version so that this option, the manage_* +# factories, and Manager.default_version all move together. register_option( - 'management.version', 'string', check_str, 'v1', + 'management.version', 'string', check_str, DEFAULT_MANAGEMENT_VERSION, 'Specifies the version for the management API.', environ=['SINGLESTOREDB_MANAGEMENT_VERSION'], ) +register_option( + 'management.trace', 'bool', check_bool, False, + 'Log the duration of every management API request and every poll ' + 'the wait_on_* loops sleep through to stderr.', + environ=['SINGLESTOREDB_MANAGEMENT_TRACE'], +) + # # External function options diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 63a06193f..af3dbd385 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -68,8 +68,8 @@ from . import rowdat_1 from . import utils from ... import connection -from ... import manage_workspaces from ...config import get_option +from ...management.stage import get_stage from ...mysql.constants import FIELD_TYPE as ft from ..signature import get_signature from ..signature import signature_to_sql @@ -1992,21 +1992,20 @@ def to_environment( if not url.path or url.path == '/': raise ValueError(f'no stage path was specified: {destination}') - mgr = manage_workspaces() - if url.hostname: - wsg = mgr.get_workspace_group(url.hostname) - elif os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP'): - wsg = mgr.get_workspace_group( - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'], - ) - else: - raise ValueError(f'no workspace group specified: {destination}') + # The host names the deployment whose Stage is wanted: a cluster at + # v2, a workspace group at v1. With no host, get_stage falls back to + # the deployment named by the environment -- SINGLESTOREDB_WORKSPACE + # at v2, SINGLESTOREDB_WORKSPACE_GROUP at v1. + try: + stage = get_stage(url.hostname or None) + except RuntimeError: + raise ValueError(f'no deployment specified: {destination}') # Make intermediate directories if url.path.count('/') > 1: - wsg.stage.mkdirs(os.path.dirname(url.path)) + stage.mkdirs(os.path.dirname(url.path)) - wsg.stage.upload_file( + stage.upload_file( local_path, url.path + f'{name}.env', overwrite=overwrite, ) @@ -2205,21 +2204,21 @@ def main(argv: Optional[List[str]] = None) -> None: if url.path.endswith('/'): raise ValueError(f'an environment file must be specified: {f}') - mgr = manage_workspaces() - if url.hostname: - wsg = mgr.get_workspace_group(url.hostname) - elif os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP'): - wsg = mgr.get_workspace_group( - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'], - ) - else: - raise ValueError(f'no workspace group specified: {f}') + # The host names the deployment whose Stage is wanted: a cluster + # at v2, a workspace group at v1. With no host, get_stage falls + # back to the deployment named by the environment -- + # SINGLESTOREDB_WORKSPACE at v2, + # SINGLESTOREDB_WORKSPACE_GROUP at v1. + try: + stage = get_stage(url.hostname or None) + except RuntimeError: + raise ValueError(f'no deployment specified: {f}') if tmpdir is None: tmpdir = tempfile.TemporaryDirectory() local_path = os.path.join(tmpdir.name, url.path.split('/')[-1]) - wsg.stage.download_file(url.path, local_path) + stage.download_file(url.path, local_path) args.functions[i] = local_path elif f.startswith('http://') or f.startswith('https://'): diff --git a/singlestoredb/functions/ext/mmap.py b/singlestoredb/functions/ext/mmap.py index df200fa14..ca8a96005 100644 --- a/singlestoredb/functions/ext/mmap.py +++ b/singlestoredb/functions/ext/mmap.py @@ -62,8 +62,8 @@ def print_it_pandas(x2: float, x3: str) -> str: from . import asgi from . import utils -from ... import manage_workspaces from ...config import get_option +from ...management.stage import get_stage logger = utils.get_logger('singlestoredb.functions.ext.mmap') @@ -266,21 +266,21 @@ def main(argv: Optional[List[str]] = None) -> None: if url.path.endswith('/'): raise ValueError(f'an environment file must be specified: {f}') - mgr = manage_workspaces() - if url.hostname: - wsg = mgr.get_workspace_group(url.hostname) - elif os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP'): - wsg = mgr.get_workspace_group( - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'], - ) - else: - raise ValueError(f'no workspace group specified: {f}') + # The host names the deployment whose Stage is wanted: a cluster + # at v2, a workspace group at v1. With no host, get_stage falls + # back to the deployment named by the environment -- + # SINGLESTOREDB_WORKSPACE at v2, + # SINGLESTOREDB_WORKSPACE_GROUP at v1. + try: + stage = get_stage(url.hostname or None) + except RuntimeError: + raise ValueError(f'no deployment specified: {f}') if tmpdir is None: tmpdir = tempfile.TemporaryDirectory() local_path = os.path.join(tmpdir.name, url.path.split('/')[-1]) - wsg.stage.download_file(url.path, local_path) + stage.download_file(url.path, local_path) args.functions[i] = local_path elif f.startswith('http://') or f.startswith('https://'): diff --git a/singlestoredb/fusion/README.md b/singlestoredb/fusion/README.md index e7868615a..530e1b9e6 100644 --- a/singlestoredb/fusion/README.md +++ b/singlestoredb/fusion/README.md @@ -207,7 +207,9 @@ ShowMonthHandler.register() ## Example Here is a more complete example demonstrating optional values, selection groups, -and repeated values. +and repeated values. It is abridged from `handlers/workspace.py`, which speaks +the deprecated management API v1 vocabulary; see `handlers/cluster.py` for the +current `CLUSTER` commands. ```python class CreateWorkspaceGroupHandler(SQLHandler): diff --git a/singlestoredb/fusion/handler.py b/singlestoredb/fusion/handler.py index 929b8a6a6..18b2f5458 100644 --- a/singlestoredb/fusion/handler.py +++ b/singlestoredb/fusion/handler.py @@ -22,6 +22,7 @@ from . import result from ..connection import Connection +from ..warnings import DeprecatedFeatureWarning from ..warnings import PreviewFeatureWarning CORE_GRAMMAR = r''' @@ -584,6 +585,12 @@ class SQLHandler(NodeVisitor): _enabled: bool = True _preview: bool = False + #: Command that replaces this one, e.g. ``'SHOW CLUSTERS'``. When set, the + #: command still runs but warns on every execution. Used for the management + #: API v1 vocabulary (``handlers/workspace.py``), which v2 replaced with the + #: flat ``CLUSTER`` commands. Empty means not deprecated. + _deprecated_by: str = '' + def __init__(self, connection: Connection): self.connection = connection self._handled: Set[str] = set() @@ -665,6 +672,17 @@ def execute(self, sql: str) -> result.FusionSQLResult: ) type(self).compile() + + if type(self)._deprecated_by: + # After compile(), so that command_key is populated -- naming the + # command the user actually typed is the point of the message. + warnings.warn( + f'{" ".join(type(self).command_key).upper()} is a management ' + 'API v1 command and is deprecated. Use ' + f'{type(self)._deprecated_by} instead.', + DeprecatedFeatureWarning, stacklevel=2, + ) + self._handled = set() try: params = self.visit(type(self).grammar.parse(sql)) @@ -725,12 +743,15 @@ def visit_qs(self, node: Node, visited_children: Iterable[Any]) -> Any: def visit_compound(self, node: Node, visited_children: Iterable[Any]) -> Any: """Compound name.""" - print(visited_children) return flatten(visited_children)[0] def visit_number(self, node: Node, visited_children: Iterable[Any]) -> Any: """Numeric value.""" - return float(flatten(visited_children)[0]) + # The `number` rule is ` ws*`, so node.text carries the trailing + # whitespace *and* any trailing /* comment */. Take the regex child's + # text: unlike flatten(visited_children)[0] it is not confused by the + # optional fraction group, which matches empty for a bare integer. + return float(node.children[0].text) def visit_integer(self, node: Node, visited_children: Iterable[Any]) -> Any: """Integer value.""" diff --git a/singlestoredb/fusion/handlers/cluster.py b/singlestoredb/fusion/handlers/cluster.py new file mode 100644 index 000000000..4c38b40fb --- /dev/null +++ b/singlestoredb/fusion/handlers/cluster.py @@ -0,0 +1,1108 @@ +#!/usr/bin/env python3 +""" +Fusion SQL handlers for the management API cluster vocabulary. + +Every handler here reaches the API through :func:`get_cluster_manager`, so the +``CLUSTER`` commands always address the cluster resource whatever the +``management.version`` option is set to. + +A cluster is a single flat deployment resource: one ``CREATE CLUSTER`` +statement provisions it, and the compute settings (size, auto-suspend, cache) +and the deployment-wide settings (firewall, update window, expiration) all live +on the one object. A region is identified by its ``(provider, region_name)`` +pair rather than by an ID, so ``IN REGION`` takes a name and there is no +``IN REGION ID`` alternate. +""" +import json +from typing import Any +from typing import Dict +from typing import Optional + +from .. import result +from ...management.cluster import ClusterManager +from ..handler import SQLHandler +from ..result import FusionSQLResult +from .utils import dt_isoformat +from .utils import get_cluster +from .utils import get_cluster_manager +from .utils import get_project +from .utils import get_starter_cluster + +#: Seconds per unit for the ``AUTO SUSPEND AFTER`` clause. +_SUSPEND_UNIT_SECONDS = dict( + SECONDS=1, + MINUTES=60, + HOURS=60 * 60, + DAYS=60 * 60 * 24, +) + + +def _auto_suspend(params: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Convert an ``AUTO SUSPEND AFTER`` clause to API parameters.""" + if not params.get('auto_suspend'): + return None + # The clause parses to one flat dict, not a list of one dict per + # sub-rule, so index it directly. + clause = params['auto_suspend'] + units = clause['suspend_after_units'].upper() + return dict( + suspend_after_seconds=( + clause['suspend_after_value'] * _SUSPEND_UNIT_SECONDS[units] + ), + suspend_type=clause['suspend_type'].upper(), + ) + + +def _update_window(params: Dict[str, Any]) -> Optional[Dict[str, int]]: + """Convert a ``WITH UPDATE WINDOW ':'`` clause to a dict.""" + if not params.get('with_update_window'): + return None + day, hour = params['with_update_window'].split(':', 1) + return dict(day=int(day), hour=int(hour)) + + +def _cluster_region(cluster: Any) -> Optional[str]: + """Return a cluster's provider region name, e.g. ``us-east-1``.""" + region = cluster.region + if region is None: + return None + return region.region_name or region.name + + +def _cluster_project_name(cluster: Any) -> Optional[str]: + """Return the name of the project a deployment belongs to.""" + project = cluster.project + if project is None: + return None + return project.name + + +def _resolve_region( + params: Dict[str, Any], + manager: ClusterManager, +) -> Dict[str, Any]: + """ + Resolve an ``IN REGION`` clause to ``create_cluster`` keywords. + + A region is identified by its ``(provider, region_name)`` pair rather than + by an ID. ``GET /v2/regions`` reports both a + display name (``region``, e.g. ``US East 1 (N. Virginia)``) and a provider + slug (``regionName``, e.g. ``us-east-1``), and a cluster's own ``region`` + field is the *slug* -- so a display name has to be translated before it is + posted. Matching accepts either spelling, case-insensitively: the display + names are mixed case, ``USING PROVIDER`` is already case-insensitive, and + ``SHOW CLUSTER REGIONS`` matches its ``LIKE`` pattern case-insensitively + too, so a name that command finds has to be a name this clause accepts. A + match is returned in the API's own spelling, not the caller's. + + An unmatched literal is passed through untouched rather than rejected: the + region list is cached, and the API gives a clearer error for an unknown + region than a stale local list can. Note what a miss costs, which is why + the match is lenient -- the provider is only ever recovered *from* a match, + so an unmatched region is posted with no provider at all unless the caller + also wrote ``USING PROVIDER``. + + Takes the caller's ``manager`` rather than building its own, because + :attr:`ClusterManager.regions` caches on the manager instance, not on the + class (see ``management.utils.TTLProperty``). With a throwaway manager here + the region list would be fetched a second time, the first being the lookup + ``Cluster.from_dict`` does on the caller's manager to fill in the display + name of the region of the cluster just created. + """ + region_name = params['in_region']['region_name'] + provider = params.get('using_provider') or None + + wanted = region_name.casefold() + matches = [ + x for x in manager.regions + if wanted in [ + y.casefold() for y in (x.name, x.region_name) if y + ] + ] + if provider: + matches = [ + x for x in matches + if (x.provider or '').upper() == provider.upper() + ] + + if len(matches) > 1: + found = ', '.join( + f'{x.provider} {x.region_name}' for x in matches + ) + raise ValueError( + f'more than one region matches "{region_name}": {found}; ' + 'use the USING PROVIDER clause to select one', + ) + + if matches: + return dict( + provider=matches[0].provider, + region=matches[0].region_name, + ) + + # Unknown to the cached region list; let the API rule on it. + return dict(provider=provider, region=region_name) + + +class ShowClustersHandler(SQLHandler): + """ + SHOW CLUSTERS [ ] + [ ] [ ] + [ ]; + + Description + ----------- + Displays information on clusters. A cluster is a single flat deployment + resource: its compute settings and its deployment-wide settings all live + on the one object. + + Arguments + --------- + * ````: A pattern similar to SQL LIKE clause. + Uses ``%`` as the wildcard character. + + Remarks + ------- + * Use the ``LIKE`` clause to specify a pattern and return only the + clusters that match the specified pattern. + * The ``LIMIT`` clause limits the number of results to the + specified number. + * Use the ``ORDER BY`` clause to sort the results by the specified + key. By default, the results are sorted in the ascending order. + * To return more information about the clusters, use the + ``EXTENDED`` clause. + + Example + ------- + The following command displays a list of clusters with names that + match the specified pattern:: + + SHOW CLUSTERS LIKE 'analytics%' EXTENDED ORDER BY Name; + + See Also + -------- + * ``SHOW STARTER CLUSTERS`` + * ``CREATE CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + res = FusionSQLResult() + res.add_field('Name', result.STRING) + res.add_field('ID', result.STRING) + res.add_field('Region', result.STRING) + res.add_field('Size', result.STRING) + res.add_field('State', result.STRING) + + if params['extended']: + res.add_field('Provider', result.STRING) + res.add_field('Endpoint', result.STRING) + res.add_field('DeploymentType', result.STRING) + res.add_field('FirewallRanges', result.JSON) + res.add_field('ProjectName', result.STRING) + res.add_field('CreatedAt', result.DATETIME) + res.add_field('TerminatedAt', result.DATETIME) + + def fields(x: Any) -> Any: + return ( + x.name, x.id, _cluster_region(x), x.size, x.state, + x.provider, x.endpoint, x.deployment_type, + json.dumps(x.firewall_ranges or []), + _cluster_project_name(x), + dt_isoformat(x.created_at), + dt_isoformat(x.terminated_at), + ) + else: + def fields(x: Any) -> Any: + # Report the provider slug, not the region's display name. + return (x.name, x.id, _cluster_region(x), x.size, x.state) + + res.set_rows([fields(x) for x in manager.clusters]) + + if params['like']: + res = res.like(Name=params['like']) + + return res.order_by(**params['order_by']).limit(params['limit']) + + +ShowClustersHandler.register(overwrite=True) + + +class ShowClusterRegionsHandler(SQLHandler): + """ + SHOW CLUSTER REGIONS [ ] + [ ] + [ ]; + + Description + ----------- + Returns the regions available for creating clusters. + + Arguments + --------- + * ````: A pattern similar to SQL LIKE clause. + Uses ``%`` as the wildcard character. + + Remarks + ------- + * Use the ``LIKE`` clause to specify a pattern and return only the + regions that match the specified pattern. + * The ``LIMIT`` clause limits the number of results to the + specified number. + * Use the ``ORDER BY`` clause to sort the results by the specified + key. By default, the results are sorted in the ascending order. + * There is no ``ID`` column. The API assigns no region IDs; a region is + identified by its provider and region name, which is why + ``CREATE CLUSTER`` has no ``IN REGION ID`` clause. + * ``Name`` is the display name, for example + ``US East 1 (N. Virginia)``. ``RegionName`` is the cloud provider's + own name for it, for example ``us-east-1``. Either may be given to + ``CREATE CLUSTER``. + + Example + ------- + The following command returns the regions in the US, sorted by name:: + + SHOW CLUSTER REGIONS LIKE 'US%' ORDER BY Name; + + See Also + -------- + * ``CREATE CLUSTER``, whose ``IN REGION`` clause takes one of these names. + * ``SHOW STARTER CLUSTER REGIONS``, the subset of these that a starter + cluster can use. + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + res = FusionSQLResult() + res.add_field('Name', result.STRING) + res.add_field('Provider', result.STRING) + res.add_field('RegionName', result.STRING) + + res.set_rows([ + (x.name, x.provider, x.region_name) + for x in manager.regions + ]) + + if params['like']: + res = res.like(Name=params['like']) + + return res.order_by(**params['order_by']).limit(params['limit']) + + +ShowClusterRegionsHandler.register(overwrite=True) + + +class ShowProjectsHandler(SQLHandler): + """ + SHOW PROJECTS [ ] + [ ] + [ ]; + + Description + ----------- + Displays the projects in the current organization. + + Arguments + --------- + * ````: A pattern similar to SQL LIKE clause. + Uses ``%`` as the wildcard character. + + Remarks + ------- + * Use the ``LIKE`` clause to specify a pattern and return only the + projects that match the specified pattern. + * The ``LIMIT`` clause limits the number of results to the + specified number. + * Use the ``ORDER BY`` clause to sort the results by the specified + key. By default, the results are sorted in the ascending order. + * Projects cannot be created or dropped from Fusion SQL. This command + exists so that the project required by ``CREATE CLUSTER`` can be + discovered. + * ``CREATE CLUSTER`` needs a project. If the organization has exactly + one, it is used automatically; otherwise name one with the + ``IN PROJECT`` clause. + + Example + ------- + The following command displays the projects in the current + organization:: + + SHOW PROJECTS ORDER BY Name; + + See Also + -------- + * ``CREATE CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + res = FusionSQLResult() + res.add_field('Name', result.STRING) + res.add_field('ID', result.STRING) + res.add_field('Edition', result.STRING) + res.add_field('CreatedAt', result.DATETIME) + + res.set_rows([ + (x.name, x.id, x.edition, dt_isoformat(x.created_at)) + for x in manager.projects + ]) + + if params['like']: + res = res.like(Name=params['like']) + + return res.order_by(**params['order_by']).limit(params['limit']) + + +ShowProjectsHandler.register(overwrite=True) + + +class CreateClusterHandler(SQLHandler): + """ + CREATE CLUSTER [ if_not_exists ] cluster_name + in_region + [ using_provider ] + [ in_project ] + [ with_size ] + [ using_scale_factor ] + [ auto_suspend ] + [ enable_kai ] + [ with_cache_config ] + [ with_firewall_ranges ] + [ allow_all_traffic ] + [ with_update_window ] + [ expires_at ] + [ wait_on_active ] + ; + + # Only create the cluster if it doesn't exist already + if_not_exists = IF NOT EXISTS + + # Name of the cluster + cluster_name = '' + + # Region to create the cluster in + in_region = IN REGION region_name + region_name = '' + + # Cloud provider, to disambiguate a region name + using_provider = USING PROVIDER '' + + # Project to create the cluster in + in_project = IN PROJECT { project_id | project_name } + project_id = ID '' + project_name = '' + + # Runtime size + with_size = WITH SIZE '' + + # Scale factor + using_scale_factor = USING SCALE FACTOR + + # Auto-suspend + auto_suspend = AUTO SUSPEND AFTER suspend_after_value suspend_after_units suspend_type + suspend_after_value = + suspend_after_units = { SECONDS | MINUTES | HOURS | DAYS } + suspend_type = WITH TYPE { IDLE | SCHEDULED | DISABLED } + + # Enable Kai + enable_kai = ENABLE KAI + + # Cache config + with_cache_config = WITH CACHE CONFIG + + # Incoming IP ranges + with_firewall_ranges = WITH FIREWALL RANGES '',... + + # Allow all incoming traffic + allow_all_traffic = ALLOW ALL TRAFFIC + + # Update window + with_update_window = WITH UPDATE WINDOW ':' + + # Datetime or interval for expiration date/time of the cluster + expires_at = EXPIRES AT '' + + # Wait for the cluster to be active before continuing + wait_on_active = WAIT ON ACTIVE + + Description + ----------- + Creates a cluster. + + Arguments + --------- + * ````: The name of the cluster. Must be 1-32 characters + of lowercase letters, digits and hyphens, and must start and end with + a letter or digit. + * ````: The display name or the cloud provider name of the + region to create the cluster in, as reported by + ``SHOW CLUSTER REGIONS``. Matched without regard to case. + * ````: The cloud provider (AWS, GCP or Azure), if the region + name alone is ambiguous. Matched without regard to case. + * ```` or ````: The ID or name of the project + to create the cluster in. + * ````: The size of the cluster in cluster size notation, for + example ``S-1``. + * ``:``: The day of the week (0-6) and the hour of the day + (0-23) when engine updates are applied. + * ````: A list of allowed IP addresses or CIDR ranges. + + Remarks + ------- + * Specify the ``IF NOT EXISTS`` clause to create the cluster only if one + with the given name does not already exist. + * ``IN PROJECT`` is optional in an organization with a single project, + which is then used automatically. In an organization with several, the + clause is required; ``SHOW PROJECTS`` lists the candidates. + * To allow incoming traffic from any IP address, use the + ``ALLOW ALL TRAFFIC`` clause. + * The ``WAIT ON ACTIVE`` clause pauses execution until the cluster + reaches the ``ACTIVE`` state. + * This command returns a row. The admin password is generated by the API + and reported when the cluster is created and at no later point, so it is + returned here; a cluster created without capturing it has no reachable + ``admin`` user. + * There are no KMS key or ``SMART DR`` clauses. The API takes no such + parameters, so those clauses would be silently dropped. + * The API's ``deploymentType`` and ``multiAZ`` are not surfaced as clauses; + reach them through ``ClusterManager.create_cluster``, which takes both. + + Example + ------- + The following command creates a cluster named **analytics** in the + ``US East 1 (N. Virginia)`` region and waits for it to become active:: + + CREATE CLUSTER 'analytics' IN REGION 'US East 1 (N. Virginia)' + WITH SIZE 'S-00' WAIT ON ACTIVE; + + See Also + -------- + * ``SHOW CLUSTERS`` + * ``SHOW CLUSTER REGIONS`` + * ``DROP CLUSTER`` + + """ # noqa: E501 + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + # Only create if a live one doesn't exist. A terminated cluster keeps + # its name in the listing, so it must not count as existing or the + # name would be unusable ever after. + if params['if_not_exists']: + live = [ + x for x in manager.clusters + if x.name == params['cluster_name'] + and x.terminated_at is None + ] + if live: + return None + + project = get_project(params) + region = _resolve_region(params, manager) + + cluster = manager.create_cluster( + params['cluster_name'], + provider=region['provider'], + region=region['region'], + size=params['with_size'], + scale_factor=params['using_scale_factor'], + firewall_ranges=params['with_firewall_ranges'], + allow_all_traffic=params['allow_all_traffic'], + auto_suspend=_auto_suspend(params), + cache_config=params['with_cache_config'], + expires_at=params['expires_at'], + update_window=_update_window(params), + kai=params['enable_kai'], + project=project, + wait_on_active=params['wait_on_active'], + ) + + res = FusionSQLResult() + res.add_field('Name', result.STRING) + res.add_field('ID', result.STRING) + res.add_field('Endpoint', result.STRING) + res.add_field('AdminPassword', result.STRING) + res.set_rows([ + ( + cluster.name, cluster.id, cluster.endpoint, + cluster.admin_password, + ), + ]) + return res + + +CreateClusterHandler.register(overwrite=True) + + +class SuspendClusterHandler(SQLHandler): + """ + SUSPEND CLUSTER cluster + [ wait_on_suspended ]; + + # Cluster + cluster = { cluster_id | cluster_name } + + # ID of the cluster + cluster_id = ID '' + + # Name of the cluster + cluster_name = '' + + # Wait for the cluster to be suspended before continuing + wait_on_suspended = WAIT ON SUSPENDED + + Description + ----------- + Suspends a cluster. + + Arguments + --------- + * ````: The ID of the cluster to suspend. + * ````: The name of the cluster to suspend. + + Remarks + ------- + * Use the ``WAIT ON SUSPENDED`` clause to pause query execution + until the cluster is in the ``SUSPENDED`` state. + * There is no ``IN GROUP`` clause. A cluster is flat, so there is no + containing group to name. + + Example + ------- + The following example suspends a cluster named **analytics**:: + + SUSPEND CLUSTER 'analytics' WAIT ON SUSPENDED; + + See Also + -------- + * ``RESUME CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + cluster = get_cluster(params) + cluster.suspend(wait_on_suspended=params['wait_on_suspended']) + return None + + +SuspendClusterHandler.register(overwrite=True) + + +class ResumeClusterHandler(SQLHandler): + """ + RESUME CLUSTER cluster + [ disable_auto_suspend ] + [ wait_on_resumed ]; + + # Cluster + cluster = { cluster_id | cluster_name } + + # ID of the cluster + cluster_id = ID '' + + # Name of the cluster + cluster_name = '' + + # Disable auto-suspend + disable_auto_suspend = DISABLE AUTO SUSPEND + + # Wait for the cluster to be resumed before continuing + wait_on_resumed = WAIT ON RESUMED + + Description + ----------- + Resumes a cluster. + + Arguments + --------- + * ````: The ID of the cluster to resume. + * ````: The name of the cluster to resume. + + Remarks + ------- + * Use the ``WAIT ON RESUMED`` clause to pause query execution + until the cluster is in the ``RESUMED`` state. + * Specify the ``DISABLE AUTO SUSPEND`` clause to disable + auto-suspend for the resumed cluster. + + Example + ------- + The following example resumes a cluster named **analytics** and + disables its auto-suspend setting:: + + RESUME CLUSTER 'analytics' DISABLE AUTO SUSPEND WAIT ON RESUMED; + + See Also + -------- + * ``SUSPEND CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + cluster = get_cluster(params) + cluster.resume( + wait_on_resumed=params['wait_on_resumed'], + disable_auto_suspend=params['disable_auto_suspend'], + ) + return None + + +ResumeClusterHandler.register(overwrite=True) + + +class DropClusterHandler(SQLHandler): + """ + DROP CLUSTER [ if_exists ] + cluster + [ wait_on_terminated ]; + + # Only run the command if the cluster exists + if_exists = IF EXISTS + + # Cluster + cluster = { cluster_id | cluster_name } + + # ID of the cluster to delete + cluster_id = ID '' + + # Name of the cluster to delete + cluster_name = '' + + # Wait for termination to complete before continuing + wait_on_terminated = WAIT ON TERMINATED + + Description + ----------- + Deletes the specified cluster. + + Arguments + --------- + * ````: The ID of the cluster to delete. + * ````: The name of the cluster to delete. + + Remarks + ------- + * Specify the ``IF EXISTS`` clause to attempt the delete operation + only if a cluster with the specified ID or name exists. + * Use the ``WAIT ON TERMINATED`` clause to pause query execution until + the cluster is in the ``TERMINATED`` state. + * There is no ``FORCE`` clause. ``DELETE /v2/clusters`` does take a + ``force`` query parameter, which ``Cluster.terminate()`` documents as + "even if it is in use", but that meaning has not been confirmed against + the live API, so the clause is withheld rather than guessed at. + * All databases attached to the cluster are detached when the cluster + is deleted. + + Example + ------- + The following example deletes a cluster named **analytics** if it + exists, waiting for the termination to finish:: + + DROP CLUSTER IF EXISTS 'analytics' WAIT ON TERMINATED; + + See Also + -------- + * ``CREATE CLUSTER`` + * ``DROP STARTER CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + try: + cluster = get_cluster(params) + if cluster.terminated_at is not None: + raise KeyError('cluster is already terminated') + cluster.terminate( + wait_on_terminated=params['wait_on_terminated'], + ) + + except KeyError: + if not params['if_exists']: + raise + + return None + + +DropClusterHandler.register(overwrite=True) + + +class UseClusterHandler(SQLHandler): + """ + USE CLUSTER cluster [ with_database ]; + + # Cluster + cluster = { cluster_id | cluster_name | current_cluster } + + # ID of the cluster + cluster_id = ID '' + + # Name of the cluster + cluster_name = '' + + # Current cluster + current_cluster = @@CURRENT + + # Name of database + with_database = WITH DATABASE '' + + Description + ----------- + Change the cluster and database in the notebook. + + Arguments + --------- + * ````: The ID of the cluster to use. + * ````: The name of the cluster to use. + * ````: The name of the database to select. + + Remarks + ------- + * If you want to specify a database in the current cluster, the + cluster name can be specified as ``@@CURRENT``. + * Specify the ``WITH DATABASE`` clause to select a default + database for the session. + * There is no ``IN GROUP`` clause. A cluster is flat, so there is no + containing group to search in. + * This command only works in a notebook session in the + Managed Service. + + Example + ------- + The following command sets the cluster to ``analytics`` and selects + ``dbname`` as the default database:: + + USE CLUSTER 'analytics' WITH DATABASE 'dbname'; + + See Also + -------- + * ``SHOW CLUSTERS`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + from singlestoredb.notebook import portal + + # Handle current cluster case + if params['cluster'].get('current_cluster'): + if params.get('with_database'): + portal.default_database = params['with_database'] + return None + + cluster_name = params['cluster'].get('cluster_name') + cluster_id = params['cluster'].get('cluster_id') + + try: + if params.get('with_database'): + portal.connection = ( + cluster_name or cluster_id, + params['with_database'], + ) + else: + portal.workspace = cluster_name or cluster_id + + except RuntimeError as exc: + if 'timeout' not in str(exc): + raise + + return None + + +UseClusterHandler.register(overwrite=True) + + +class ShowStarterClustersHandler(SQLHandler): + """ + SHOW STARTER CLUSTERS [ ] + [ ] [ ] + [ ]; + + Description + ----------- + Displays information on starter clusters, the shared-tier deployment + resource. + + Arguments + --------- + * ````: A pattern similar to SQL LIKE clause. + Uses ``%`` as the wildcard character. + + Remarks + ------- + * Use the ``LIKE`` clause to specify a pattern and return only the + starter clusters that match the specified pattern. + * The ``LIMIT`` clause limits the number of results to the + specified number. + * Use the ``ORDER BY`` clause to sort the results by the specified + key. By default, the results are sorted in the ascending order. + * To return more information about the starter clusters, use the + ``EXTENDED`` clause. + + Example + ------- + The following command displays the starter clusters, sorted by name:: + + SHOW STARTER CLUSTERS ORDER BY Name; + + See Also + -------- + * ``SHOW CLUSTERS`` + * ``CREATE STARTER CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + res = FusionSQLResult() + res.add_field('Name', result.STRING) + res.add_field('ID', result.STRING) + res.add_field('DatabaseName', result.STRING) + + if params['extended']: + res.add_field('Endpoint', result.STRING) + res.add_field('ProjectName', result.STRING) + + def fields(x: Any) -> Any: + return ( + x.name, x.id, x.database_name, + x.endpoint, _cluster_project_name(x), + ) + else: + def fields(x: Any) -> Any: + return (x.name, x.id, x.database_name) + + res.set_rows([fields(x) for x in manager.starter_clusters]) + + if params['like']: + res = res.like(Name=params['like']) + + return res.order_by(**params['order_by']).limit(params['limit']) + + +ShowStarterClustersHandler.register(overwrite=True) + + +class ShowStarterClusterRegionsHandler(SQLHandler): + """ + SHOW STARTER CLUSTER REGIONS [ ] + [ ] + [ ]; + + Description + ----------- + Returns the regions available for creating starter clusters. These are a + subset of the regions ``SHOW CLUSTER REGIONS`` reports: the shared-tier + route accepts only the regions listed here. + + Arguments + --------- + * ````: A pattern similar to SQL LIKE clause. + Uses ``%`` as the wildcard character. + + Remarks + ------- + * Use the ``LIKE`` clause to specify a pattern and return only the + regions that match the specified pattern. + * The ``LIMIT`` clause limits the number of results to the + specified number. + * Use the ``ORDER BY`` clause to sort the results by the specified + key. By default, the results are sorted in the ascending order. + * The columns are those of ``SHOW CLUSTER REGIONS``: ``Name`` is the + display name, for example ``US East 1 (N. Virginia)``, and + ``RegionName`` is the cloud provider's own name for it, for example + ``us-east-1``. + * ``CREATE STARTER CLUSTER`` needs both the ``RegionName`` and the + ``Provider`` from this listing, and accepts no region ID. + + Example + ------- + The following command returns the starter cluster regions in the US, + sorted by name:: + + SHOW STARTER CLUSTER REGIONS LIKE 'US%' ORDER BY Name; + + See Also + -------- + * ``CREATE STARTER CLUSTER`` + * ``SHOW CLUSTER REGIONS``, the regions a full cluster can use. + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + res = FusionSQLResult() + res.add_field('Name', result.STRING) + res.add_field('Provider', result.STRING) + res.add_field('RegionName', result.STRING) + + res.set_rows([ + (x.name, x.provider, x.region_name) + for x in manager.shared_tier_regions + ]) + + if params['like']: + res = res.like(Name=params['like']) + + return res.order_by(**params['order_by']).limit(params['limit']) + + +ShowStarterClusterRegionsHandler.register(overwrite=True) + + +class CreateStarterClusterHandler(SQLHandler): + """ + CREATE STARTER CLUSTER [ if_not_exists ] cluster_name + with_database + in_region + using_provider + ; + + # Only create the starter cluster if it doesn't exist already + if_not_exists = IF NOT EXISTS + + # Name of the starter cluster + cluster_name = '' + + # Database to create in the starter cluster + with_database = WITH DATABASE '' + + # Region to create the starter cluster in + in_region = IN REGION '' + + # Cloud provider to create the starter cluster in + using_provider = USING PROVIDER '' + + Description + ----------- + Creates a starter cluster, the shared-tier deployment resource. + + Arguments + --------- + * ````: The name of the starter cluster. + * ````: The name of the database to create in it. + * ````: The cloud provider name of the region, for + example ``us-east-1``. Unlike ``CREATE CLUSTER``, this is sent to the + API as written rather than matched against a region listing, so it must + be spelled exactly as ``SHOW STARTER CLUSTER REGIONS`` reports it, + including case. + * ````: The cloud provider: AWS, GCP or Azure. Any + capitalization is accepted. + + Remarks + ------- + * Specify the ``IF NOT EXISTS`` clause to create the starter cluster + only if one with the given name does not already exist. + * Not every region supports starter clusters. Only the regions reported + by ``SHOW STARTER CLUSTER REGIONS`` are accepted -- that is a subset of + ``SHOW CLUSTER REGIONS``, so a region valid for ``CREATE CLUSTER`` is + not necessarily valid here. Both the provider and the region name are + required because there is nothing to infer them from. + + Example + ------- + The following command creates a starter cluster named **scratch** with + a database named **scratchdb**:: + + CREATE STARTER CLUSTER 'scratch' WITH DATABASE 'scratchdb' + IN REGION 'us-east-1' USING PROVIDER 'AWS'; + + See Also + -------- + * ``SHOW STARTER CLUSTER REGIONS`` + * ``SHOW STARTER CLUSTERS`` + * ``DROP STARTER CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + manager = get_cluster_manager() + + if params['if_not_exists']: + try: + get_starter_cluster( + {'cluster_name': params['cluster_name']}, + ) + return None + except (ValueError, KeyError): + pass + + manager.create_starter_cluster( + params['cluster_name'], + database_name=params['with_database'], + provider=params['using_provider'], + region=params['in_region'], + ) + + return None + + +CreateStarterClusterHandler.register(overwrite=True) + + +class DropStarterClusterHandler(SQLHandler): + """ + DROP STARTER CLUSTER [ if_exists ] cluster; + + # Only run the command if the starter cluster exists + if_exists = IF EXISTS + + # Starter cluster + cluster = { cluster_id | cluster_name } + + # ID of the starter cluster to delete + cluster_id = ID '' + + # Name of the starter cluster to delete + cluster_name = '' + + Description + ----------- + Deletes the specified starter cluster. + + Arguments + --------- + * ````: The ID of the starter cluster to delete. + * ````: The name of the starter cluster to delete. + + Remarks + ------- + * Specify the ``IF EXISTS`` clause to attempt the delete operation + only if a starter cluster with the specified ID or name exists. + * There is no ``WAIT ON TERMINATED`` clause. The shared-tier + termination route reports no state to wait on. + + Example + ------- + The following example deletes a starter cluster named **scratch** if + it exists:: + + DROP STARTER CLUSTER IF EXISTS 'scratch'; + + See Also + -------- + * ``CREATE STARTER CLUSTER`` + * ``DROP CLUSTER`` + + """ + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + try: + get_starter_cluster(params).terminate() + + except KeyError: + if not params['if_exists']: + raise + + return None + + +DropStarterClusterHandler.register(overwrite=True) diff --git a/singlestoredb/fusion/handlers/export.py b/singlestoredb/fusion/handlers/export.py index 7a879b4da..6416c0e03 100644 --- a/singlestoredb/fusion/handlers/export.py +++ b/singlestoredb/fusion/handlers/export.py @@ -1,4 +1,25 @@ #!/usr/bin/env python3 +""" +Fusion SQL handlers for the table egress (EXPORT) service. + +Pinned to management API v2, so an export is owned by a +:class:`~singlestoredb.management.v2.cluster.Cluster`. The version is named at +the import line rather than left to the ``management.version`` option, for the +same reason ``handlers/utils.py`` pins its managers: the egress routes differ by +version (``clusters/{id}/egress/...`` at v2 against +``workspaceGroups/{id}/egress/...`` at v1) and these handlers must not follow an +unrelated option onto the other one. + +Every handler here resolves its target with ``get_cluster({})``, which reads +``SINGLESTOREDB_WORKSPACE``. At v1 it was ``get_workspace_group({})``, reading +``SINGLESTOREDB_WORKSPACE_GROUP`` -- so the environment variable that names the +export target changed with the version. There is deliberately no ``IN`` clause: +none of these commands took a target clause at v1 either, and adding one is a +grammar change rather than part of the version move. + +All handlers are hidden (``_enabled = False``), so they only register under +``SINGLESTOREDB_FUSION_ENABLE_HIDDEN``. +""" import datetime import json from typing import Any @@ -6,12 +27,12 @@ from typing import Optional from .. import result -from ...management.export import _get_exports -from ...management.export import ExportService -from ...management.export import ExportStatus +from ...management.v2.export import _get_exports +from ...management.v2.export import ExportService +from ...management.v2.export import ExportStatus from ..handler import SQLHandler from ..result import FusionSQLResult -from .utils import get_workspace_group +from .utils import get_cluster class CreateClusterIdentity(SQLHandler): @@ -82,13 +103,10 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: storage_config['provider'] = 'S3' - wsg = get_workspace_group({}) - - if wsg._manager is None: - raise TypeError('no workspace manager is associated with workspace group') + cluster = get_cluster({}) out = ExportService( - wsg, + cluster, 'none', 'none', dict(**catalog_config, **catalog_creds), @@ -124,14 +142,11 @@ def _start_export(params: Dict[str, Any]) -> Optional[FusionSQLResult]: storage_config['provider'] = 'S3' - wsg = get_workspace_group({}) + cluster = get_cluster({}) if from_database is None: raise ValueError('database name must be specified for source table') - if wsg._manager is None: - raise TypeError('no workspace manager is associated with workspace group') - partition_by = [] if params['partition_by']: for key in params['partition_by']: @@ -178,7 +193,7 @@ def _start_export(params: Dict[str, Any]) -> Optional[FusionSQLResult]: raise ValueError('invalid refresh interval time unit') out = ExportService( - wsg, + cluster, from_database, from_table, dict(**catalog_config, **catalog_creds), @@ -420,9 +435,9 @@ class ShowExport(SQLHandler): _enabled = False def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: - wsg = get_workspace_group({}) + cluster = get_cluster({}) return _format_status( - params['export_id'], ExportStatus(params['export_id'], wsg), + params['export_id'], ExportStatus(params['export_id'], cluster), ) @@ -441,21 +456,25 @@ class ShowExports(SQLHandler): _enabled = False def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: - wsg = get_workspace_group({}) + cluster = get_cluster({}) - exports = _get_exports(wsg, params.get('scope', 'all')) + exports = _get_exports(cluster, params.get('scope', 'all')) res = FusionSQLResult() res.add_field('ExportID', result.STRING) res.add_field('Status', result.STRING) res.add_field('Message', result.STRING) + # The ID comes from the ExportStatus rather than from its ``_info()`` + # body: ``_info()`` is a per-export status GET, which is not documented + # to echo ``egressID`` back. ``_get_exports`` already read the ID from + # the listing to build each object, so it is known here either way. res.set_rows([ ( - info['egressID'], + x.export_id, info.get('status', 'Unknown'), info.get('statusMsg', ''), ) - for info in [x._info() for x in exports] + for x, info in [(x, x._info()) for x in exports] ]) return res @@ -476,8 +495,8 @@ class SuspendExport(SQLHandler): _enabled = False def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: - wsg = get_workspace_group({}) - service = ExportService.from_export_id(wsg, params['export_id']) + cluster = get_cluster({}) + service = ExportService.from_export_id(cluster, params['export_id']) return _format_status(params['export_id'], service.suspend()) @@ -496,8 +515,8 @@ class ResumeExport(SQLHandler): _enabled = False def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: - wsg = get_workspace_group({}) - service = ExportService.from_export_id(wsg, params['export_id']) + cluster = get_cluster({}) + service = ExportService.from_export_id(cluster, params['export_id']) return _format_status(params['export_id'], service.resume()) @@ -516,8 +535,8 @@ class DropExport(SQLHandler): _enabled = False def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: - wsg = get_workspace_group({}) - service = ExportService.from_export_id(wsg, params['export_id']) + cluster = get_cluster({}) + service = ExportService.from_export_id(cluster, params['export_id']) service.drop() return None diff --git a/singlestoredb/fusion/handlers/files.py b/singlestoredb/fusion/handlers/files.py index 7f848611b..4136fad54 100644 --- a/singlestoredb/fusion/handlers/files.py +++ b/singlestoredb/fusion/handlers/files.py @@ -198,9 +198,11 @@ class UploadFileHandler(SQLHandler): def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: file_space = get_file_space(params) - file_space.upload_file( + # Nothing here reads the uploaded file's metadata, so don't pay the + # request that fetching it costs. + file_space._upload_local_file( params['local_path'], params['path'], - overwrite=params['overwrite'], + overwrite=params['overwrite'], fetch_info=False, ) return None diff --git a/singlestoredb/fusion/handlers/job.py b/singlestoredb/fusion/handlers/job.py index 9da298f3d..8d23fef8e 100644 --- a/singlestoredb/fusion/handlers/job.py +++ b/singlestoredb/fusion/handlers/job.py @@ -10,7 +10,7 @@ from ..handler import SQLHandler from ..result import FusionSQLResult from .utils import dt_isoformat -from .utils import get_workspace_manager +from .utils import get_cluster_manager from singlestoredb.management.job import Mode @@ -128,7 +128,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res = FusionSQLResult() res.add_field('JobID', result.STRING) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs parameters = None if params.get('with_parameters'): @@ -228,7 +228,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res = FusionSQLResult() res.add_field('JobID', result.STRING) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs parameters = None if params.get('with_parameters'): @@ -290,7 +290,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res = FusionSQLResult() res.add_field('Success', result.BOOL) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs timeout_in_secs = None if params.get('with_timeout'): @@ -367,7 +367,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res.add_field('TargetID', result.STRING) res.add_field('TargetType', result.STRING) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs jobs = [] for job_id in params['job_ids']: @@ -496,7 +496,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res.add_field('StartedAt', result.DATETIME) res.add_field('FinishedAt', result.DATETIME) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs executionsData = jobs_manager.get_executions( params['job_id'], @@ -562,7 +562,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res.add_field('Value', result.STRING) res.add_field('Type', result.STRING) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs parameters = jobs_manager.get_parameters(params['job_id']) @@ -601,7 +601,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res.add_field('Name', result.STRING) res.add_field('Description', result.STRING) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs runtimes = jobs_manager.runtimes() @@ -646,7 +646,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res.add_field('JobID', result.STRING) res.add_field('Success', result.BOOL) - jobs_manager = get_workspace_manager().organizations.current.jobs + jobs_manager = get_cluster_manager().organizations.current.jobs results: List[Tuple[Any, ...]] = [] for job_id in params['job_ids']: diff --git a/singlestoredb/fusion/handlers/models.py b/singlestoredb/fusion/handlers/models.py index 8bb618d7a..722048e2b 100644 --- a/singlestoredb/fusion/handlers/models.py +++ b/singlestoredb/fusion/handlers/models.py @@ -5,6 +5,7 @@ from typing import Optional from .. import result +from ...management.utils import normalize_remote_path from ..handler import SQLHandler from ..result import FusionSQLResult from .files import ShowFilesHandler @@ -13,6 +14,14 @@ from .utils import get_inference_api_manager +# Every handler in this module is hidden -- ``_enabled = False``, so it only +# registers when SINGLESTOREDB_FUSION_ENABLE_HIDDEN is set. The models surface +# is v1-only: the START/STOP/SHOW/DROP MODEL commands ride the ``inferenceapis/`` +# routes, which do not exist past v1, and the CUSTOM MODEL commands are the same +# generation of API. Rather than let the grammar advertise commands that have no +# v2 equivalent, none of them are registered by default. + + class ShowCustomModelsHandler(ShowFilesHandler): """ SHOW CUSTOM MODELS @@ -70,6 +79,8 @@ class ShowCustomModelsHandler(ShowFilesHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: params['file_location'] = 'MODELS' @@ -122,6 +133,8 @@ class UploadCustomModelHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: params['file_location'] = 'MODELS' @@ -130,17 +143,23 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: file_space = get_file_space(params) + # Remote paths always use '/', so they can't be built with os.path.join if os.path.isdir(local_path): file_space.upload_folder( local_path=local_path, - path=os.path.join(model_name, ''), + path=model_name, overwrite=params['overwrite'], ) else: - file_space.upload_file( + # Nothing here reads the uploaded file's metadata, so don't pay + # the request that fetching it costs. + file_space._upload_local_file( local_path=local_path, - path=os.path.join(model_name, local_path), + path=normalize_remote_path( + f'{model_name}/{os.path.basename(local_path)}', + ), overwrite=params['overwrite'], + fetch_info=False, ) return None @@ -199,6 +218,8 @@ class DownloadCustomModelHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: params['file_location'] = 'MODELS' @@ -206,7 +227,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: model_name = params['model_name'] file_space.download_folder( - path=os.path.join(model_name, ''), + path=model_name, local_path=params['local_path'] or model_name, overwrite=params['overwrite'], ) @@ -240,9 +261,12 @@ class DropCustomModelHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: params['file_location'] = 'MODELS' - path = os.path.join(params['model_name'], '') + # Remote paths always use '/', so they can't be built with os.path.join + path = normalize_remote_path(params['model_name']) + '/' file_space = get_file_space(params) file_space.removedirs(path=path) @@ -281,6 +305,8 @@ class StartModelHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: inference_api = get_inference_api(params) operation_result = inference_api.start() @@ -329,6 +355,8 @@ class StopModelHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: inference_api = get_inference_api(params) operation_result = inference_api.stop() @@ -371,6 +399,8 @@ class ShowModelsHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: inference_api_manager = get_inference_api_manager() models = inference_api_manager.show() @@ -422,6 +452,8 @@ class DropModelHandler(SQLHandler): """ # noqa: E501 + _enabled = False + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: inference_api = get_inference_api(params) operation_result = inference_api.drop() diff --git a/singlestoredb/fusion/handlers/stage.py b/singlestoredb/fusion/handlers/stage.py index 6cbd4cd6a..87560685d 100644 --- a/singlestoredb/fusion/handlers/stage.py +++ b/singlestoredb/fusion/handlers/stage.py @@ -1,4 +1,21 @@ #!/usr/bin/env python3 +""" +Fusion SQL handlers for Stage. + +Every handler names its Stage owner through the same ``in`` clause. A bare +``IN`` is the spelling to use, and it needs no keyword to say what kind of +owner it names: the value is resolved as a deployment against management API +v2, and failing that as a workspace group against v1, where Stage is attached +to the group rather than to a workspace. Both are silent, because naming a +group this way is what Stage statements always did -- a group was the only kind +of Stage owner before v2 -- and which kind a given name belongs to is a fact +about the org rather than about the statement. + +``IN GROUP`` names a workspace group explicitly, and is the one deprecated +spelling here: it goes away with ``management/v1/``, and dropping the keyword +is an edit that works today either way. :func:`.utils.get_deployment` resolves +all of this, and everything it can return exposes ``.stage``. +""" from typing import Any from typing import Dict from typing import Optional @@ -19,7 +36,7 @@ class ShowStageFilesHandler(SQLHandler): # Deployment in = { in_group | in_deployment } - in_group = IN GROUP { deployment_id | deployment_name } + in_group = IN GROUP { group_id | group_name } in_deployment = IN { deployment_id | deployment_name } # ID of deployment @@ -28,6 +45,12 @@ class ShowStageFilesHandler(SQLHandler): # Name of deployment deployment_name = '' + # ID of workspace group + group_id = ID '' + + # Name of workspace group + group_name = '' + # Stage path to list at_path = AT '' @@ -50,6 +73,10 @@ class ShowStageFilesHandler(SQLHandler): the Stage is attached. * ````: The name of the deployment in which which the Stage is attached. + * ````: The ID of the workspace group in which the + Stage is attached. + * ````: The name of the workspace group in which + the Stage is attached. * ````: A path in the Stage. * ````: A pattern similar to SQL LIKE clause. Uses ``%`` as the wildcard character. @@ -64,8 +91,13 @@ class ShowStageFilesHandler(SQLHandler): key. By default, the results are sorted in the ascending order. * The ``AT`` clause specifies the path in the Stage to list the files from. - * The ``IN`` clause specifies the ID or the name of the - deployment in which the Stage is attached. + * The ``IN`` clause specifies the ID or the name of the deployment -- + or, for a Stage that has not moved off one, the workspace group -- + in which the Stage is attached. + * The ``IN GROUP`` clause names a workspace group explicitly. It is + deprecated and goes away with management API v1, which is the version + workspace groups belong to: drop the ``GROUP`` keyword, since a bare + ``IN`` resolves a workspace group too. * Use the ``RECURSIVE`` clause to list the files recursively. * To return more information about the files, use the ``EXTENDED`` clause. @@ -142,7 +174,7 @@ class UploadStageFileHandler(SQLHandler): # Deployment in = { in_group | in_deployment } - in_group = IN GROUP { deployment_id | deployment_name } + in_group = IN GROUP { group_id | group_name } in_deployment = IN { deployment_id | deployment_name } # ID of deployment @@ -151,6 +183,12 @@ class UploadStageFileHandler(SQLHandler): # Name of deployment deployment_name = '' + # ID of workspace group + group_id = ID '' + + # Name of workspace group + group_name = '' + # Path to local file local_path = '' @@ -171,13 +209,22 @@ class UploadStageFileHandler(SQLHandler): is attached. * ````: The name of the deployment in which which the Stage is attached. + * ````: The ID of the workspace group in which the + Stage is attached. + * ````: The name of the workspace group in which + the Stage is attached. * ````: The path to the file to upload in the local directory. Remarks ------- - * The ``IN`` clause specifies the ID or the name of the workspace - group in which the Stage is attached. + * The ``IN`` clause specifies the ID or the name of the deployment -- + or, for a Stage that has not moved off one, the workspace group -- + in which the Stage is attached. + * The ``IN GROUP`` clause names a workspace group explicitly. It is + deprecated and goes away with management API v1, which is the version + workspace groups belong to: drop the ``GROUP`` keyword, since a bare + ``IN`` resolves a workspace group too. * If the ``OVERWRITE`` clause is specified, any existing file at the specified path in the Stage is overwritten. @@ -197,9 +244,11 @@ class UploadStageFileHandler(SQLHandler): def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: wg = get_deployment(params) - wg.stage.upload_file( + # Nothing here reads the uploaded file's metadata, so don't pay the + # request that fetching it costs. + wg.stage._upload_local_file( params['local_path'], params['stage_path'], - overwrite=params['overwrite'], + overwrite=params['overwrite'], fetch_info=False, ) return None @@ -220,7 +269,7 @@ class DownloadStageFileHandler(SQLHandler): # Deployment in = { in_group | in_deployment } - in_group = IN GROUP { deployment_id | deployment_name } + in_group = IN GROUP { group_id | group_name } in_deployment = IN { deployment_id | deployment_name } # ID of deployment @@ -229,6 +278,12 @@ class DownloadStageFileHandler(SQLHandler): # Name of deployment deployment_name = '' + # ID of workspace group + group_id = ID '' + + # Name of workspace group + group_name = '' + # Path to local file local_path = TO '' @@ -252,6 +307,10 @@ class DownloadStageFileHandler(SQLHandler): Stage is attached. * ````: The name of the deployment in which which the Stage is attached. + * ````: The ID of the workspace group in which the + Stage is attached. + * ````: The name of the workspace group in which + the Stage is attached. * ````: The encoding to apply to the downloaded file. * ````: Specifies the path in the local directory where the file is downloaded. @@ -260,8 +319,13 @@ class DownloadStageFileHandler(SQLHandler): ------- * If the ``OVERWRITE`` clause is specified, any existing file at the download location is overwritten. - * The ``IN`` clause specifies the ID or the name of the - deployment in which the Stage is attached. + * The ``IN`` clause specifies the ID or the name of the deployment -- + or, for a Stage that has not moved off one, the workspace group -- + in which the Stage is attached. + * The ``IN GROUP`` clause names a workspace group explicitly. It is + deprecated and goes away with management API v1, which is the version + workspace groups belong to: drop the ``GROUP`` keyword, since a bare + ``IN`` resolves a workspace group too. * By default, files are downloaded in binary encoding. To view the contents of the file on the standard output, use the ``ENCODING`` clause and specify an encoding. @@ -322,7 +386,7 @@ class DropStageFileHandler(SQLHandler): # Deployment in = { in_group | in_deployment } - in_group = IN GROUP { deployment_id | deployment_name } + in_group = IN GROUP { group_id | group_name } in_deployment = IN { deployment_id | deployment_name } # ID of deployment @@ -331,6 +395,12 @@ class DropStageFileHandler(SQLHandler): # Name of deployment deployment_name = '' + # ID of workspace group + group_id = ID '' + + # Name of workspace group + group_name = '' + Description ----------- Deletes a file from a Stage. @@ -345,11 +415,20 @@ class DropStageFileHandler(SQLHandler): Stage is attached. * ````: The name of the deployment in which which the Stage is attached. + * ````: The ID of the workspace group in which the + Stage is attached. + * ````: The name of the workspace group in which + the Stage is attached. Remarks ------- - * The ``IN`` clause specifies the ID or the name of the - deployment in which the Stage is attached. + * The ``IN`` clause specifies the ID or the name of the deployment -- + or, for a Stage that has not moved off one, the workspace group -- + in which the Stage is attached. + * The ``IN GROUP`` clause names a workspace group explicitly. It is + deprecated and goes away with management API v1, which is the version + workspace groups belong to: drop the ``GROUP`` keyword, since a bare + ``IN`` resolves a workspace group too. Example -------- @@ -384,7 +463,7 @@ class DropStageFolderHandler(SQLHandler): # Deployment in = { in_group | in_deployment } - in_group = IN GROUP { deployment_id | deployment_name } + in_group = IN GROUP { group_id | group_name } in_deployment = IN { deployment_id | deployment_name } # ID of deployment @@ -393,6 +472,12 @@ class DropStageFolderHandler(SQLHandler): # Name of deployment deployment_name = '' + # ID of workspace group + group_id = ID '' + + # Name of workspace group + group_name = '' + # Should folders be deleted recursively? recursive = RECURSIVE @@ -410,11 +495,22 @@ class DropStageFolderHandler(SQLHandler): Stage is attached. * ````: The name of the deployment in which which the Stage is attached. + * ````: The ID of the workspace group in which the + Stage is attached. + * ````: The name of the workspace group in which + the Stage is attached. Remarks ------- * The ``RECURSIVE`` clause indicates that the specified folder is deleted recursively. + * The ``IN`` clause specifies the ID or the name of the deployment -- + or, for a Stage that has not moved off one, the workspace group -- + in which the Stage is attached. + * The ``IN GROUP`` clause names a workspace group explicitly. It is + deprecated and goes away with management API v1, which is the version + workspace groups belong to: drop the ``GROUP`` keyword, since a bare + ``IN`` resolves a workspace group too. Example ------- @@ -449,7 +545,7 @@ class CreateStageFolderHandler(SQLHandler): # Deployment in = { in_group | in_deployment } - in_group = IN GROUP { deployment_id | deployment_name } + in_group = IN GROUP { group_id | group_name } in_deployment = IN { deployment_id | deployment_name } # ID of deployment @@ -458,6 +554,12 @@ class CreateStageFolderHandler(SQLHandler): # Name of deployment deployment_name = '' + # ID of workspace group + group_id = ID '' + + # Name of workspace group + group_name = '' + # Path to stage folder stage_path = '' @@ -476,13 +578,22 @@ class CreateStageFolderHandler(SQLHandler): the Stage is attached. * ````: The name of the deployment in which the Stage is attached. + * ````: The ID of the workspace group in which the + Stage is attached. + * ````: The name of the workspace group in which + the Stage is attached. Remarks ------- * If the ``OVERWRITE`` clause is specified, any existing folder at the specified path is overwritten. - * The ``IN`` clause specifies the ID or the name of - the deployment in which the Stage is attached. + * The ``IN`` clause specifies the ID or the name of the deployment -- + or, for a Stage that has not moved off one, the workspace group -- + in which the Stage is attached. + * The ``IN GROUP`` clause names a workspace group explicitly. It is + deprecated and goes away with management API v1, which is the version + workspace groups belong to: drop the ``GROUP`` keyword, since a bare + ``IN`` resolves a workspace group too. Example ------- diff --git a/singlestoredb/fusion/handlers/utils.py b/singlestoredb/fusion/handlers/utils.py index f5e82b039..c1cc9e2cd 100644 --- a/singlestoredb/fusion/handlers/utils.py +++ b/singlestoredb/fusion/handlers/utils.py @@ -1,33 +1,71 @@ #!/usr/bin/env python import datetime import os +import warnings from typing import Any from typing import Dict from typing import Optional +from typing import Tuple from typing import Union from ...exceptions import ManagementError from ...management import files as mgmt_files -from ...management import manage_workspaces +from ...management.cluster import Cluster +from ...management.cluster import ClusterManager +from ...management.cluster import manage_clusters +from ...management.cluster import Project +from ...management.cluster import StarterCluster from ...management.files import FilesManager from ...management.files import FileSpace from ...management.files import manage_files -from ...management.inference_api import InferenceAPIInfo -from ...management.inference_api import InferenceAPIManager +from ...management.utils import get_cluster_id +from ...management.utils import get_workspace_id +from ...management.v1.inference_api import InferenceAPIInfo +from ...management.v1.inference_api import InferenceAPIManager +from ...management.workspace import _manage_workspaces_v1 from ...management.workspace import StarterWorkspace from ...management.workspace import Workspace from ...management.workspace import WorkspaceGroup from ...management.workspace import WorkspaceManager +from ...warnings import DeprecatedFeatureWarning def get_workspace_manager() -> WorkspaceManager: - """Return a new workspace manager.""" - return manage_workspaces() + """ + Return a new workspace manager. + + Pinned to v1. The ``WORKSPACE`` and ``WORKSPACE GROUP`` commands are the + v1 vocabulary -- v2 replaced both with the flat ``Cluster`` -- so they must + not follow the ``management.version`` option out of v1. The v2 equivalent + is :func:`get_cluster_manager`. + """ + return _manage_workspaces_v1() + + +def get_cluster_manager() -> ClusterManager: + """ + Return a new cluster manager. + + Pinned to v2 for the mirror image of the reason + :func:`get_workspace_manager` is pinned to v1: the ``CLUSTER`` commands + *are* the v2 vocabulary, so they must not follow the ``management.version`` + option out of v2 -- at v1 there is no cluster resource at all. + """ + return manage_clusters(version='v2') def get_files_manager() -> FilesManager: - """Return a new files manager.""" - return manage_files() + """ + Return a new files manager. + + Pinned to v2. ``management/files.py`` is version-neutral -- the personal, + shared and models spaces are the same resource at both versions and only + the URL differs -- so the pin is about which URL the Fusion FILES commands + address, not about which implementation they get. It is explicit rather + than left to the ``management.version`` option so that the FILES commands + do not change which API they talk to when an unrelated option is set. + """ + return manage_files(version='v2') def dt_isoformat(dt: Optional[datetime.datetime]) -> Optional[str]: @@ -51,7 +89,11 @@ def get_workspace_group(params: Dict[str, Any]) -> WorkspaceGroup: * params['in_group']['group_name'] * params['in_group']['group_id'] - Or, from the SINGLESTOREDB_WORKSPACE_GROUP environment variable. + Or, from the SINGLESTOREDB_WORKSPACE_GROUP environment variable, which the + notebook environment sets to the deployment's group ID. This resolves it + against v1, where a group is a resource in its own right; at v2 the same ID + is only reported back as ``Cluster.group`` and cannot be looked up, which + is why :func:`get_deployment` refuses it rather than guessing. """ manager = get_workspace_manager() @@ -102,9 +144,6 @@ def get_workspace_group(params: Dict[str, Any]) -> WorkspaceGroup: ) raise - if os.environ.get('SINGLESTOREDB_CLUSTER'): - raise ValueError('clusters and shared workspaces are not currently supported') - raise KeyError('no workspace group was specified') @@ -120,7 +159,11 @@ def get_workspace(params: Dict[str, Any]) -> Workspace: * params['workspace']['workspace_name'] * params['workspace']['workspace_id'] - Or, from the SINGLESTOREDB_WORKSPACE environment variable. + Or, from the SINGLESTOREDB_WORKSPACE environment variable, which the + notebook environment sets to the current deployment. Its value is a + *workspace* ID only in a v1 environment; from v2 onward the same variable + carries a cluster ID, which these v1 commands cannot resolve -- use the + ``CLUSTER`` commands, or :func:`get_cluster`, there. """ manager = get_workspace_manager() @@ -156,147 +199,542 @@ def get_workspace(params: Dict[str, Any]) -> Workspace: raise KeyError(f'no workspace found with ID: {workspace_id}') raise - if os.environ.get('SINGLESTOREDB_WORKSPACE'): + from_env = get_workspace_id() + if from_env: try: - return manager.get_workspace( - os.environ['SINGLESTOREDB_WORKSPACE'], - ) + return manager.get_workspace(from_env) except ManagementError as exc: if exc.errno == 404: raise KeyError( - 'no workspace found with ID: ' - f'{os.environ["SINGLESTOREDB_WORKSPACE"]}', + f'no workspace found with ID: {from_env}', ) raise - if os.environ.get('SINGLESTOREDB_CLUSTER'): - raise ValueError('clusters and shared workspaces are not currently supported') - raise KeyError('no workspace was specified') +def _is_missing(exc: ManagementError) -> bool: + """ + Return True if ``exc`` means "no such deployment". + + A well-formed but unknown ID draws ``404``, but a *malformed* one draws + ``400 uuid: incorrect UUID length`` from the v2 routes, which v1's + non-UUID IDs never did. Both mean the caller named something that does not + exist, so both become a ``KeyError`` rather than leaking a raw 400 for + what is usually a typo. Other 400s -- a real request-body problem -- are + left alone. + """ + if exc.errno == 404: + return True + return exc.errno == 400 and 'uuid' in str(exc.msg or '').lower() + + +def get_cluster(params: Dict[str, Any]) -> Cluster: + """ + Retrieve the specified cluster. + + The v2 counterpart of :func:`get_workspace`, and flat where that one is + nested: a cluster has no containing group, so there is nothing to resolve + first. + + This function will get a cluster name or ID from the following parameters: + + * params['cluster_name'] + * params['cluster_id'] + * params['cluster']['cluster_name'] + * params['cluster']['cluster_id'] + + Or, from ``SINGLESTOREDB_WORKSPACE``, which is what the notebook + environment calls the current deployment whatever the API version calls it. + + """ + manager = get_cluster_manager() + + cluster_name = params.get('cluster_name') or \ + (params.get('cluster') or {}).get('cluster_name') + if cluster_name: + clusters = [x for x in manager.clusters if x.name == cluster_name] + + if not clusters: + raise KeyError(f'no cluster found with name: {cluster_name}') + + if len(clusters) > 1: + ids = ', '.join(x.id for x in clusters) + raise ValueError( + f'more than one cluster with given name was found: {ids}', + ) + + return clusters[0] + + cluster_id = params.get('cluster_id') or \ + (params.get('cluster') or {}).get('cluster_id') + if cluster_id: + try: + return manager.get_cluster(cluster_id) + except ManagementError as exc: + if _is_missing(exc): + raise KeyError(f'no cluster found with ID: {cluster_id}') + raise + + from_env = get_cluster_id() + if from_env: + try: + return manager.get_cluster(from_env) + except ManagementError as exc: + if _is_missing(exc): + raise KeyError( + f'no cluster found with ID: {from_env} ' + '(from SINGLESTOREDB_WORKSPACE)', + ) + raise + + raise KeyError('no cluster was specified') + + +def get_starter_cluster(params: Dict[str, Any]) -> StarterCluster: + """ + Retrieve the specified starter cluster. + + This function will get a starter cluster name or ID from the following + parameters: + + * params['cluster_name'] + * params['cluster_id'] + * params['cluster']['cluster_name'] + * params['cluster']['cluster_id'] + + """ + manager = get_cluster_manager() + + cluster_name = params.get('cluster_name') or \ + (params.get('cluster') or {}).get('cluster_name') + if cluster_name: + clusters = [ + x for x in manager.starter_clusters + if x.name == cluster_name + ] + + if not clusters: + raise KeyError( + f'no starter cluster found with name: {cluster_name}', + ) + + if len(clusters) > 1: + ids = ', '.join(x.id for x in clusters) + raise ValueError( + 'more than one starter cluster with given name was ' + f'found: {ids}', + ) + + return clusters[0] + + cluster_id = params.get('cluster_id') or \ + (params.get('cluster') or {}).get('cluster_id') + if cluster_id: + try: + return manager.get_starter_cluster(cluster_id) + except ManagementError as exc: + if _is_missing(exc): + raise KeyError( + f'no starter cluster found with ID: {cluster_id}', + ) + raise + + raise KeyError('no starter cluster was specified') + + +def get_project(params: Dict[str, Any]) -> Optional[Project]: + """ + Resolve an ``IN PROJECT`` clause. + + Returns ``None`` when the clause is absent, so that ``CREATE CLUSTER`` falls + through to ``ClusterManager._resolve_project_id``, which reads the project + off the deployment the command is running in, else picks the organization's + only project, else raises naming the candidates. The clause is therefore + needed only to override that, or in an organization with several projects + reached from outside a deployment. + + This function will get a project name or ID from the following parameters: + + * params['project_name'] + * params['project_id'] + * params['in_project']['project_name'] + * params['in_project']['project_id'] + + """ + project_name = params.get('project_name') or \ + (params.get('in_project') or {}).get('project_name') + project_id = params.get('project_id') or \ + (params.get('in_project') or {}).get('project_id') + + if not project_name and not project_id: + return None + + manager = get_cluster_manager() + + if project_name: + projects = [x for x in manager.projects if x.name == project_name] + + if not projects: + raise KeyError(f'no project found with name: {project_name}') + + if len(projects) > 1: + ids = ', '.join(x.id for x in projects) + raise ValueError( + f'more than one project with given name was found: {ids}', + ) + + return projects[0] + + assert project_id is not None + try: + return manager.get_project(project_id) + except ManagementError as exc: + if _is_missing(exc): + raise KeyError(f'no project found with ID: {project_id}') + raise + + +# +# The parameter keys :func:`get_deployment` accepts, in resolution order. An +# empty path means the value sits directly on ``params``. The ``GROUP`` +# spellings are not here: they name a different resource and are resolved by +# :func:`_get_stage_group` before any of these are consulted. A value that +# arrives through one of these can still end up at a workspace group, but only +# as the fallback in :func:`_group_fallback`. +# +_DEPLOYMENT_KEYS: Tuple[Tuple[str, ...], ...] = ( + (), + ('in_deployment',), + ('in', 'in_deployment'), +) + +# +# The parameter keys the ``IN GROUP`` spelling arrives under, in resolution +# order. These carry ``group_id``/``group_name`` rather than the +# ``deployment_*`` fields, because a workspace group is a different resource +# from a deployment rather than another way of naming one. +# +_GROUP_KEYS: Tuple[Tuple[str, ...], ...] = ( + ('group',), + ('in', 'in_group'), +) + + +def _first_param( + params: Dict[str, Any], + paths: Tuple[Tuple[str, ...], ...], + field: str, +) -> Optional[str]: + """Return the first value of ``field`` found along ``paths``.""" + for path in paths: + container: Any = params + for key in path: + container = container.get(key) or {} + value = container.get(field) + if value: + return str(value) + return None + + +def _workspace_group( + name: Optional[str] = None, + id: Optional[str] = None, +) -> Optional[Union[WorkspaceGroup, StarterWorkspace]]: + """ + Look a workspace group up by name or ID, or return None if there is none. + + A workspace group is a management API v1 resource, so this goes through the + v1 manager whatever the rest of the statement addresses. Stage is attached + to the group itself at v1 -- the route is ``stage/{group_id}/fs/`` -- so a + group names a Stage on its own, with no workspace to add. A starter + workspace owns its Stage the same way and is the fallback for a name or ID + that is no group's, because both were reachable this way before. + + Returns None rather than raising, so the caller can say whether a miss + means "no such group" or "no such deployment either". + """ + manager = get_workspace_manager() + + if name: + groups = [x for x in manager.workspace_groups if x.name == name] + + if len(groups) == 1: + return groups[0] + + elif len(groups) > 1: + ids = ', '.join(x.id for x in groups) + raise ValueError( + f'more than one workspace group with given name was ' + f'found: {ids}', + ) + + starters = [x for x in manager.starter_workspaces if x.name == name] + + if len(starters) == 1: + return starters[0] + + elif len(starters) > 1: + ids = ', '.join(x.id for x in starters) + raise ValueError( + 'more than one starter workspace with given name was ' + f'found: {ids}', + ) + + return None + + assert id is not None + try: + return manager.get_workspace_group(id) + except ManagementError as exc: + if not _is_missing(exc): + raise + try: + return manager.get_starter_workspace(id) + except ManagementError as exc: + if not _is_missing(exc): + raise + return None + + +def _get_stage_group( + params: Dict[str, Any], +) -> Optional[Union[WorkspaceGroup, StarterWorkspace]]: + """ + Resolve the ``IN GROUP`` spelling, or return None if it was not used. + + Returns None when no group was named, which is the caller's signal to + resolve a deployment instead. A named group that does not exist raises: the + clause says what resource was meant, so there is nothing else to try. + """ + group_name = _first_param(params, _GROUP_KEYS, 'group_name') + group_id = _first_param(params, _GROUP_KEYS, 'group_id') + + if not group_name and not group_id: + return None + + # Warned before the lookup, so a caller who named a group that is gone + # still hears that the spelling itself is going. stacklevel reaches the + # handler method: user code is an unknown number of execute() frames + # further up, so there is no frame count that lands on it. + # + # The warning is about the clause, not the resource: a bare IN resolves a + # workspace group too, so dropping the GROUP keyword is an edit the caller + # can make today whether or not their Stage has moved to a cluster. + warnings.warn( + 'IN GROUP is deprecated: it names a workspace group explicitly, and ' + 'workspace groups are a management API v1 resource that goes away ' + 'with v1. Use a bare IN instead, which names a deployment or a ' + 'workspace group.', + DeprecatedFeatureWarning, stacklevel=3, + ) + + group = _workspace_group(name=group_name, id=group_id) + if group is None: + raise KeyError( + 'no workspace group found with ' + f'{"name" if group_name else "ID"}: {group_name or group_id}', + ) + return group + + +def _group_fallback( + name: Optional[str] = None, + id: Optional[str] = None, +) -> Optional[Union[WorkspaceGroup, StarterWorkspace]]: + """ + Try a name or ID that matched no deployment as a workspace group. + + A bare ``IN`` named a workspace group before the Stage commands moved to + v2, because a group was the only kind of Stage owner there was. So a value + that matches no cluster is tried as a group rather than reported missing, + and ``IN`` names either kind of Stage owner. + + This is deliberately silent. Naming a group with a bare ``IN`` was always + how a Stage was addressed, so there is no statement to correct: the + spelling the user wrote is the one to keep writing, and whether it lands on + a cluster or a group is a fact about their org, not about their SQL. The + group resource does go away with ``management/v1/``, but a warning here + would ask for a migration that no edit to the statement can perform -- + the same reason :func:`.workspace._manage_workspaces_v1` exists. ``IN + GROUP`` still warns, because that spelling *is* something the user can + change. + + The deployment lookup goes first, so a name that is both a cluster's and a + group's is the cluster's, and nothing that resolves today changes meaning. + """ + return _workspace_group(name=name, id=id) + + def get_deployment( params: Dict[str, Any], -) -> Union[WorkspaceGroup, StarterWorkspace]: +) -> Union[Cluster, StarterCluster, WorkspaceGroup, StarterWorkspace]: """ - Find a starter workspace matching deployment_id or deployment_name. + Find the Stage owner named by the statement. - This function will get a starter workspace or ID from the - following parameters: + ``stage.py`` is the only consumer, and it touches nothing but + ``.stage``, which every class returned here provides. + + A bare ``IN`` names a deployment, resolved against management API v2, so it + yields a :class:`Cluster` or a :class:`StarterCluster`. It is the spelling + to use: a deployment is named the same way whatever kind it is, so there is + nothing for a qualified spelling to disambiguate. It is read from: * params['deployment_name'] * params['deployment_id'] - * params['group']['deployment_name'] - * params['group']['deployment_id'] * params['in_deployment']['deployment_name'] * params['in_deployment']['deployment_id'] - * params['in']['in_group']['deployment_name'] - * params['in']['in_group']['deployment_id'] * params['in']['in_deployment']['deployment_name'] * params['in']['in_deployment']['deployment_id'] - Or, from the SINGLESTOREDB_WORKSPACE_GROUP - or SINGLESTOREDB_CLUSTER environment variables. + ``IN GROUP`` names a workspace group instead, resolved against v1 by + :func:`_get_stage_group` from: + + * params['group']['group_name'] + * params['group']['group_id'] + * params['in']['in_group']['group_name'] + * params['in']['in_group']['group_id'] + + ``IN GROUP`` is not a second way of naming a deployment: it names a + different resource at a different version, and it goes away with + ``management/v1/``. It is checked first, so a group is never looked for + among clusters. + + It is not the only way to reach a group, though. A bare ``IN`` that matches + no deployment falls back to :func:`_group_fallback`, so ``IN`` names either + kind of Stage owner and needs no keyword to say which -- a bare ``IN`` + named a workspace group before the Stage commands moved to v2, and that is + still what it does when that is what the name belongs to. The fallback is + second, so a name that is both a cluster's and a group's is the cluster's. + + With neither clause, the deployment comes from ``SINGLESTOREDB_WORKSPACE``, + which is what the notebook environment calls the current deployment + whatever the API version calls it. That path does not fall back -- see + :func:`_deployment_by_id`. """ - manager = get_workspace_manager() + group = _get_stage_group(params) + if group is not None: + return group + + manager = get_cluster_manager() # # Search for deployment by name # - deployment_name = params.get('deployment_name') or \ - (params.get('in_deployment') or {}).get('deployment_name') or \ - (params.get('group') or {}).get('deployment_name') or \ - ((params.get('in') or {}).get('in_group') or {}).get('deployment_name') or \ - ((params.get('in') or {}).get('in_deployment') or {}).get('deployment_name') + deployment_name = _first_param(params, _DEPLOYMENT_KEYS, 'deployment_name') if deployment_name: - # Standard workspace group - workspace_groups = [ - x for x in manager.workspace_groups + # Standard cluster + clusters = [ + x for x in manager.clusters if x.name == deployment_name ] - if len(workspace_groups) == 1: - return workspace_groups[0] + if len(clusters) == 1: + return clusters[0] - elif len(workspace_groups) > 1: - ids = ', '.join(x.id for x in workspace_groups) + elif len(clusters) > 1: + ids = ', '.join(x.id for x in clusters) raise ValueError( - f'more than one workspace group with given name was found: {ids}', + f'more than one cluster with given name was found: {ids}', ) - # Starter workspace - starter_workspaces = [ - x for x in manager.starter_workspaces + # Starter cluster + starter_clusters = [ + x for x in manager.starter_clusters if x.name == deployment_name ] - if len(starter_workspaces) == 1: - return starter_workspaces[0] + if len(starter_clusters) == 1: + return starter_clusters[0] - elif len(starter_workspaces) > 1: - ids = ', '.join(x.id for x in starter_workspaces) + elif len(starter_clusters) > 1: + ids = ', '.join(x.id for x in starter_clusters) raise ValueError( - f'more than one starter workspace with given name was found: {ids}', + 'more than one starter cluster with given name was ' + f'found: {ids}', ) + # No cluster of that name: try it as a workspace group, which is what + # a bare IN named before the Stage commands moved to v2. + group = _group_fallback(name=deployment_name) + if group is not None: + return group + raise KeyError(f'no deployment found with name: {deployment_name}') # # Search for deployment by ID # - deployment_id = params.get('deployment_id') or \ - (params.get('in_deployment') or {}).get('deployment_id') or \ - (params.get('group') or {}).get('deployment_id') or \ - ((params.get('in') or {}).get('in_group') or {}).get('deployment_id') or \ - ((params.get('in') or {}).get('in_deployment') or {}).get('deployment_id') + deployment_id = _first_param(params, _DEPLOYMENT_KEYS, 'deployment_id') if deployment_id: - try: - # Standard workspace group - return manager.get_workspace_group(deployment_id) - except ManagementError as exc: - if exc.errno == 404: - try: - # Starter workspace - return manager.get_starter_workspace(deployment_id) - except ManagementError as exc: - if exc.errno == 404: - raise KeyError(f'no deployment found with ID: {deployment_id}') - raise - else: - raise - - # Use workspace group from environment + return _deployment_by_id(manager, deployment_id, fall_back=True) + + # + # Use the deployment named by the environment. There is one deployment + # resource and the environment names it once, so one lookup tries cluster + # then starter cluster. + # + from_env = get_cluster_id() + if from_env: + return _deployment_by_id( + manager, from_env, 'SINGLESTOREDB_WORKSPACE', + ) + if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP'): - try: - return manager.get_workspace_group( - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'], - ) - except ManagementError as exc: - if exc.errno == 404: - raise KeyError( - 'no workspace found with ID: ' - f'{os.environ["SINGLESTOREDB_WORKSPACE_GROUP"]}', - ) - raise + # Deliberately not resolved. The value is a group ID, which v2 exposes + # only as the read-only Cluster.group attribute -- there is no group + # route to look it up with, so guessing which cluster was meant could + # target the wrong deployment. + # + # Unreachable from a notebook, which never publishes this variable + # without SINGLESTOREDB_WORKSPACE, resolved above. It is here for a + # value set by hand. + raise KeyError( + 'SINGLESTOREDB_WORKSPACE_GROUP holds a group ID, which management ' + 'API v2 reports as a cluster attribute rather than something that ' + 'can be looked up -- clusters are flat. Set ' + 'SINGLESTOREDB_WORKSPACE to the cluster ID instead, or name the ' + 'deployment with IN.', + ) - # Use cluster from environment - if os.environ.get('SINGLESTOREDB_CLUSTER'): - try: - return manager.get_starter_workspace( - os.environ['SINGLESTOREDB_CLUSTER'], - ) - except ManagementError as exc: - if exc.errno == 404: - raise KeyError( - 'no starter workspace found with ID: ' - f'{os.environ["SINGLESTOREDB_CLUSTER"]}', - ) + raise KeyError('no deployment was specified') + + +def _deployment_by_id( + manager: ClusterManager, + deployment_id: str, + envvar: Optional[str] = None, + fall_back: bool = False, +) -> Union[Cluster, StarterCluster, WorkspaceGroup, StarterWorkspace]: + """ + Look an ID up as a cluster, then as a starter cluster. + + ``fall_back`` then tries it as a workspace group, for an ID the statement + named itself. An ID from the environment does not fall back: at v1 that + variable held a *workspace* ID, which is no group's, so the lookup could + only ever add a wasted round trip to a failure. + """ + source = f' (from {envvar})' if envvar else '' + try: + return manager.get_cluster(deployment_id) + except ManagementError as exc: + if not _is_missing(exc): + raise + try: + return manager.get_starter_cluster(deployment_id) + except ManagementError as exc: + if not _is_missing(exc): raise - raise KeyError('no deployment was specified') + if fall_back: + group = _group_fallback(id=deployment_id) + if group is not None: + return group + + raise KeyError(f'no deployment found with ID: {deployment_id}{source}') def get_file_space(params: Dict[str, Any]) -> FileSpace: @@ -327,7 +765,17 @@ def get_file_space(params: Dict[str, Any]) -> FileSpace: def get_inference_api_manager() -> InferenceAPIManager: - """Return the inference API manager for the current project.""" + """ + Return the inference API manager for the current project. + + Stays on the v1 manager while files and jobs move to v2, because unlike + those two there is no v2 route to move to: ``Organization.inference_apis`` + raises for every version past v1, and the implementation is imported from + ``management/v1/inference_api.py`` by that name -- there is deliberately no + version-neutral alias for it. The handlers this feeds are hidden for the + same reason (see ``handlers/models.py``). Revisit when the models and + inference surface gains a v2 equivalent. + """ wm = get_workspace_manager() return wm.organization.inference_apis diff --git a/singlestoredb/fusion/handlers/workspace.py b/singlestoredb/fusion/handlers/workspace.py index 24870fd6d..9bb44ce44 100644 --- a/singlestoredb/fusion/handlers/workspace.py +++ b/singlestoredb/fusion/handlers/workspace.py @@ -1,4 +1,18 @@ #!/usr/bin/env python3 +""" +Fusion SQL handlers for the management API v1 workspace vocabulary. + +**Deprecated.** ``handlers/cluster.py`` is the v2 replacement, and v2 is the +default everywhere else in the SDK. Every command here sets ``_deprecated_by`` +naming its ``CLUSTER`` counterpart, so it still runs but warns once per +execution. Nothing is removed and no grammar changed -- an existing v1 script +keeps working, it just says where to go. This module is what gets deleted when +``management/v1/`` goes. + +Pinned to v1 through :func:`.utils.get_workspace_manager`: these commands *are* +the v1 vocabulary, so they must not follow the ``management.version`` option onto +a version that has no workspaces. +""" import json from typing import Any from typing import Dict @@ -77,6 +91,8 @@ class UseWorkspaceHandler(SQLHandler): USE WORKSPACE 'examplews' IN GROUP 'my-workspace-group'; """ + _deprecated_by = 'USE CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: from singlestoredb.notebook import portal @@ -168,6 +184,9 @@ class ShowRegionsHandler(SQLHandler): specified number. * Use the ``ORDER BY`` clause to sort the results by the specified key. By default, the results are sorted in the ascending order. + * The ``ID`` column has no counterpart in ``SHOW CLUSTER REGIONS``: v2 + assigns no region IDs and identifies a region by its provider and + region name instead. Example ------- @@ -176,8 +195,20 @@ class ShowRegionsHandler(SQLHandler): SHOW REGIONS LIKE 'US%' ORDER BY Name; + See Also + -------- + * ``SHOW CLUSTER REGIONS``, the management API v2 replacement + """ + # Not a column-for-column replacement, unlike the rest of this module: v2 + # has no region IDs, so ``SHOW CLUSTER REGIONS`` reports ``Provider`` and + # ``RegionName`` where this reports ``ID``. Deprecated anyway, because this + # command reads the v1 API and that is what is going away -- a caller + # holding a v1 region ID needs to hear that now, not when the route stops + # answering. + _deprecated_by = 'SHOW CLUSTER REGIONS' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: manager = get_workspace_manager() @@ -236,6 +267,8 @@ class ShowWorkspaceGroupsHandler(SQLHandler): """ + _deprecated_by = 'SHOW CLUSTERS' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: manager = get_workspace_manager() @@ -327,6 +360,8 @@ class ShowWorkspacesHandler(SQLHandler): """ + _deprecated_by = 'SHOW CLUSTERS' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res = FusionSQLResult() res.add_field('Name', result.STRING) @@ -471,6 +506,8 @@ class CreateWorkspaceGroupHandler(SQLHandler): """ + _deprecated_by = 'CREATE CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: manager = get_workspace_manager() @@ -606,6 +643,8 @@ class CreateWorkspaceHandler(SQLHandler): """ # noqa: E501 + _deprecated_by = 'CREATE CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: workspace_group = get_workspace_group(params) @@ -708,6 +747,8 @@ class SuspendWorkspaceHandler(SQLHandler): """ # noqa: E501 + _deprecated_by = 'SUSPEND CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: ws = get_workspace(params) ws.suspend(wait_on_suspended=params['wait_on_suspended']) @@ -783,6 +824,8 @@ class ResumeWorkspaceHandler(SQLHandler): """ # noqa: E501 + _deprecated_by = 'RESUME CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: ws = get_workspace(params) ws.resume( @@ -851,6 +894,8 @@ class DropWorkspaceGroupHandler(SQLHandler): """ + _deprecated_by = 'DROP CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: try: workspace_group = get_workspace_group(params) @@ -939,6 +984,8 @@ class DropWorkspaceHandler(SQLHandler): """ + _deprecated_by = 'DROP CLUSTER' + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: try: ws = get_workspace(params) diff --git a/singlestoredb/management/__init__.py b/singlestoredb/management/__init__.py index 8a87d2840..cc99bb495 100644 --- a/singlestoredb/management/__init__.py +++ b/singlestoredb/management/__init__.py @@ -1,9 +1,15 @@ #!/usr/bin/env python -from .cluster import manage_cluster +# Everything exported here is version-neutral: an explicit ``version=`` wins, +# otherwise the ``management.version`` option (the +# SINGLESTOREDB_MANAGEMENT_VERSION environment variable) decides which version +# package answers the call. Import from .v1/.v2 -- or from the version-locked +# shims .workspace and .cluster -- to pin a version instead. +from .cluster import manage_clusters from .files import manage_files from .manager import get_token +from .organization import get_organization +from .organization import get_secret from .region import manage_regions -from .workspace import get_organization -from .workspace import get_secret -from .workspace import get_stage +from .stage import get_stage +from .timing import trace as trace_timing from .workspace import manage_workspaces diff --git a/singlestoredb/management/_version_import.py b/singlestoredb/management/_version_import.py new file mode 100644 index 000000000..254560e05 --- /dev/null +++ b/singlestoredb/management/_version_import.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python +"""Importer for version-specific management API modules.""" +import importlib +import re +import warnings +from typing import Any +from typing import Optional + +from .. import config +from .._management_version import DEFAULT_MANAGEMENT_VERSION +from .._management_version import DEPRECATED_MANAGEMENT_VERSION +from ..exceptions import ManagementError + + +_VERSION_RE = re.compile(r'^v\d+$') + +#: API version used when neither the caller nor the ``management.version`` +#: option names one -- i.e. when the option has been explicitly blanked out, +#: since it otherwise carries this same default itself. Everything that should +#: follow the current version reads this, so it is fixed in one place: +#: :data:`singlestoredb._management_version.DEFAULT_MANAGEMENT_VERSION`. +DEFAULT_VERSION = DEFAULT_MANAGEMENT_VERSION + +#: The version this SDK is winding down. Everything under +#: ``singlestoredb.management.v1`` goes away with it, so any *public* entry +#: point that resolves to it warns -- see :func:`_warn_if_deprecated_version`. +DEPRECATED_VERSION = DEPRECATED_MANAGEMENT_VERSION + + +def _warn_if_deprecated_version(version: str, stacklevel: int = 3) -> None: + """ + Warn if ``version`` names a management API version being wound down. + + Called from the public version-neutral entry points -- the ``manage_*`` + factories and the ``get_organization``/``get_secret``/``get_stage`` + helpers -- *after* the version has been resolved, so it fires whether v1 + was named by the caller or inherited from the ``management.version`` + option. + + Deliberately not called from :func:`_resolve_version` itself. Several + internal paths are v1-only by design and resolve v1 with no v2 route to + move to -- ``workspace._manage_workspaces_v1`` and the inference API + behind it -- so warning at the resolver would emit noise the caller can do + nothing about. :func:`manage_workspaces` is likewise excluded: it raises + its own, more specific warning naming ``manage_clusters``. + + Parameters + ---------- + version : str + The already-resolved version + stacklevel : int, optional + Passed through to :func:`warnings.warn`. The default of 3 is right for + a public entry point calling this directly: 1 is this function, 2 is + the entry point, 3 is the user. Add one per intervening frame. + + """ + if version != DEPRECATED_VERSION: + return + warnings.warn( + f'management API {DEPRECATED_VERSION} is deprecated and will be ' + 'removed; it has been replaced by ' + f'{DEFAULT_VERSION}. Stop passing version=' + f'"{DEPRECATED_VERSION}", and unset the management.version option ' + '(the SINGLESTOREDB_MANAGEMENT_VERSION environment variable) if it ' + f'names {DEPRECATED_VERSION}.', + DeprecationWarning, + stacklevel=stacklevel + 1, + ) + + +def _resolve_version( + version: Optional[str] = None, + default: Optional[str] = None, +) -> str: + """ + Resolve the management API version to use. + + An explicit argument wins; otherwise the ``management.version`` option (the + ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment variable) decides. This is + the one place that rule is written down; every version-neutral entry point + goes through here so they cannot drift apart. + + Parameters + ---------- + version : str, optional + Version named by the caller, if any + default : str, optional + Version to use when neither the caller nor the option names one. + Defaults to :data:`DEFAULT_VERSION`. + + Returns + ------- + str + + """ + return version or config.get_option('management.version') \ + or default or DEFAULT_VERSION + + +def _import_versioned_package(version: str) -> Any: + """Import a version package, raising a friendly error if not found.""" + if not _VERSION_RE.match(version): + raise ManagementError( + msg=f"Invalid API version format: '{version}'", + ) + try: + return importlib.import_module(f'singlestoredb.management.{version}') + except ModuleNotFoundError: + raise ManagementError( + msg=f"Unsupported API version: '{version}'", + ) + + +def _versioned_attr(name: str, version: Optional[str] = None) -> Any: + """ + Look a name up in the resolved version package. + + Version-neutral helpers dispatch through here rather than naming the module + that holds their implementation, because that module differs by version -- + the v1 helpers hang off workspaces, the v2 helpers off clusters -- and a + future version is free to put them somewhere else again. Each version + package re-exports its own, so this layer only has to resolve the version. + + Every caller is a public entry point one frame up + (``organization.get_organization``, ``organization.get_secret``, + ``stage.get_stage``), so the deprecated-version warning is raised here + rather than repeated in each of them. + + Parameters + ---------- + name : str + Name to look up in the version package + version : str, optional + Version of the API to use. Defaults to the ``management.version`` + option. + + Returns + ------- + Any + + Raises + ------ + :class:`ManagementError` + If the resolved version does not provide the name + + """ + ver = _resolve_version(version) + # +1 for this frame sitting between the helper and the user. + _warn_if_deprecated_version(ver, stacklevel=4) + pkg = _import_versioned_package(ver) + try: + return getattr(pkg, name) + except AttributeError: + raise ManagementError( + msg=f"management API {ver} does not provide '{name}'", + ) + + +def _import_versioned_module(version: str, module_name: str) -> Any: + """Import a versioned module, raising a friendly error if not found.""" + if not _VERSION_RE.match(version): + raise ManagementError( + msg=f"Invalid API version format: '{version}'", + ) + version_pkg = f'singlestoredb.management.{version}' + path = f'{version_pkg}.{module_name}' + try: + return importlib.import_module(path) + except ModuleNotFoundError as e: + if e.name is None or (e.name != path and not path.startswith(e.name)): + # Failure originated deeper than the requested module + # (e.g., a transitive import inside a valid module). Don't mask. + raise + try: + importlib.import_module(version_pkg) + except ModuleNotFoundError: + raise ManagementError( + msg=f"Unsupported API version: '{version}'", + ) + raise ManagementError( + msg=f"API version '{version}' does not provide " + f"module '{module_name}'", + ) diff --git a/singlestoredb/management/billing.py b/singlestoredb/management/billing.py new file mode 100644 index 000000000..e73eecd50 --- /dev/null +++ b/singlestoredb/management/billing.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python +""" +SingleStoreDB billing information. + +``GET /v1/billing/usage`` and ``GET /v2/billing/usage`` are the same route with +the same response, so this is version-neutral. +""" +import datetime +from typing import List +from typing import Optional + +from .billing_usage import BillingUsageItem +from .manager import Manager +from .utils import from_datetime +from .utils import snake_to_camel + + +class Billing(object): + """Billing information.""" + + COMPUTE_CREDIT = 'compute_credit' + STORAGE_AVG_BYTE = 'storage_avg_byte' + + HOUR = 'hour' + DAY = 'day' + MONTH = 'month' + + def __init__(self, manager: Manager): + self._manager = manager + + def usage( + self, + start_time: datetime.datetime, + end_time: datetime.datetime, + metric: Optional[str] = None, + aggregate_by: Optional[str] = None, + ) -> List[BillingUsageItem]: + """ + Get usage information. + + Parameters + ---------- + start_time : datetime.datetime + Start time for usage interval + end_time : datetime.datetime + End time for usage interval + metric : str, optional + Possible metrics are ``mgr.billing.COMPUTE_CREDIT`` and + ``mgr.billing.STORAGE_AVG_BYTE`` (default is all) + aggregate_by : str, optional + Aggregate type used to group usage: ``mgr.billing.HOUR``, + ``mgr.billing.DAY``, or ``mgr.billing.MONTH`` + + Returns + ------- + List[BillingUsage] + + """ + res = self._manager._get( + 'billing/usage', + params={ + k: v for k, v in dict( + metric=snake_to_camel(metric), + startTime=from_datetime(start_time), + endTime=from_datetime(end_time), + aggregateBy=aggregate_by.lower() if aggregate_by else None, + ).items() if v is not None + }, + ) + return [ + BillingUsageItem.from_dict(x, self._manager) + for x in res.json()['billingUsage'] + ] diff --git a/singlestoredb/management/billing_usage.py b/singlestoredb/management/billing_usage.py index 24c8683dc..43c80b376 100644 --- a/singlestoredb/management/billing_usage.py +++ b/singlestoredb/management/billing_usage.py @@ -8,10 +8,11 @@ from .manager import Manager from .utils import camel_to_snake +from .utils import to_datetime_strict from .utils import vars_to_str -class UsageItem(object): +class UsageItem: """Usage statistics.""" def __init__( @@ -67,9 +68,9 @@ def from_dict( Parameters ---------- obj : dict - Key-value pairs to retrieve billling usage information from - manager : WorkspaceManager, optional - The WorkspaceManager the UsageItem belongs to + Key-value pairs to retrieve billing usage information from + manager : ClusterManager, optional + The ClusterManager the UsageItem belongs to Returns ------- @@ -77,19 +78,19 @@ def from_dict( """ out = cls( - end_time=datetime.datetime.fromisoformat(obj['endTime']), - start_time=datetime.datetime.fromisoformat(obj['startTime']), + end_time=to_datetime_strict(obj['endTime']), + start_time=to_datetime_strict(obj['startTime']), 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 return out -class BillingUsageItem(object): +class BillingUsageItem: """Billing usage item.""" def __init__( @@ -98,7 +99,7 @@ def __init__( metric: str, usage: List[UsageItem], ): - """Use :attr:`WorkspaceManager.billing.usage` instead.""" + """Use :attr:`ClusterManager.billing.usage` instead.""" #: Description of the usage metric self.description = description @@ -118,7 +119,7 @@ def __repr__(self) -> str: """Return string representation.""" return str(self) - @ classmethod + @classmethod def from_dict( cls, obj: Dict[str, Any], @@ -130,9 +131,9 @@ def from_dict( Parameters ---------- obj : dict - Key-value pairs to retrieve billling usage information from - manager : WorkspaceManager, optional - The WorkspaceManager the BillingUsageItem belongs to + Key-value pairs to retrieve billing usage information from + manager : ClusterManager, optional + The ClusterManager the BillingUsageItem belongs to Returns ------- @@ -142,7 +143,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/cluster.py b/singlestoredb/management/cluster.py index 8fa6ae2c1..8d9f28d76 100644 --- a/singlestoredb/management/cluster.py +++ b/singlestoredb/management/cluster.py @@ -1,431 +1,37 @@ #!/usr/bin/env python -"""SingleStoreDB Cluster Management.""" -import datetime -import warnings -from typing import Any -from typing import Dict -from typing import List -from typing import Optional -from typing import Union - -from .. import config -from .. import connection -from ..exceptions import ManagementError -from .manager import Manager -from .region import Region -from .utils import NamedList -from .utils import to_datetime -from .utils import vars_to_str - - -class Cluster(object): - """ - SingleStoreDB cluster definition. - - This object is not instantiated directly. It is used in the results - of API calls on the :class:`ClusterManager`. Clusters are created using - :meth:`ClusterManager.create_cluster`, or existing clusters are accessed by either - :attr:`ClusterManager.clusters` or by calling :meth:`ClusterManager.get_cluster`. - - See Also - -------- - :meth:`ClusterManager.create_cluster` - :meth:`ClusterManager.get_cluster` - :attr:`ClusterManager.clusters` - - """ - - def __init__( - self, name: str, id: str, region: Region, size: str, - units: float, state: str, version: str, - created_at: Union[str, datetime.datetime], - expires_at: Optional[Union[str, datetime.datetime]] = None, - firewall_ranges: Optional[List[str]] = None, - terminated_at: Optional[Union[str, datetime.datetime]] = None, - endpoint: Optional[str] = None, - ): - """Use :attr:`ClusterManager.clusters` or :meth:`ClusterManager.get_cluster`.""" - #: Name of the cluster - self.name = name.strip() - - #: Unique ID of the cluster - self.id = id - - #: Region of the cluster (see :class:`Region`) - self.region = region - - #: Size of the cluster in cluster size notation (S-00, S-1, etc.) - self.size = size - - #: Size of the cluster in units such as 0.25, 1.0, etc. - self.units = units - - #: State of the cluster: PendingCreation, Transitioning, Active, - #: Terminated, Suspended, Resuming, Failed - self.state = state.strip() - - #: Version of the SingleStoreDB server - self.version = version.strip() - - #: Timestamp of when the cluster was created - self.created_at = to_datetime(created_at) - - #: Timestamp of when the cluster expires - self.expires_at = to_datetime(expires_at) - - #: List of allowed incoming IP addresses / ranges - self.firewall_ranges = firewall_ranges - - #: Timestamp of when the cluster was terminated - self.terminated_at = to_datetime(terminated_at) - - #: Hostname (or IP address) of the cluster database server - self.endpoint = endpoint - - self._manager: Optional[ClusterManager] = 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: 'ClusterManager') -> 'Cluster': - """ - Construct a Cluster from a dictionary of values. - - Parameters - ---------- - obj : dict - Dictionary of values - manager : ClusterManager, optional - The ClusterManager the Cluster belongs to - - Returns - ------- - :class:`Cluster` - - """ - out = cls( - name=obj['name'], id=obj['clusterID'], - region=Region.from_dict(obj['region'], manager), - size=obj.get('size', 'Unknown'), units=obj.get('units', float('nan')), - state=obj['state'], version=obj['version'], - created_at=obj['createdAt'], expires_at=obj.get('expiresAt'), - firewall_ranges=obj.get('firewallRanges'), - terminated_at=obj.get('terminatedAt'), - endpoint=obj.get('endpoint'), - ) - out._manager = manager - return out - - def refresh(self) -> 'Cluster': - """Update the object to the current state.""" - if self._manager is None: - raise ManagementError( - msg='No cluster manager is associated with this object.', - ) - new_obj = self._manager.get_cluster(self.id) - for name, value in vars(new_obj).items(): - setattr(self, name, value) - return self - - def update( - self, name: Optional[str] = None, - admin_password: Optional[str] = None, - expires_at: Optional[str] = None, - size: Optional[str] = None, firewall_ranges: Optional[List[str]] = None, - ) -> None: - """ - Update the cluster definition. - - Parameters - ---------- - name : str, optional - Cluster name - admim_password : str, optional - Admin password for the cluster - expires_at : str, optional - Timestamp when the cluster expires - size : str, optional - Cluster size in cluster size notation (S-00, S-1, etc.) - firewall_ranges : Sequence[str], optional - List of allowed incoming IP addresses - - """ - if self._manager is None: - raise ManagementError( - msg='No cluster manager is associated with this object.', - ) - data = { - k: v for k, v in dict( - name=name, adminPassword=admin_password, - expiresAt=expires_at, size=size, - firewallRanges=firewall_ranges, - ).items() if v is not None - } - self._manager._patch(f'clusters/{self.id}', json=data) - self.refresh() - - def suspend( - self, - wait_on_suspended: bool = False, - wait_interval: int = 20, - wait_timeout: int = 600, - ) -> None: - """ - Suspend the cluster. - - Parameters - ---------- - wait_on_suspended : bool, optional - Wait for the cluster to go into 'Suspended' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No cluster manager is associated with this object.', - ) - self._manager._post( - f'clusters/{self.id}/suspend', - headers={'Content-Type': 'application/x-www-form-urlencoded'}, - ) - if wait_on_suspended: - self._manager._wait_on_state( - self._manager.get_cluster(self.id), - 'Suspended', interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() - - def resume( - self, - wait_on_resumed: bool = False, - wait_interval: int = 20, - wait_timeout: int = 600, - ) -> None: - """ - Resume the cluster. - - Parameters - ---------- - wait_on_resumed : bool, optional - Wait for the cluster to go into 'Resumed' or 'Active' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No cluster manager is associated with this object.', - ) - self._manager._post( - f'clusters/{self.id}/resume', - headers={'Content-Type': 'application/x-www-form-urlencoded'}, - ) - if wait_on_resumed: - self._manager._wait_on_state( - self._manager.get_cluster(self.id), - ['Resumed', 'Active'], interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() - - def terminate( - self, - wait_on_terminated: bool = False, - wait_interval: int = 10, - wait_timeout: int = 600, - ) -> None: - """ - Terminate the cluster. - - Parameters - ---------- - wait_on_terminated : bool, optional - Wait for the cluster to go into 'Terminated' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No cluster manager is associated with this object.', - ) - self._manager._delete(f'clusters/{self.id}') - if wait_on_terminated: - self._manager._wait_on_state( - self._manager.get_cluster(self.id), - 'Terminated', interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() - - def connect(self, **kwargs: Any) -> connection.Connection: - """ - Create a connection to the database server for this cluster. - - Parameters - ---------- - **kwargs : keyword-arguments, optional - Parameters to the SingleStoreDB `connect` function except host - and port which are supplied by the cluster object - - Returns - ------- - :class:`Connection` - - """ - if not self.endpoint: - raise ManagementError( - msg='An endpoint has not been set in ' - 'this cluster configuration', - ) - kwargs['host'] = self.endpoint - return connection.connect(**kwargs) - - -class ClusterManager(Manager): - """ - SingleStoreDB cluster manager. - - This class should be instantiated using :func:`singlestoredb.manage_cluster`. - - Parameters - ---------- - access_token : str, optional - The API key or other access token for the cluster management API - version : str, optional - Version of the API to use - base_url : str, optional - Base URL of the cluster management API - - See Also - -------- - :func:`singlestoredb.manage_cluster` - - """ - - #: Cluster management API version if none is specified. - default_version = 'v0beta' - - #: Base URL if none is specified. - default_base_url = config.get_option('management.base_url') \ - or 'https://api.singlestore.com' - - #: Object type - obj_type = 'cluster' - - @property - def clusters(self) -> NamedList[Cluster]: - """Return a list of available clusters.""" - res = self._get('clusters') - return NamedList([Cluster.from_dict(item, self) for item in res.json()]) - - @property - def regions(self) -> NamedList[Region]: - """Return a list of available regions.""" - res = self._get('regions') - return NamedList([Region.from_dict(item, self) for item in res.json()]) - - def create_cluster( - self, name: str, region: Union[str, Region], admin_password: str, - firewall_ranges: List[str], expires_at: Optional[str] = None, - size: Optional[str] = None, plan: Optional[str] = None, - wait_on_active: bool = False, wait_timeout: int = 600, - wait_interval: int = 20, - ) -> Cluster: - """ - Create a new cluster. - - Parameters - ---------- - name : str - Name of the cluster - region : str or Region - The region ID of the cluster - admin_password : str - Admin password for the cluster - firewall_ranges : Sequence[str], optional - List of allowed incoming IP addresses - expires_at : str, optional - Timestamp of when the cluster expires - size : str, optional - Cluster size in cluster size notation (S-00, S-1, etc.) - plan : str, optional - Internal use only - wait_on_active : bool, optional - Wait for the cluster to be active before returning - wait_timeout : int, optional - Maximum number of seconds to wait before raising an exception - if wait=True - wait_interval : int, optional - Number of seconds between each polling interval - - Returns - ------- - :class:`Cluster` - - """ - if isinstance(region, Region) and region.id: - region = region.id - res = self._post( - 'clusters', json=dict( - name=name, regionID=region, adminPassword=admin_password, - expiresAt=expires_at, size=size, firewallRanges=firewall_ranges, - plan=plan, - ), - ) - out = self.get_cluster(res.json()['clusterID']) - if wait_on_active: - out = self._wait_on_state( - out, 'Active', interval=wait_interval, - timeout=wait_timeout, - ) - return out - - def get_cluster(self, id: str) -> Cluster: - """ - Retrieve a cluster definition. - - Parameters - ---------- - id : str - ID of the cluster - - Returns - ------- - :class:`Cluster` - - """ - res = self._get(f'clusters/{id}') - return Cluster.from_dict(res.json(), manager=self) +""" +SingleStoreDB Cluster Management. +Clusters are the flat deployment resource introduced by management API v2, so +the names below come from :mod:`singlestoredb.management.v2.cluster`. There is +no v1 cluster resource; :func:`manage_clusters` defaults to v2 accordingly. +""" +from typing import Optional -def manage_cluster( +from ._version_import import _import_versioned_module +from ._version_import import DEFAULT_VERSION +from .v2.cluster import Cluster as Cluster +from .v2.cluster import ClusterManager as ClusterManager +from .v2.cluster import get_cluster as get_cluster +from .v2.cluster import get_organization as get_organization +from .v2.cluster import get_secret as get_secret +from .v2.cluster import get_stage as get_stage +from .v2.cluster import Project as Project +from .v2.cluster import PROJECT_ID_RE as PROJECT_ID_RE +from .v2.cluster import SHAREDTIER_PATH as SHAREDTIER_PATH +from .v2.cluster import Stage as Stage +from .v2.cluster import StageObject as StageObject +from .v2.cluster import StarterCluster as StarterCluster + +#: API version used by :func:`manage_clusters` when neither the caller nor the +#: ``management.version`` option names one. Clusters exist at every version +#: from v2 on, so this follows +#: :data:`~singlestoredb.management._version_import.DEFAULT_VERSION` rather +#: than naming a version of its own; v1 is rejected below instead. +DEFAULT_CLUSTER_VERSION = DEFAULT_VERSION + + +def manage_clusters( access_token: Optional[str] = None, version: Optional[str] = None, base_url: Optional[str] = None, @@ -440,23 +46,43 @@ def manage_cluster( access_token : str, optional The API key or other access token for the cluster management API version : str, optional - Version of the API to use + Version of the API to use. Defaults to the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment + variable), or to :data:`DEFAULT_CLUSTER_VERSION` when that is unset. base_url : str, optional Base URL of the cluster management API - organization_id: str, optional + organization_id : str, optional ID of organization, if using a JWT for authentication Returns ------- :class:`ClusterManager` + Raises + ------ + :class:`ManagementError` + If ``v1`` is the resolved version, whether requested by the caller or + by the ``management.version`` option. Clusters were introduced in v2; + the v1 equivalents are workspaces, reached with + :func:`singlestoredb.manage_workspaces`. + """ - warnings.warn( - 'The cluster management API is deprecated; ' - 'use manage_workspaces instead.', - category=DeprecationWarning, - ) - return ClusterManager( + from ..exceptions import ManagementError + from ._version_import import _resolve_version + # Follows the management.version option like the other public entry points + # rather than pinning the front door to one version, so a future version is + # picked up from the environment. A bare call lands on the current default, + # which has clusters; an explicit 'v1' does not and raises below. + ver = _resolve_version(version, default=DEFAULT_CLUSTER_VERSION) + if ver == 'v1': + raise ManagementError( + msg='clusters do not exist in management API v1; they replaced ' + 'workspaces in v2. Use manage_workspaces() instead, or ask ' + 'for v2, either with version="v2" here or by setting the ' + 'management.version option.', + ) + mod = _import_versioned_module(ver, 'cluster') + return mod.ClusterManager( access_token=access_token, base_url=base_url, - version=version, organization_id=organization_id, + version=ver, organization_id=organization_id, ) diff --git a/singlestoredb/management/export.py b/singlestoredb/management/export.py index a84efbb9c..8ed98879b 100644 --- a/singlestoredb/management/export.py +++ b/singlestoredb/management/export.py @@ -1,295 +1,17 @@ #!/usr/bin/env python -"""SingleStoreDB export service.""" -from __future__ import annotations - -import copy -import json -from typing import Any -from typing import Dict -from typing import List -from typing import Optional -from typing import Union - -from .. import ManagementError -from .utils import vars_to_str -from .workspace import WorkspaceGroup -from .workspace import WorkspaceManager - - -class ExportService(object): - """Export service.""" - - database: str - table: str - catalog_info: Dict[str, Any] - storage_info: Dict[str, Any] - columns: Optional[List[str]] - partition_by: Optional[List[Dict[str, str]]] - order_by: Optional[List[Dict[str, Dict[str, str]]]] - properties: Optional[Dict[str, Any]] - incremental: bool - refresh_interval: Optional[int] - export_id: Optional[str] - - def __init__( - self, - workspace_group: WorkspaceGroup, - database: str, - table: str, - catalog_info: Union[str, Dict[str, Any]], - storage_info: Union[str, Dict[str, Any]], - columns: Optional[List[str]] = None, - partition_by: Optional[List[Dict[str, str]]] = None, - order_by: Optional[List[Dict[str, Dict[str, str]]]] = None, - incremental: bool = False, - refresh_interval: Optional[int] = None, - properties: Optional[Dict[str, Any]] = None, - ): - #: Workspace group - self.workspace_group = workspace_group - - #: Name of SingleStoreDB database - self.database = database - - #: Name of SingleStoreDB table - self.table = table - - #: List of columns to export - self.columns = columns - - #: Catalog - if isinstance(catalog_info, str): - self.catalog_info = json.loads(catalog_info) - else: - self.catalog_info = copy.copy(catalog_info) - - #: Storage - if isinstance(storage_info, str): - self.storage_info = json.loads(storage_info) - else: - self.storage_info = copy.copy(storage_info) - - self.partition_by = partition_by or None - self.order_by = order_by or None - self.properties = properties or None - - self.incremental = incremental - self.refresh_interval = refresh_interval - - self.export_id = None - - self._manager: Optional[WorkspaceManager] = workspace_group._manager - - @classmethod - def from_export_id( - self, - workspace_group: WorkspaceGroup, - export_id: str, - ) -> ExportService: - """Create export service from export ID.""" - out = ExportService( - workspace_group=workspace_group, - database='', - table='', - catalog_info={}, - storage_info={}, - ) - out.export_id = export_id - return out - - def __str__(self) -> str: - """Return string representation.""" - return vars_to_str(self) - - def __repr__(self) -> str: - """Return string representation.""" - return str(self) - - def create_cluster_identity(self) -> Dict[str, Any]: - """Create a cluster identity.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - out = self._manager._post( - f'workspaceGroups/{self.workspace_group.id}/' - 'egress/createEgressClusterIdentity', - json=dict( - catalogInfo=self.catalog_info, - storageInfo=self.storage_info, - ), - ) - - return out.json() - - def start(self, tags: Optional[List[str]] = None) -> 'ExportStatus': - """Start the export process.""" - if not self.table or not self.database: - raise ManagementError( - msg='Database and table must be set before starting the export.', - ) - - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - partition_spec = None - if self.partition_by: - partition_spec = dict(partitions=self.partition_by) - - sort_order_spec = None - if self.order_by: - sort_order_spec = dict(keys=self.order_by) - - out = self._manager._post( - f'workspaceGroups/{self.workspace_group.id}/egress/startTableEgress', - json={ - k: v for k, v in dict( - databaseName=self.database, - tableName=self.table, - storageInfo=self.storage_info, - catalogInfo=self.catalog_info, - partitionSpec=partition_spec, - sortOrderSpec=sort_order_spec, - properties=self.properties, - incremental=self.incremental or None, - refreshInterval=self.refresh_interval - if self.refresh_interval is not None else None, - ).items() if v is not None - }, - ) - - self.export_id = str(out.json()['egressID']) - - return ExportStatus(self.export_id, self.workspace_group) - - def suspend(self) -> 'ExportStatus': - """Suspend the export process.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - if self.export_id is None: - raise ManagementError( - msg='Export ID is not set. You must start the export first.', - ) - - self._manager._post( - f'workspaceGroups/{self.workspace_group.id}/egress/suspendTableEgress', - json=dict(egressID=self.export_id), - ) - - return ExportStatus(self.export_id, self.workspace_group) - - def resume(self) -> 'ExportStatus': - """Resume the export process.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - if self.export_id is None: - raise ManagementError( - msg='Export ID is not set. You must start the export first.', - ) - - self._manager._post( - f'workspaceGroups/{self.workspace_group.id}/egress/resumeTableEgress', - json=dict(egressID=self.export_id), - ) - - return ExportStatus(self.export_id, self.workspace_group) - - def drop(self) -> None: - """Drop the export process.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - if self.export_id is None: - raise ManagementError( - msg='Export ID is not set. You must start the export first.', - ) - - self._manager._delete( - f'workspaceGroups/{self.workspace_group.id}/egress/dropTableEgress', - json=dict(egressID=self.export_id), - ) - - return None - - def status(self) -> ExportStatus: - """Get the status of the export process.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - if self.export_id is None: - raise ManagementError( - msg='Export ID is not set. You must start the export first.', - ) - - return ExportStatus(self.export_id, self.workspace_group) - - -class ExportStatus(object): - - export_id: str - - def __init__(self, export_id: str, workspace_group: WorkspaceGroup): - self.export_id = export_id - self.workspace_group = workspace_group - self._manager: Optional[WorkspaceManager] = workspace_group._manager - - def _info(self) -> Dict[str, Any]: - """Return export status.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - out = self._manager._get( - f'workspaceGroups/{self.workspace_group.id}/egress/tableEgressStatus', - json=dict(egressID=self.export_id), - ) - - return out.json() - - @property - def status(self) -> str: - """Return export status.""" - return self._info().get('status', 'Unknown') - - @property - def message(self) -> str: - """Return export status message.""" - return self._info().get('statusMsg', '') - - def __str__(self) -> str: - return self.status - - def __repr__(self) -> str: - return self.status - - -def _get_exports( - workspace_group: WorkspaceGroup, - scope: str = 'all', -) -> List[ExportStatus]: - """Get all exports in the workspace group.""" - if workspace_group._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - out = workspace_group._manager._get( - f'workspaceGroups/{workspace_group.id}/egress/tableEgressStatus', - json=dict(scope=scope), - ) - - return out.json() +""" +SingleStoreDB export service. + +The names below come from :mod:`singlestoredb.management.v2.export`, matching +:mod:`singlestoredb.management.cluster`: table egress is driven through +``clusters/{id}/egress/...``, so an export is owned by a +:class:`~singlestoredb.management.v2.cluster.Cluster`. + +This is a version-locked shim, not a version-neutral one -- it does not consult +the ``management.version`` option, because the two implementations take +different objects (a ``Cluster`` at v2, a ``WorkspaceGroup`` at v1) and so +cannot be swapped behind one name. Import from :mod:`.v1.export` to pin v1. +""" +from .v2.export import _get_exports as _get_exports +from .v2.export import ExportService as ExportService +from .v2.export import ExportStatus as ExportStatus diff --git a/singlestoredb/management/files.py b/singlestoredb/management/files.py index 593f7e398..5a1b26f46 100644 --- a/singlestoredb/management/files.py +++ b/singlestoredb/management/files.py @@ -3,7 +3,6 @@ from __future__ import annotations import datetime -import glob import io import os import re @@ -21,7 +20,10 @@ from .. import config from ..exceptions import ManagementError from .manager import Manager +from .utils import ensure_within +from .utils import normalize_remote_path from .utils import PathLike +from .utils import resolve_ignore_files from .utils import to_datetime from .utils import vars_to_str @@ -30,14 +32,15 @@ MODELS_SPACE = 'models' -class FilesObject(object): +class FilesObject: """ File / folder object. - It can belong to either a workspace stage or personal/shared space. + It can belong to either a deployment's stage or personal/shared space. This object is not instantiated directly. It is used in the results - of various operations in ``WorkspaceGroup.stage``, ``FilesManager.personal_space``, + of various operations in ``Cluster.stage`` (``WorkspaceGroup.stage`` at + management API v1), ``FilesManager.personal_space``, ``FilesManager.shared_space`` and ``FilesManager.models_space`` methods. """ @@ -89,12 +92,13 @@ def __init__( self.content: List[str] = content or [] self._location: Optional[FileLocation] = None + self._manager: Optional[Manager] = None @classmethod def from_dict( cls, obj: Dict[str, Any], - location: FileLocation, + location: Optional[FileLocation] = None, ) -> FilesObject: """ Construct a FilesObject from a dictionary of values. @@ -123,6 +127,8 @@ def from_dict( writable=bool(obj['writable']), ) out._location = location + if location is not None: + out._manager = location._manager return out def __str__(self) -> str: @@ -352,6 +358,8 @@ class FilesObjectBytesReader(io.BytesIO): class FileLocation(ABC): + _manager: Manager + @abstractmethod def open( self, @@ -391,9 +399,60 @@ def _upload( path: PathLike, *, overwrite: bool = False, - ) -> FilesObject: + fetch_info: bool = True, + ) -> Optional[FilesObject]: pass + def _upload_local_file( + self, + local_path: Union[PathLike, io.IOBase], + path: PathLike, + *, + overwrite: bool = False, + fetch_info: bool = True, + ) -> Optional[FilesObject]: + """ + Upload a local file or open file object to a remote path. + + This is what ``upload_file`` does, minus the return type promise, so + that callers which discard the result -- the Fusion upload handlers -- + can pass ``fetch_info=False`` and save the metadata request that + building a :class:`FilesObject` costs. + + Parameters + ---------- + local_path : Path or str or file-like + Path to the local file or an open file object + path : Path or str + Path to the remote file + overwrite : bool, optional + Should the ``path`` be overwritten if it exists already? + fetch_info : bool, optional + Should the metadata of the uploaded file be fetched and returned? + + Returns + ------- + FilesObject - ``fetch_info`` is True + None - ``fetch_info`` is False + + """ + if isinstance(local_path, io.IOBase): + return self._upload( + local_path, path, + overwrite=overwrite, fetch_info=fetch_info, + ) + + if not os.path.isfile(local_path): + raise IsADirectoryError(f'local path is not a file: {local_path}') + + # The handle has to close even when ``_upload`` raises on a + # non-overwrite conflict, which it does before touching the content. + with open(local_path, 'rb') as infile: + return self._upload( + infile, path, + overwrite=overwrite, fetch_info=fetch_info, + ) + @abstractmethod def mkdir(self, path: PathLike, overwrite: bool = False) -> FilesObject: pass @@ -412,6 +471,33 @@ def rename( def info(self, path: PathLike) -> FilesObject: pass + def _info_or_none(self, path: PathLike) -> Optional[FilesObject]: + """ + Return the metadata of ``path``, or ``None`` if it does not exist. + + This is what :meth:`exists` asks and then throws away. A caller that + goes on to branch on *what* the path is -- ``_upload`` does, on whether + it is a directory -- reads the object instead and so pays for one + request rather than one per question. + + Parameters + ---------- + path : Path or str + Path to the remote object + + Returns + ------- + FilesObject - the path exists + None - it does not + + """ + try: + return self.info(path) + except ManagementError as exc: + if exc.errno == 404: + return None + raise + @abstractmethod def exists(self, path: PathLike) -> bool: pass @@ -469,7 +555,7 @@ def download_file( def download_folder( self, path: PathLike, - local_path: PathLike = '.', + local_path: Optional[PathLike] = None, *, overwrite: bool = False, ) -> None: @@ -517,8 +603,9 @@ class FilesManager(Manager): """ - #: Management API version if none is specified. - default_version = config.get_option('management.version') or 'v1' + # The Files routes are the same at every version, so ``default_version`` + # is inherited from ``Manager`` rather than pinned here: it picks the URL, + # not the implementation. #: Base URL if none is specified. default_base_url = config.get_option('management.base_url') \ @@ -558,7 +645,10 @@ def manage_files( access_token : str, optional The API key or other access token for the files management API version : str, optional - Version of the API to use + Version of the API to use. Defaults to the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment + variable). ``'v1'`` is deprecated and raises a + :class:`DeprecationWarning`. base_url : str, optional Base URL of the files management API organization_id : str, optional @@ -569,9 +659,15 @@ def manage_files( :class:`FilesManager` """ - return FilesManager( + from ._version_import import _import_versioned_module + from ._version_import import _resolve_version + from ._version_import import _warn_if_deprecated_version + ver = _resolve_version(version) + _warn_if_deprecated_version(ver) + mod = _import_versioned_module(ver, 'files') + return mod.FilesManager( access_token=access_token, base_url=base_url, - version=version, organization_id=organization_id, + version=ver, organization_id=organization_id, ) @@ -669,21 +765,10 @@ def upload_file( Should the ``path`` be overwritten if it exists already? """ - if isinstance(local_path, io.IOBase): - pass - elif not os.path.isfile(local_path): - raise IsADirectoryError(f'local path is not a file: {local_path}') - - if self.exists(path): - if not overwrite: - raise OSError(f'file path already exists: {path}') - - self.remove(path) - - if isinstance(local_path, io.IOBase): - return self._upload(local_path, path, overwrite=overwrite) - - return self._upload(open(local_path, 'rb'), path, overwrite=overwrite) + return cast( + FilesObject, + self._upload_local_file(local_path, path, overwrite=overwrite), + ) def upload_folder( self, @@ -714,8 +799,10 @@ def upload_folder( include_root : bool, optional Should the local root folder itself be uploaded as the top folder? ignore : Path or str or List[Path] or List[str], optional - Glob patterns of files to ignore, for example, '**/*.pyc` will - ignore all '*.pyc' files in the directory tree + Glob patterns of files or folders to ignore, for example, + ``**/*.pyc`` will ignore all ``*.pyc`` files in the directory + tree, and ``**/__pycache__`` will ignore those folders entirely. + Relative patterns are resolved against ``local_path``. """ if not os.path.isdir(local_path): @@ -724,30 +811,41 @@ def upload_folder( if not path: path = local_path - ignore_files = set() - if ignore: - if isinstance(ignore, list): - for item in ignore: - ignore_files.update(glob.glob(str(item), recursive=recursive)) - else: - ignore_files.update(glob.glob(str(ignore), recursive=recursive)) + ignore_files = resolve_ignore_files(local_path, ignore) - for dir_path, _, files in os.walk(str(local_path)): + local_root = os.path.normpath(str(local_path)) + root_name = os.path.basename(local_root) + remote_prefix = normalize_remote_path(path, strip_leading=True) + + for dir_path, dirs, files in os.walk(local_root): + if ignore_files: + # Prune ignored folders so their contents are skipped too + dirs[:] = [ + d for d in dirs + if os.path.normpath(os.path.join(dir_path, d)) + not in ignore_files + ] for fname in files: - if ignore_files and fname in ignore_files: + # Normalized so it compares equal to the normalized + # glob results in ignore_files (e.g. local_path='.') + local_file_path = os.path.normpath(os.path.join(dir_path, fname)) + if ignore_files and local_file_path in ignore_files: continue - local_file_path = os.path.join(dir_path, fname) - remote_path = os.path.join( - path, - local_file_path.lstrip(str(local_path)), - ) + rel = os.path.relpath(local_file_path, local_root) + if include_root: + rel = os.path.join(root_name, rel) + # Remote paths always use '/', whatever the local platform + rel = rel.replace(os.sep, '/') + remote_path = f'{remote_prefix}/{rel}' if remote_prefix else rel self.upload_file( local_path=local_file_path, path=remote_path, overwrite=overwrite, ) - return self.info(path) + if not recursive: + break + return self.info(remote_prefix) def _upload( self, @@ -755,7 +853,8 @@ def _upload( path: PathLike, *, overwrite: bool = False, - ) -> FilesObject: + fetch_info: bool = True, + ) -> Optional[FilesObject]: """ Upload content to a file. @@ -767,12 +866,22 @@ def _upload( Path to the file overwrite : bool, optional Should the ``path`` be overwritten if it exists already? + fetch_info : bool, optional + Should the metadata of the uploaded file be fetched and returned? + The write response carries only the name and path, so a + :class:`FilesObject` costs an extra request. """ - if self.exists(path): + # One metadata request, not two: exists() and remove()'s is_dir() are + # the same GET on the same path, so the object is fetched once here and + # every branch reads it. + existing = self._info_or_none(path) + if existing is not None: if not overwrite: raise OSError(f'file path already exists: {path}') - self.remove(path) + if existing.type == 'directory': + raise IsADirectoryError('file path is a directory') + self._manager._delete(f'files/fs/{self._location}/{path}') self._manager._put( f'files/fs/{self._location}/{path}', @@ -780,7 +889,7 @@ def _upload( headers={'Content-Type': None}, ) - return self.info(path) + return self.info(path) if fetch_info else None def mkdir(self, path: PathLike, overwrite: bool = False) -> FilesObject: """ @@ -1020,8 +1129,7 @@ def listdir( List[str] or List[FilesObject] """ - path = re.sub(r'^(\./|/)+', r'', str(path)) - path = re.sub(r'/+$', r'', path) + '/' + path = normalize_remote_path(path, strip_leading=True) + '/' # Validate via listing GET; if response lacks 'content', it's not a directory try: @@ -1136,29 +1244,43 @@ def _download_file( def download_folder( self, path: PathLike, - local_path: PathLike = '.', + local_path: Optional[PathLike] = None, *, overwrite: bool = False, ) -> None: """ Download a FileSpace folder to a local directory. + The contents of ``path`` are written into ``local_path``, which is + created as the destination folder. + Parameters ---------- path : Path or str Directory path - local_path : Path or str - Path to local directory target location + local_path : Path or str, optional + Local directory to create and download into. Defaults to the + name of the ``path`` folder in the current directory. overwrite : bool, optional Should an existing directory / files be overwritten if they exist? """ + # Remote paths always use '/', whatever the local platform + remote_prefix = normalize_remote_path(path, strip_leading=True) + + if local_path is None: + local_path = os.path.basename(remote_prefix) + if not local_path: + raise ValueError( + 'local_path must be specified when downloading ' + 'the root folder', + ) - if local_path is not None and not overwrite and os.path.exists(local_path): + if not overwrite and os.path.exists(local_path): raise OSError('target path already exists; use overwrite=True to replace') # listdir validates directory; no extra info call needed - entries = self.listdir(path, recursive=True, return_objects=True) + entries = self.listdir(remote_prefix, recursive=True, return_objects=True) for entry in entries: # Each entry is a FilesObject with path relative to root and type if not isinstance(entry, FilesObject): # defensive: skip unexpected @@ -1166,14 +1288,18 @@ def download_folder( rel_path = entry.path if entry.type == 'directory': # Ensure local directory exists; no remote call needed - target_dir = os.path.normpath(os.path.join(local_path, rel_path)) + target_dir = ensure_within( + local_path, os.path.join(local_path, rel_path), + ) os.makedirs(target_dir, exist_ok=True) continue - remote_path = os.path.join(path, rel_path) - target_file = os.path.normpath( - os.path.join(local_path, rel_path), + remote_path = ( + f'{remote_prefix}/{rel_path}' if remote_prefix else rel_path + ) + target_file = ensure_within( + local_path, os.path.join(local_path, rel_path), ) - os.makedirs(os.path.dirname(target_file), exist_ok=True) + os.makedirs(os.path.dirname(target_file) or '.', exist_ok=True) self._download_file( remote_path, target_file, overwrite=overwrite, _skip_dir_check=True, diff --git a/singlestoredb/management/job.py b/singlestoredb/management/job.py index e449c0fe8..03be0754e 100644 --- a/singlestoredb/management/job.py +++ b/singlestoredb/management/job.py @@ -1,7 +1,6 @@ #!/usr/bin/env python """SingleStoreDB Cloud Scheduled Notebook Job.""" import datetime -import time from enum import Enum from typing import Any from typing import Dict @@ -10,11 +9,11 @@ from typing import Type from typing import Union +from . import timing from ..exceptions import ManagementError from .manager import Manager from .utils import camel_to_snake from .utils import from_datetime -from .utils import get_cluster_id from .utils import get_database_name from .utils import get_virtual_workspace_id from .utils import get_workspace_id @@ -52,9 +51,35 @@ def __repr__(self) -> str: class TargetType(Enum): + """ + Job target type, spanning both management API versions. + + The wire vocabulary differs by version, and ``'Cluster'`` unhelpfully + means *different things* at each: at v1 it is a legacy self-managed + cluster, at v2 it is the resource that v1 called a workspace. This enum + holds the union so the read path (:meth:`from_str`) round-trips either + version's value without having to know which version produced it. The + *write* path is version-specific -- see ``JobsManager._resolve_target``. + + ========================== ========= =================================== + Value Versions Meaning + ========================== ========= =================================== + ``'Workspace'`` v1 Workspace + ``'VirtualWorkspace'`` v1 Starter (shared tier) workspace + ``'Cluster'`` v1 Legacy self-managed cluster + ``'Cluster'`` v2 Cluster (the v1 "workspace") + ``'VirtualCluster'`` v2 Starter (shared tier) cluster + ========================== ========= =================================== + + Only the read path ever sees the v1 sense of ``'Cluster'``: the write path + takes its target from the notebook environment, which names a workspace or + starter deployment and never a legacy self-managed cluster. + """ + WORKSPACE = 'Workspace' CLUSTER = 'Cluster' VIRTUAL_WORKSPACE = 'VirtualWorkspace' + VIRTUAL_CLUSTER = 'VirtualCluster' @classmethod def from_str(cls, s: str) -> 'TargetType': @@ -545,7 +570,7 @@ def __repr__(self) -> str: return str(self) -class Job(object): +class Job: """ Scheduled Notebook Job definition. @@ -634,9 +659,9 @@ def wait(self, timeout: Optional[int] = None) -> bool: return self._manager._wait_for_job(self, timeout) def get_executions( - self, - start_execution_number: int, - end_execution_number: int, + self, + start_execution_number: int, + end_execution_number: int, ) -> ExecutionsData: """Get executions for the job.""" if self._manager is None: @@ -668,7 +693,7 @@ def __repr__(self) -> str: return str(self) -class JobsManager(object): +class JobsManager: """ SingleStoreDB scheduled notebook jobs manager. @@ -676,17 +701,49 @@ class JobsManager(object): Parameters ---------- - manager : WorkspaceManager, optional - The WorkspaceManager the JobsManager belongs to + manager : ClusterManager, optional + The ClusterManager the JobsManager belongs to See Also -------- :attr:`Organization.jobs` """ + #: ``targetType`` sent for a regular deployment. This is the only part of + #: the jobs API whose vocabulary changed at v2 (``'Workspace'`` became + #: ``'Cluster'``), so the v1 subclass overrides these three attributes + #: instead of reimplementing ``schedule``. + _deployment_target_type = TargetType.CLUSTER + + #: ``targetType`` sent for a starter / shared-tier deployment. + _starter_target_type = TargetType.VIRTUAL_CLUSTER + def __init__(self, manager: Optional[Manager]): self._manager = manager + def _resolve_target(self, target_config: Dict[str, Any]) -> None: + """ + Fill in ``targetID`` / ``targetType`` from the ambient environment. + + The deployment the job should run against is taken from the + environment variables set by the notebook runtime: + ``SINGLESTOREDB_VIRTUAL_WORKSPACE`` for a starter deployment, and + ``SINGLESTOREDB_WORKSPACE`` for a regular one -- the latter holds a + cluster ID at v2 and a workspace ID at v1. Which ``targetType`` string + names each kind of deployment is version-specific; see the + ``_*_target_type`` class attributes. + """ + starter_id = get_virtual_workspace_id() + deployment_id = get_workspace_id() + + if starter_id is not None: + target_config['targetID'] = starter_id + target_config['targetType'] = self._starter_target_type.value + + elif deployment_id is not None: + target_config['targetID'] = deployment_id + target_config['targetType'] = self._deployment_target_type.value + def schedule( self, notebook_path: str, @@ -699,6 +756,7 @@ def schedule( runtime_name: Optional[str] = None, resume_target: Optional[bool] = None, parameters: Optional[Dict[str, Any]] = None, + max_allowed_execution_duration_in_minutes: Optional[int] = None, ) -> Job: """Creates and returns a scheduled notebook job.""" if self._manager is None: @@ -722,6 +780,10 @@ def schedule( if runtime_name is not None: execution_config['runtimeName'] = runtime_name + if max_allowed_execution_duration_in_minutes is not None: + execution_config['maxAllowedExecutionDurationInMinutes'] = \ + max_allowed_execution_duration_in_minutes + target_config = None # type: Optional[Dict[str, Any]] database_name = get_database_name() if database_name is not None: @@ -732,20 +794,7 @@ def schedule( if resume_target is not None: target_config['resumeTarget'] = resume_target - workspace_id = get_workspace_id() - virtual_workspace_id = get_virtual_workspace_id() - cluster_id = get_cluster_id() - if virtual_workspace_id is not None: - target_config['targetID'] = virtual_workspace_id - target_config['targetType'] = TargetType.VIRTUAL_WORKSPACE.value - - elif workspace_id is not None: - target_config['targetID'] = workspace_id - target_config['targetType'] = TargetType.WORKSPACE.value - - elif cluster_id is not None: - target_config['targetID'] = cluster_id - target_config['targetType'] = TargetType.CLUSTER.value + self._resolve_target(target_config) job_run_json = dict( schedule=schedule, @@ -832,7 +881,7 @@ def _wait_for_job(self, job: Union[str, Job], timeout: Optional[int] = None) -> return True if job.schedule.mode == Mode.RECURRING: raise ValueError(f'Cannot wait for recurring job {job_id}') - time.sleep(5) + timing.sleep(5, 'job completion') def get(self, job_id: str) -> Job: """Get a job by its ID.""" @@ -843,10 +892,10 @@ def get(self, job_id: str) -> Job: return Job.from_dict(res, self) def get_executions( - self, - job_id: str, - start_execution_number: int, - end_execution_number: int, + self, + job_id: str, + start_execution_number: int, + end_execution_number: int, ) -> ExecutionsData: """Get executions for a job by its ID.""" if self._manager is None: diff --git a/singlestoredb/management/manager.py b/singlestoredb/management/manager.py index 575df0876..37ba35708 100644 --- a/singlestoredb/management/manager.py +++ b/singlestoredb/management/manager.py @@ -7,14 +7,19 @@ from typing import Dict from typing import List from typing import Optional +from typing import Tuple from typing import Union from urllib.parse import urljoin import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry +from . import timing from .. import config from ..exceptions import ManagementError from ..exceptions import OperationalError +from ._version_import import DEFAULT_VERSION from .utils import get_token @@ -30,6 +35,59 @@ def set_organization(kwargs: Dict[str, Any]) -> None: kwargs['params']['organizationID'] = org +#: Methods that may be replayed after a transport-level failure. POST is +#: absent on purpose: a dropped connection does not say whether the server +#: acted on the request, and replaying ``POST /clusters`` would deploy twice. +#: Everything the long ``wait_on_*`` loops issue is a GET, so the retries +#: cover the failure mode that actually shows up -- a keep-alive connection +#: the far end closed while the client was sleeping between polls, which +#: surfaces as ``RemoteDisconnected`` on the next request. +RETRY_METHODS = frozenset(['GET', 'HEAD', 'OPTIONS', 'PUT', 'DELETE']) + +#: Status codes worth retrying. These are the transient ones; a 4xx other +#: than 429 is a client error that will fail again identically. +RETRY_STATUSES = frozenset([429, 500, 502, 503, 504]) + + +def build_retry( + total: Optional[int] = None, + backoff_factor: Optional[float] = None, +) -> Retry: + """Build the retry policy used by every manager session.""" + if total is None: + total = int(os.environ.get('SINGLESTOREDB_MANAGEMENT_RETRIES', '4')) + if backoff_factor is None: + backoff_factor = float( + os.environ.get('SINGLESTOREDB_MANAGEMENT_RETRY_BACKOFF', '0.5'), + ) + return Retry( + total=total, + connect=total, + read=total, + status=total, + allowed_methods=RETRY_METHODS, + status_forcelist=RETRY_STATUSES, + backoff_factor=backoff_factor, + # Let ``Manager._check`` raise the error with the response body in it + # rather than urllib3 raising a bare MaxRetryError. + raise_on_status=False, + respect_retry_after_header=True, + ) + + +def default_timeout() -> Tuple[float, float]: + """ + Return the (connect, read) timeout applied when a caller gives none. + + Without this a stalled connection hangs the client forever instead of + failing and being retried. + """ + return ( + float(os.environ.get('SINGLESTOREDB_MANAGEMENT_CONNECT_TIMEOUT', '10')), + float(os.environ.get('SINGLESTOREDB_MANAGEMENT_READ_TIMEOUT', '180')), + ) + + def is_jwt(token: str) -> bool: """Is the given token a JWT?""" import jwt @@ -40,11 +98,18 @@ def is_jwt(token: str) -> bool: return False -class Manager(object): +class Manager: """SingleStoreDB manager base class.""" - #: Management API version if none is specified. - default_version = config.get_option('management.version') or 'v1' + #: Management API version if none is specified. The shared + #: :data:`~singlestoredb.management._version_import.DEFAULT_VERSION`, which + #: also supplies the ``management.version`` option default, so the two + #: cannot drift. Deliberately not a reading of that option: it is read by + #: the ``manage_*`` factories at call time, and reading it here would let a + #: version-specific class declare itself to be whatever the option happened + #: to say. A class that implements one specific version pins that version + #: as a literal instead of inheriting this. + default_version = DEFAULT_VERSION #: Base URL if none is specified. default_base_url = config.get_option('management.base_url') \ @@ -64,8 +129,17 @@ def __init__( if not new_access_token: raise ManagementError(msg='No management token was configured.') + base_url_root = ( + base_url + or config.get_option('management.base_url') + or type(self).default_base_url + ) + self._is_jwt = not access_token and new_access_token and is_jwt(new_access_token) self._sess = requests.Session() + adapter = HTTPAdapter(max_retries=build_retry()) + self._sess.mount('http://', adapter) + self._sess.mount('https://', adapter) self._sess.headers.update({ 'Authorization': f'Bearer {new_access_token}', 'Content-Type': 'application/json', @@ -74,9 +148,7 @@ def __init__( }) self._base_url = urljoin( - base_url - or config.get_option('management.base_url') - or type(self).default_base_url, + base_url_root, version or type(self).default_version, ) + '/' @@ -126,9 +198,31 @@ def _doit( # Refresh the JWT as needed if self._is_jwt: self._sess.headers.update({'Authorization': f'Bearer {get_token()}'}) - return getattr(self._sess, method.lower())( - urljoin(self._base_url, path), *args, **kwargs, + kwargs.setdefault('timeout', default_timeout()) + url = urljoin(self._base_url, path) + # Every management HTTP call comes through here, so this is the one + # place request time has to be recorded. See management.timing. + started_at = time.monotonic() + try: + res = getattr(self._sess, method.lower())(url, *args, **kwargs) + except requests.exceptions.RequestException as exc: + timing.record_request( + method, path, time.monotonic() - started_at, started_at, + error=exc, + ) + # A transport failure otherwise escapes as a bare + # requests.ConnectionError / ReadTimeout naming neither the route + # nor the method, which makes it indistinguishable from a bug in + # the caller. Retries for the replayable methods are already + # exhausted by the time this is reached. + raise ManagementError( + msg=f'{type(exc).__name__} on {method.upper()} {url}: {exc}', + ) from exc + timing.record_request( + method, path, time.monotonic() - started_at, started_at, + response=res, ) + return res def _get(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: """ @@ -298,17 +392,24 @@ def _wait_on_state( ), ) + remaining = float(timeout) while True: if getattr(out, 'state').lower() in states: break - if timeout <= 0: + if remaining <= 0: raise ManagementError( msg=f'Exceeded waiting time for {self.obj_type} to become ' '{}.'.format(', '.join(states)), ) - time.sleep(interval) - timeout -= interval + started_at = timing.now() + timing.sleep( + interval, + '{} state -> {}'.format(self.obj_type, ', '.join(states)), + ) out = getattr(self, f'get_{self.obj_type}')(out.id) + # Charged after the refetch, and by measured time: the refetch is + # part of what the iteration cost. See timing.poll_cost. + remaining -= timing.poll_cost(started_at, interval) return out @@ -324,7 +425,8 @@ def _wait_on_endpoint( Parameters ---------- out : Any - Workspace object with a connect method + Deployment object with a connect method -- a ``Cluster`` or + ``StarterCluster`` at v2, a ``Workspace`` at v1 interval : int, optional Interval between each connection attempt (default: 10 seconds) timeout : int, optional @@ -351,23 +453,32 @@ def _wait_on_endpoint( msg=f'{type(out).__name__} object does not have a valid endpoint', ) + remaining = float(timeout) while True: + started_at = timing.now() try: # Try to establish a connection to the endpoint using context manager - with out.connect(connect_timeout=5): - pass + with timing.timed(f'{self.obj_type} endpoint connect'): + with out.connect(connect_timeout=5): + pass + # Connected, so the endpoint is ready. Without this the loop + # reconnects forever on success and only ever leaves through + # the 1045 branch or the timeout. + break except Exception as exc: # If we get an 'access denied' error, that means that the server is # up and we just aren't authenticating. if isinstance(exc, OperationalError) and exc.errno == 1045: break # If connection fails, check timeout and retry - if timeout <= 0: + if remaining <= 0: raise ManagementError( msg=f'Exceeded waiting time for {self.obj_type} endpoint ' 'to become ready', ) - time.sleep(interval) - timeout -= interval + timing.sleep(interval, f'{self.obj_type} endpoint') + # The failed connect attempt is part of what the iteration + # cost: connect_timeout is 5 seconds on top of the sleep. + remaining -= timing.poll_cost(started_at, interval) return out diff --git a/singlestoredb/management/organization.py b/singlestoredb/management/organization.py index 2c0f917df..a330023dc 100644 --- a/singlestoredb/management/organization.py +++ b/singlestoredb/management/organization.py @@ -1,18 +1,61 @@ #!/usr/bin/env python """SingleStoreDB Cloud Organization.""" import datetime +from typing import Any from typing import Dict from typing import List from typing import Optional +from typing import Type from typing import Union from ..exceptions import ManagementError -from .inference_api import InferenceAPIManager +from ._version_import import _versioned_attr from .job import JobsManager from .manager import Manager +from .utils import to_datetime from .utils import vars_to_str +def get_organization(version: Optional[str] = None) -> 'Organization': + """ + Get the current organization. + + Parameters + ---------- + version : str, optional + Version of the API to use. Defaults to the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment + variable). + + Returns + ------- + :class:`Organization` + + """ + return _versioned_attr('get_organization', version)() + + +def get_secret(name: str, version: Optional[str] = None) -> Optional[str]: + """ + Get the value of a secret in the current organization. + + Parameters + ---------- + name : str + Name of the secret + version : str, optional + Version of the API to use. Defaults to the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment + variable). + + Returns + ------- + str or None + + """ + return _versioned_attr('get_secret', version)(name) + + def listify(x: Union[str, List[str]]) -> List[str]: if isinstance(x, list): return x @@ -38,9 +81,9 @@ def __init__( id: str, name: str, created_by: str, - created_at: Union[str, datetime.datetime], + created_at: Optional[Union[str, datetime.datetime]], last_updated_by: str, - last_updated_at: Union[str, datetime.datetime], + last_updated_at: Optional[Union[str, datetime.datetime]], value: Optional[str] = None, deleted_by: Optional[str] = None, deleted_at: Optional[Union[str, datetime.datetime]] = None, @@ -91,12 +134,12 @@ def from_dict(cls, obj: Dict[str, str]) -> 'Secret': id=obj['secretID'], name=obj['name'], created_by=obj['createdBy'], - created_at=obj['createdAt'], + created_at=to_datetime(obj.get('createdAt')), last_updated_by=obj['lastUpdatedBy'], - last_updated_at=obj['lastUpdatedAt'], + last_updated_at=to_datetime(obj.get('lastUpdatedAt')), value=obj.get('value'), deleted_by=obj.get('deletedBy'), - deleted_at=obj.get('deletedAt'), + deleted_at=to_datetime(obj.get('deletedAt')), ) return out @@ -110,16 +153,16 @@ def __repr__(self) -> str: return str(self) -class Organization(object): +class Organization: """ Organization in SingleStoreDB Cloud portal. This object is not directly instantiated. It is used in results - of ``WorkspaceManager`` API calls. + of ``ClusterManager`` API calls. See Also -------- - :attr:`WorkspaceManager.organization` + :attr:`ClusterManager.organization` """ @@ -127,8 +170,19 @@ class Organization(object): name: str firewall_ranges: List[str] + #: Sub-manager classes reached through this organization. The + #: ``organizations/current`` and ``secrets`` routes are identical at v1 and + #: v2, so ``Organization`` itself is version-neutral; only the managers it + #: hands out differ. These name the current-version managers, and + #: ``v1/organization.py`` repoints them back to the v1 classes. + _jobs_manager_class: Type[JobsManager] = JobsManager + + #: Inference API manager class, or ``None`` if the version has no + #: inference routes. There are none from v2 onward. + _inference_api_manager_class: Optional[Type[Any]] = None + def __init__(self, id: str, name: str, firewall_ranges: List[str]): - """Use :attr:`WorkspaceManager.organization` instead.""" + """Use :attr:`ClusterManager.organization` instead.""" #: Unique ID of the organization self.id = id @@ -177,8 +231,8 @@ def from_dict( ---------- obj : dict Key-value pairs to retrieve organization information from - manager : WorkspaceManager, optional - The WorkspaceManager the Organization belongs to + manager : ClusterManager, optional + The ClusterManager the Organization belongs to Returns ------- @@ -200,27 +254,51 @@ def jobs(self) -> JobsManager: Parameters ---------- - manager : WorkspaceManager, optional - The WorkspaceManager the JobsManager belongs to + manager : ClusterManager, optional + The ClusterManager the JobsManager belongs to Returns ------- :class:`JobsManager` """ - return JobsManager(self._manager) + return self._jobs_manager_class(self._manager) @property - def inference_apis(self) -> InferenceAPIManager: + def inference_apis(self) -> Any: """ Retrieve a SingleStoreDB inference api manager. - Parameters - ---------- - manager : WorkspaceManager, optional - The WorkspaceManager the InferenceAPIManager belongs to - Returns ------- :class:`InferenceAPIManager` + + Raises + ------ + ManagementError + If the API version has no inference routes + """ - return InferenceAPIManager(self._manager) + if self._inference_api_manager_class is None: + raise ManagementError( + msg='The inference API is not available in this version of ' + 'the management API. None of the inferenceapis/ routes ' + 'exist past v1.', + ) + return self._inference_api_manager_class(self._manager) + + +class Organizations(object): + """Organizations.""" + + #: The ``Organization`` class this hands out. Version subclasses repoint + #: this so the organization carries the right sub-managers. + _organization_class: Type[Organization] = Organization + + def __init__(self, manager: Manager): + self._manager = manager + + @property + def current(self) -> Organization: + """Get current organization.""" + res = self._manager._get('organizations/current').json() + return self._organization_class.from_dict(res, self._manager) diff --git a/singlestoredb/management/project.py b/singlestoredb/management/project.py new file mode 100644 index 000000000..800214a9a --- /dev/null +++ b/singlestoredb/management/project.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python +""" +SingleStoreDB Project Management. + +Projects are only addressed by the v2 wrappers -- ``POST /v2/clusters`` +requires a ``projectID`` where the v1 workspace group route assigned one +implicitly -- so the implementation lives in +:mod:`singlestoredb.management.v2.project` and this module is a stable import +path for it, the same arrangement :mod:`singlestoredb.management.cluster` uses. +""" +from .v2.project import Project as Project diff --git a/singlestoredb/management/region.py b/singlestoredb/management/region.py index 7bc39a7ec..f5737ea06 100644 --- a/singlestoredb/management/region.py +++ b/singlestoredb/management/region.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -"""SingleStoreDB Cluster Management.""" +"""SingleStoreDB Region Management.""" from typing import Dict from typing import Optional @@ -8,16 +8,16 @@ from .utils import vars_to_str -class Region(object): +class Region: """ Cluster region information. This object is not directly instantiated. It is used in results - of ``WorkspaceManager`` API calls. + of ``ClusterManager`` API calls. See Also -------- - :attr:`WorkspaceManager.regions` + :attr:`ClusterManager.regions` """ @@ -25,7 +25,7 @@ def __init__( self, name: str, provider: str, id: Optional[str] = None, region_name: Optional[str] = None, ) -> None: - """Use :attr:`WorkspaceManager.regions` instead.""" + """Use :attr:`ClusterManager.regions` instead.""" #: Unique ID of the region self.id = id @@ -57,8 +57,8 @@ def from_dict(cls, obj: Dict[str, str], manager: Manager) -> 'Region': ---------- obj : dict Key-value pairs to retrieve region information from - manager : WorkspaceManager, optional - The WorkspaceManager the Region belongs to + manager : ClusterManager, optional + The ClusterManager the Region belongs to Returns ------- @@ -87,11 +87,11 @@ class RegionManager(Manager): Parameters ---------- access_token : str, optional - The API key or other access token for the workspace management API + 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 workspace management API + Base URL of the management API See Also -------- @@ -122,17 +122,18 @@ def list_regions(self) -> NamedList[Region]: def list_shared_tier_regions(self) -> NamedList[Region]: """ - List regions that support shared tier workspaces. + List regions that support shared tier deployments. Returns ------- NamedList[Region] - List of regions that support shared tier workspaces + List of regions that support shared tier deployments Raises ------ ManagementError If there is an error getting the regions + """ res = self._get('regions/sharedtier') return NamedList( @@ -151,19 +152,28 @@ def manage_regions( Parameters ---------- access_token : str, optional - The API key or other access token for the workspace management API + The API key or other access token for the management API version : str, optional - Version of the API to use + Version of the API to use. Defaults to the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment + variable). ``'v1'`` is deprecated and raises a + :class:`DeprecationWarning`. base_url : str, optional - Base URL of the workspace management API + Base URL of the management API Returns ------- :class:`RegionManager` """ - return RegionManager( + from ._version_import import _import_versioned_module + from ._version_import import _resolve_version + from ._version_import import _warn_if_deprecated_version + ver = _resolve_version(version) + _warn_if_deprecated_version(ver) + mod = _import_versioned_module(ver, 'region') + return mod.RegionManager( access_token=access_token, - version=version, + version=ver, base_url=base_url, ) diff --git a/singlestoredb/management/stage.py b/singlestoredb/management/stage.py new file mode 100644 index 000000000..ac724b35e --- /dev/null +++ b/singlestoredb/management/stage.py @@ -0,0 +1,788 @@ +#!/usr/bin/env python +""" +SingleStoreDB Stage management. + +Stage is version-neutral apart from where it hangs off the API: at v1 the +filesystem lives under ``stage/{deployment_id}/fs/...``, at v2 it moved under +the cluster resource (``clusters/{cluster_id}/stage/fs/...``). Every request +this class makes routes through :meth:`Stage._fs_path`, so the version +difference is a one-line override in the v2 subclass rather than a copy of +every method. +""" +from __future__ import annotations + +import io +import os +import re +from typing import Any +from typing import cast +from typing import List +from typing import Literal +from typing import Optional +from typing import overload +from typing import Union + +from ..exceptions import ManagementError +from ._version_import import _versioned_attr +from .files import FileLocation +from .files import FilesObject +from .files import FilesObjectBytesReader +from .files import FilesObjectBytesWriter +from .files import FilesObjectTextReader +from .files import FilesObjectTextWriter +from .manager import Manager +from .utils import ensure_within +from .utils import normalize_remote_path +from .utils import PathLike +from .utils import resolve_ignore_files +from .utils import vars_to_str + + +def get_stage( + deployment: Optional[Any] = None, + version: Optional[str] = None, +) -> 'Stage': + """ + Get the stage of a deployment. + + Parameters + ---------- + deployment : Cluster or WorkspaceGroup or str, optional + The deployment whose stage is wanted, or its name or ID. What counts + as a deployment is version-specific: a cluster at v2, a workspace + group at v1. If not given, the deployment named by the environment is + used -- ``SINGLESTOREDB_WORKSPACE_GROUP`` at v1, and + ``SINGLESTOREDB_WORKSPACE`` at v2. + version : str, optional + Version of the API to use. Defaults to the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment + variable). + + Returns + ------- + :class:`Stage` + + """ + return _versioned_attr('get_stage', version)(deployment) + + +class Stage(FileLocation): + """ + Stage manager. + + This object is not instantiated directly. + It is returned by ``Cluster.stage`` or ``StarterCluster.stage``. + + """ + + def __init__(self, deployment_id: str, manager: Manager): + self._deployment_id = deployment_id + self._manager = manager + + def _fs_path(self, path: PathLike = '') -> str: + """ + Return the management API path for a Stage filesystem location. + + Overridden by the v1 ``Stage``, where Stage was a top-level resource + rather than nested under the cluster. All Stage requests go through + here so that the version difference is a one-line override rather + than a copy of every method. + + Parameters + ---------- + path : Path or str, optional + Stage path, relative to the root of the deployment's Stage + + Returns + ------- + str + + """ + return f'clusters/{self._deployment_id}/stage/fs/{path}' + + def open( + self, + stage_path: PathLike, + mode: str = 'r', + encoding: Optional[str] = None, + ) -> Union[io.StringIO, io.BytesIO]: + """ + Open a Stage path for reading or writing. + + Parameters + ---------- + stage_path : Path or str + The stage path to read / write + mode : str, optional + The read / write mode. The following modes are supported: + * 'r' open for reading (default) + * 'w' open for writing, truncating the file first + * 'x' create a new file and open it for writing + The data type can be specified by adding one of the following: + * 'b' binary mode + * 't' text mode (default) + encoding : str, optional + The string encoding to use for text + + Returns + ------- + FilesObjectBytesReader - 'rb' or 'b' mode + FilesObjectBytesWriter - 'wb' or 'xb' mode + FilesObjectTextReader - 'r' or 'rt' mode + FilesObjectTextWriter - 'w', 'x', 'wt' or 'xt' mode + + """ + if '+' in mode or 'a' in mode: + raise ValueError('modifying an existing stage file is not supported') + + if 'w' in mode or 'x' in mode: + exists = self.exists(stage_path) + if exists: + if 'x' in mode: + raise FileExistsError(f'stage path already exists: {stage_path}') + self.remove(stage_path) + if 'b' in mode: + return FilesObjectBytesWriter(b'', self, stage_path) + return FilesObjectTextWriter('', self, stage_path) + + if 'r' in mode: + content = self.download_file(stage_path) + if isinstance(content, bytes): + if 'b' in mode: + return FilesObjectBytesReader(content) + encoding = 'utf-8' if encoding is None else encoding + return FilesObjectTextReader(content.decode(encoding)) + + if isinstance(content, str): + return FilesObjectTextReader(content) + + raise ValueError(f'unrecognized file content type: {type(content)}') + + raise ValueError(f'must have one of create/read/write mode specified: {mode}') + + def upload_file( + self, + local_path: Union[PathLike, io.IOBase], + stage_path: PathLike, + *, + overwrite: bool = False, + ) -> FilesObject: + """ + Upload a local file. + + Parameters + ---------- + local_path : Path or str or file-like + Path to the local file or an open file object + stage_path : Path or str + Path to the stage file + overwrite : bool, optional + Should the ``stage_path`` be overwritten if it exists already? + + """ + return cast( + FilesObject, + self._upload_local_file(local_path, stage_path, overwrite=overwrite), + ) + + def upload_folder( + self, + local_path: PathLike, + stage_path: PathLike, + *, + overwrite: bool = False, + recursive: bool = True, + include_root: bool = False, + ignore: Optional[Union[PathLike, List[PathLike]]] = None, + ) -> FilesObject: + """ + Upload a folder recursively. + + Only the contents of the folder are uploaded. To include the + folder name itself in the target path use ``include_root=True``. + + Parameters + ---------- + local_path : Path or str + Local directory to upload + stage_path : Path or str + Path of stage folder to upload to + overwrite : bool, optional + If a file already exists, should it be overwritten? + recursive : bool, optional + Should nested folders be uploaded? + include_root : bool, optional + Should the local root folder itself be uploaded as the top folder? + ignore : Path or str or List[Path] or List[str], optional + Glob patterns of files or folders to ignore, for example, + ``**/*.pyc`` will ignore all ``*.pyc`` files in the directory + tree, and ``**/__pycache__`` will ignore those folders entirely. + Relative patterns are resolved against ``local_path``. + + """ + if not os.path.isdir(local_path): + raise NotADirectoryError(f'local path is not a directory: {local_path}') + + stage_prefix = normalize_remote_path(stage_path, strip_leading=True) + + if self.exists(stage_prefix) and not self.is_dir(stage_prefix): + raise NotADirectoryError(f'stage path is not a directory: {stage_path}') + + ignore_files = resolve_ignore_files(local_path, ignore) + + local_root = os.path.normpath(str(local_path)) + root_name = os.path.basename(local_root) + + for dir_path, dirs, files in os.walk(local_root): + if ignore_files: + # Prune ignored folders so their contents are skipped too + dirs[:] = [ + d for d in dirs + if os.path.normpath(os.path.join(dir_path, d)) + not in ignore_files + ] + for fname in files: + # Normalized so it compares equal to the normalized + # glob results in ignore_files (e.g. local_path='.') + local_file_path = os.path.normpath(os.path.join(dir_path, fname)) + if ignore_files and local_file_path in ignore_files: + continue + rel = os.path.relpath(local_file_path, local_root) + if include_root: + rel = os.path.join(root_name, rel) + # Remote paths always use '/', whatever the local platform + rel = rel.replace(os.sep, '/') + target = f'{stage_prefix}/{rel}' if stage_prefix else rel + self.upload_file(local_file_path, target, overwrite=overwrite) + if not recursive: + break + + return self.info(stage_prefix) + + def _upload( + self, + content: Union[str, bytes, io.IOBase], + stage_path: PathLike, + *, + overwrite: bool = False, + fetch_info: bool = True, + ) -> Optional[FilesObject]: + """ + Upload content to a stage file. + + Parameters + ---------- + content : str or bytes or file-like + Content to upload to stage + stage_path : Path or str + Path to the stage file + overwrite : bool, optional + Should the ``stage_path`` be overwritten if it exists already? + fetch_info : bool, optional + Should the metadata of the uploaded file be fetched and returned? + The write response carries only the name and path, so a + :class:`FilesObject` costs an extra request. + + """ + # One metadata request, not two: exists() and remove()'s is_dir() are + # the same GET on the same path, so the object is fetched once here and + # every branch reads it. + existing = self._info_or_none(stage_path) + if existing is not None: + if not overwrite: + raise OSError(f'stage path already exists: {stage_path}') + if existing.type == 'directory': + raise IsADirectoryError( + 'stage path is a directory, ' + f'use rmdir or removedirs: {stage_path}', + ) + self._manager._delete(self._fs_path(stage_path)) + + self._manager._put( + self._fs_path(stage_path), + files={'file': content}, + headers={'Content-Type': None}, + ) + + return self.info(stage_path) if fetch_info else None + + def mkdir(self, stage_path: PathLike, overwrite: bool = False) -> FilesObject: + """ + Make a directory in the stage. + + Parameters + ---------- + stage_path : Path or str + Path of the folder to create + overwrite : bool, optional + Should the stage path be overwritten if it exists already? + + Returns + ------- + FilesObject + + """ + stage_path = re.sub(r'/*$', r'', str(stage_path)) + '/' + + if self.exists(stage_path): + if not overwrite: + return self.info(stage_path) + + self.remove(stage_path) + + self._manager._put( + self._fs_path(stage_path) + '?isFile=false', + ) + + return self.info(stage_path) + + mkdirs = mkdir + + def rename( + self, + old_path: PathLike, + new_path: PathLike, + *, + overwrite: bool = False, + ) -> FilesObject: + """ + Move the stage file to a new location. + + Paraemeters + ----------- + old_path : Path or str + Original location of the path + new_path : Path or str + New location of the path + overwrite : bool, optional + Should the ``new_path`` be overwritten if it exists already? + + """ + if not self.exists(old_path): + raise OSError(f'stage path does not exist: {old_path}') + + if self.exists(new_path): + if not overwrite: + raise OSError(f'stage path already exists: {new_path}') + + if str(old_path).endswith('/') and not str(new_path).endswith('/'): + raise OSError('original and new paths are not the same type') + + if str(new_path).endswith('/'): + self.removedirs(new_path) + else: + self.remove(new_path) + + self._manager._patch( + self._fs_path(old_path), + json=dict(newPath=new_path), + ) + + return self.info(new_path) + + def info(self, stage_path: PathLike) -> FilesObject: + """ + Return information about a stage location. + + Parameters + ---------- + stage_path : Path or str + Path to the stage location + + Returns + ------- + FilesObject + + """ + res = self._manager._get( + re.sub(r'/+$', r'/', self._fs_path(stage_path)), + params=dict(metadata=1), + ).json() + + return FilesObject.from_dict(res, self) + + def exists(self, stage_path: PathLike) -> bool: + """ + Does the given stage path exist? + + Parameters + ---------- + stage_path : Path or str + Path to stage object + + Returns + ------- + bool + + """ + try: + self.info(stage_path) + return True + except ManagementError as exc: + if exc.errno == 404: + return False + raise + + def is_dir(self, stage_path: PathLike) -> bool: + """ + Is the given stage path a directory? + + Parameters + ---------- + stage_path : Path or str + Path to stage object + + Returns + ------- + bool + + """ + try: + return self.info(stage_path).type == 'directory' + except ManagementError as exc: + if exc.errno == 404: + return False + raise + + def is_file(self, stage_path: PathLike) -> bool: + """ + Is the given stage path a file? + + Parameters + ---------- + stage_path : Path or str + Path to stage object + + Returns + ------- + bool + + """ + try: + return self.info(stage_path).type != 'directory' + except ManagementError as exc: + if exc.errno == 404: + return False + raise + + def _listdir( + self, stage_path: PathLike, *, + recursive: bool = False, + return_objects: bool = False, + ) -> List[Union[str, 'FilesObject']]: + """ + Return the names (or FilesObject instances) of files in a directory. + + Parameters + ---------- + stage_path : Path or str + Path to the folder in Stage + recursive : bool, optional + Should folders be listed recursively? + return_objects : bool, optional + If True, return list of FilesObject instances. Otherwise just paths. + + """ + from .files import FilesObject + res = self._manager._get( + re.sub(r'/+$', r'/', self._fs_path(stage_path)), + ).json() + if recursive: + out: List[Union[str, FilesObject]] = [] + for item in res['content'] or []: + if return_objects: + out.append(FilesObject.from_dict(item, self)) + else: + out.append(item['path']) + if item['type'] == 'directory': + out.extend( + self._listdir( + item['path'], + recursive=recursive, + return_objects=return_objects, + ), + ) + return out + if return_objects: + return [ + FilesObject.from_dict(x, self) + for x in res['content'] or [] + ] + return [x['path'] for x in res['content'] or []] + + @overload + def listdir( + self, + stage_path: PathLike = '/', + *, + recursive: bool = False, + return_objects: Literal[True], + ) -> List['FilesObject']: + ... + + @overload + def listdir( + self, + stage_path: PathLike = '/', + *, + recursive: bool = False, + return_objects: Literal[False] = False, + ) -> List[str]: + ... + + def listdir( + self, + stage_path: PathLike = '/', + *, + recursive: bool = False, + return_objects: bool = False, + ) -> Union[List[str], List['FilesObject']]: + """ + List the files / folders at the given path. + + Parameters + ---------- + stage_path : Path or str, optional + Path to the stage location + recursive : bool, optional + If True, recursively list all files and folders + return_objects : bool, optional + If True, return list of FilesObject instances. Otherwise just paths. + + Returns + ------- + List[str] or List[FilesObject] + + """ + from .files import FilesObject + stage_path = normalize_remote_path(stage_path, strip_leading=True) + '/' + + if self.is_dir(stage_path): + out = self._listdir( + stage_path, + recursive=recursive, + return_objects=return_objects, + ) + if stage_path != '/': + stage_path_n = len(stage_path.split('/')) - 1 + if return_objects: + result: List[FilesObject] = [] + for item in out: + if isinstance(item, FilesObject): + rel = '/'.join(item.path.split('/')[stage_path_n:]) + item.path = rel + result.append(item) + return result + out = ['/'.join(str(x).split('/')[stage_path_n:]) for x in out] + if return_objects: + return cast(List[FilesObject], out) + return cast(List[str], out) + + raise NotADirectoryError(f'stage path is not a directory: {stage_path}') + + def download_file( + self, + stage_path: PathLike, + local_path: Optional[PathLike] = None, + *, + overwrite: bool = False, + encoding: Optional[str] = None, + ) -> Optional[Union[bytes, str]]: + """ + Download the content of a stage path. + + Parameters + ---------- + stage_path : Path or str + Path to the stage file + local_path : Path or str + Path to local file target location + overwrite : bool, optional + Should an existing file be overwritten if it exists? + encoding : str, optional + Encoding used to convert the resulting data + + Returns + ------- + bytes or str - ``local_path`` is None + None - ``local_path`` is a Path or str + + """ + return self._download_file( + stage_path, + local_path=local_path, + overwrite=overwrite, + encoding=encoding, + _skip_dir_check=False, + ) + + def _download_file( + self, + stage_path: PathLike, + local_path: Optional[PathLike] = None, + *, + overwrite: bool = False, + encoding: Optional[str] = None, + _skip_dir_check: bool = False, + ) -> Optional[Union[bytes, str]]: + """ + Internal method to download the content of a stage path. + + Parameters + ---------- + stage_path : Path or str + Path to the stage file + local_path : Path or str + Path to local file target location + overwrite : bool, optional + Should an existing file be overwritten if it exists? + encoding : str, optional + Encoding used to convert the resulting data + _skip_dir_check : bool, optional + Skip the remote directory check when the caller already knows + ``stage_path`` refers to a file (e.g. from a directory listing) + + Returns + ------- + bytes or str - ``local_path`` is None + None - ``local_path`` is a Path or str + + """ + if local_path is not None and not overwrite and os.path.exists(local_path): + raise OSError('target file already exists; use overwrite=True to replace') + if not _skip_dir_check and self.is_dir(stage_path): + raise IsADirectoryError(f'stage path is a directory: {stage_path}') + + out = self._manager._get( + self._fs_path(stage_path), + ).content + + if local_path is not None: + with open(local_path, 'wb') as outfile: + outfile.write(out) + return None + + if encoding: + return out.decode(encoding) + + return out + + def download_folder( + self, + stage_path: PathLike, + local_path: Optional[PathLike] = None, + *, + overwrite: bool = False, + ) -> None: + """ + Download a Stage folder to a local directory. + + The contents of ``stage_path`` are written into ``local_path``, + which is created as the destination folder. + + Parameters + ---------- + stage_path : Path or str + Path to the stage folder + local_path : Path or str, optional + Local directory to create and download into. Defaults to the + name of the ``stage_path`` folder in the current directory. + overwrite : bool, optional + Should an existing directory / files be overwritten if they exist? + + """ + # ``listdir`` returns paths relative to ``stage_path``, so the folder + # prefix has to be added back on before making any remote calls. + stage_prefix = normalize_remote_path(stage_path, strip_leading=True) + + if local_path is None: + local_path = os.path.basename(stage_prefix) + if not local_path: + raise ValueError( + 'local_path must be specified when downloading ' + 'the root folder', + ) + + if not overwrite and os.path.exists(local_path): + raise OSError( + 'target directory already exists; ' + 'use overwrite=True to replace', + ) + if not self.is_dir(stage_prefix): + raise NotADirectoryError(f'stage path is not a directory: {stage_path}') + + # Request objects so the file / directory type comes from the listing + # rather than an extra is_dir call per entry. + for entry in self.listdir(stage_prefix, recursive=True, return_objects=True): + rel_path = entry.path + target = ensure_within(local_path, os.path.join(local_path, rel_path)) + if entry.type == 'directory': + os.makedirs(target, exist_ok=True) + continue + remote_path = ( + f'{stage_prefix}/{rel_path}' if stage_prefix else rel_path + ) + os.makedirs(os.path.dirname(target) or '.', exist_ok=True) + self._download_file( + remote_path, target, + overwrite=overwrite, _skip_dir_check=True, + ) + + def remove(self, stage_path: PathLike) -> None: + """ + Delete a stage location. + + Parameters + ---------- + stage_path : Path or str + Path to the stage location + + """ + if self.is_dir(stage_path): + raise IsADirectoryError( + 'stage path is a directory, ' + f'use rmdir or removedirs: {stage_path}', + ) + + self._manager._delete(self._fs_path(stage_path)) + + def removedirs(self, stage_path: PathLike) -> None: + """ + Delete a stage folder recursively. + + Parameters + ---------- + stage_path : Path or str + Path to the stage location + + """ + stage_path = re.sub(r'/*$', r'', str(stage_path)) + '/' + self._manager._delete(self._fs_path(stage_path)) + + def rmdir(self, stage_path: PathLike) -> None: + """ + Delete a stage folder. + + Parameters + ---------- + stage_path : Path or str + Path to the stage location + + """ + stage_path = re.sub(r'/*$', r'', str(stage_path)) + '/' + + if self.listdir(stage_path): + raise OSError(f'stage folder is not empty, use removedirs: {stage_path}') + + self._manager._delete(self._fs_path(stage_path)) + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + +StageObject = FilesObject # alias for backward compatibility diff --git a/singlestoredb/management/timing.py b/singlestoredb/management/timing.py new file mode 100644 index 000000000..95b5bbc40 --- /dev/null +++ b/singlestoredb/management/timing.py @@ -0,0 +1,622 @@ +#!/usr/bin/env python +""" +Time accounting for the management API. + +Management calls are slow for two quite different reasons, and telling them +apart is the whole point of this module: an HTTP request that the server takes +its time answering, and a ``wait_on_*`` loop that sleeps between polls while a +deployment transitions. A stopwatch around ``create_cluster`` cannot separate +the two -- and it is the second that usually dominates -- so both are recorded +as events here. + +Every management HTTP request funnels through :meth:`Manager._doit`, and every +polling sleep through :func:`sleep`, so instrumenting those two covers all of +it. Recording is off unless something asks for it: + +.. code-block:: python + + from singlestoredb.management import timing + + with timing.trace() as t: + wm.create_cluster('my-cluster', size='S-00', wait_on_active=True) + + print(t.summary()) + +Set ``SINGLESTOREDB_MANAGEMENT_TRACE=1`` (the ``management.trace`` option) to +log every event to stderr as it finishes instead, which needs no code change. + +Traces are per-context: a trace opened on one thread does not see requests +issued on another. + +""" +import contextlib +import contextvars +import re +import sys +import threading +import time +from collections.abc import Iterator +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +from .. import config + + +#: Kinds of event. ``REQUEST`` is time spent in an HTTP call, ``WAIT`` is time +#: spent sleeping between polls of a resource that is still transitioning. +REQUEST = 'request' +WAIT = 'wait' + +#: Path segments that identify one particular resource rather than a route. +#: Collapsed to ``{id}`` so that 40 polls of one cluster aggregate into one +#: row instead of 40. UUIDs and integers cover every ID the API hands out. +_UUID_RE = re.compile( + r'^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}' + r'-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$', +) +_INT_RE = re.compile(r'^\d+$') + + +def route_of(method: str, path: str) -> str: + """ + Return the aggregation key for a request. + + Parameters + ---------- + method : str + HTTP method + path : str + Path of the resource, relative to the version root, as passed to + :meth:`Manager._get` and friends + + Returns + ------- + str + The method and path with resource IDs replaced by ``{id}``, e.g. + ``GET clusters/{id}`` + + """ + # Query strings are part of the path for some callers; they are noise here. + path = path.split('?')[0].strip('/') + parts = [ + '{id}' if _UUID_RE.match(x) or _INT_RE.match(x) else x + for x in path.split('/') + ] + return '{} {}'.format(method.upper(), '/'.join(parts)) + + +class Event: + """ + One timed operation. + + This object is not instantiated directly; :func:`record_request` and + :func:`sleep` create them. + + """ + + __slots__ = ( + 'kind', 'label', 'duration', 'started_at', 'status', + 'retries', 'request_bytes', 'response_bytes', 'error', + ) + + def __init__( + self, + kind: str, + label: str, + duration: float, + started_at: float, + status: Optional[int] = None, + retries: int = 0, + request_bytes: Optional[int] = None, + response_bytes: Optional[int] = None, + error: Optional[str] = None, + ): + #: Kind of event: REQUEST or WAIT + self.kind = kind + + #: Aggregation key: a route for a request, a reason for a wait + self.label = label + + #: Seconds the operation took + self.duration = duration + + #: Value of :func:`time.monotonic` when the operation started + self.started_at = started_at + + #: HTTP status code, if the request got a response + self.status = status + + #: Number of transport-level retries urllib3 made inside this request. + #: Non-zero here means the duration includes retry backoff. + self.retries = retries + + #: Size of the request body in bytes + self.request_bytes = request_bytes + + #: Size of the response body in bytes + self.response_bytes = response_bytes + + #: Exception type name, if the request never got a response + self.error = error + + def __str__(self) -> str: + out = f'{self.duration:7.3f}s {self.label}' + if self.error is not None: + out += f' -> {self.error}' + elif self.status is not None: + out += f' -> {self.status}' + if self.retries: + out += f' (retries={self.retries})' + return out + + def __repr__(self) -> str: + return f'' + + +class Stat: + """Aggregate of every :class:`Event` sharing a label.""" + + __slots__ = ('label', 'calls', 'total', 'min', 'max', 'retries', 'errors') + + def __init__(self, label: str): + #: The shared label + self.label = label + + #: Number of events + self.calls = 0 + + #: Total seconds across all of them + self.total = 0.0 + + #: Fastest and slowest of them, in seconds + self.min = 0.0 + self.max = 0.0 + + #: Total transport-level retries + self.retries = 0 + + #: Number of events that never got a response + self.errors = 0 + + @property + def mean(self) -> float: + """Mean seconds per event.""" + return self.total / self.calls if self.calls else 0.0 + + def add(self, event: 'Event') -> None: + """Fold an event into the aggregate.""" + self.min = event.duration if not self.calls else min(self.min, event.duration) + self.max = max(self.max, event.duration) + self.calls += 1 + self.total += event.duration + self.retries += event.retries + if event.error is not None: + self.errors += 1 + + def __str__(self) -> str: + return '{} calls={} total={:.3f}s mean={:.3f}s max={:.3f}s'.format( + self.label, self.calls, self.total, self.mean, self.max, + ) + + def __repr__(self) -> str: + return f'' + + +class Trace: + """ + Collector of management API timing events. + + Use :func:`trace` rather than instantiating this directly. + + """ + + def __init__(self) -> None: + #: Every event recorded, in completion order + self.events: List[Event] = [] + + self._lock = threading.Lock() + self._started_at: Optional[float] = None + self._stopped_at: Optional[float] = None + self._token: Optional[contextvars.Token[Tuple['Trace', ...]]] = None + + @classmethod + def of(cls, events: Any, elapsed: float) -> 'Trace': + """ + Return a stopped trace holding ``events`` and reporting ``elapsed``. + + For traces that are derived rather than collected -- a nested trace's + events subtracted from its parent's, say -- where the wall clock the + result stands for is not one this object measured. + + Parameters + ---------- + events : iterable of :class:`Event` + Events the trace should hold, in any order + elapsed : float + Seconds :attr:`elapsed` should report + + Returns + ------- + :class:`Trace` + + """ + out = cls() + out.events = sorted(events, key=lambda x: x.started_at) + out._started_at = 0.0 + out._stopped_at = max(0.0, elapsed) + return out + + @classmethod + def combine(cls, traces: Any) -> 'Trace': + """ + Return one trace holding every event from ``traces``. + + The result's :attr:`elapsed` is the sum of theirs, so it reads as the + time the traced sections covered between them rather than as wall clock + -- the sections need not have been contiguous. + + Parameters + ---------- + traces : iterable of :class:`Trace` + Traces to fold together + + Returns + ------- + :class:`Trace` + + """ + events: List[Event] = [] + elapsed = 0.0 + for one in traces: + events.extend(one.events) + elapsed += one.elapsed + return cls.of(events, elapsed) + + def start(self) -> 'Trace': + """Begin collecting events issued from this context.""" + self._started_at = time.monotonic() + self._stopped_at = None + self._token = _active.set(_active.get() + (self,)) + return self + + def stop(self) -> 'Trace': + """Stop collecting.""" + self._stopped_at = time.monotonic() + if self._token is not None: + _active.reset(self._token) + self._token = None + return self + + def add(self, event: Event) -> None: + """Record an event.""" + with self._lock: + self.events.append(event) + + @property + def elapsed(self) -> float: + """Wall clock seconds the trace covers.""" + if self._started_at is None: + return 0.0 + end = self._stopped_at if self._stopped_at is not None else time.monotonic() + return end - self._started_at + + def total(self, kind: Optional[str] = None) -> float: + """ + Return the seconds accounted for. + + Parameters + ---------- + kind : str, optional + Restrict to REQUEST or WAIT events. Defaults to all of them. + + Returns + ------- + float + + """ + return sum( + x.duration for x in self.events + if kind is None or x.kind == kind + ) + + @property + def unaccounted(self) -> float: + """ + Seconds spent neither in a request nor sleeping. + + This is the client's own work -- JSON parsing, object construction, and + whatever the caller did inside the trace. + + """ + return max(0.0, self.elapsed - self.total()) + + def stats(self, kind: Optional[str] = None) -> List[Stat]: + """ + Return per-label aggregates, slowest total first. + + Parameters + ---------- + kind : str, optional + Restrict to REQUEST or WAIT events. Defaults to all of them. + + Returns + ------- + List[:class:`Stat`] + + """ + out: Dict[str, Stat] = {} + for event in self.events: + if kind is not None and event.kind != kind: + continue + out.setdefault(event.label, Stat(event.label)).add(event) + return sorted(out.values(), key=lambda x: x.total, reverse=True) + + def summary(self) -> str: + """Return a human-readable report of where the time went.""" + elapsed = self.elapsed + lines = [f'Management API: {elapsed:.3f}s elapsed, {len(self.events)} events'] + + def share(seconds: float) -> str: + pct = 100.0 * seconds / elapsed if elapsed else 0.0 + return f'{seconds:9.3f}s {pct:5.1f}%' + + requests = self.total(REQUEST) + waits = self.total(WAIT) + lines.append( + ' requests {} {} calls'.format( + share(requests), sum(1 for x in self.events if x.kind == REQUEST), + ), + ) + lines.append( + ' waiting {} {} sleeps'.format( + share(waits), sum(1 for x in self.events if x.kind == WAIT), + ), + ) + lines.append(f' other {share(self.unaccounted)}') + + for kind, heading in ((REQUEST, 'route'), (WAIT, 'waiting on')): + stats = self.stats(kind) + if not stats: + continue + lines.append('') + lines.append( + ' {:<38} {:>5} {:>9} {:>9} {:>9}'.format( + heading, 'calls', 'total', 'mean', 'max', + ), + ) + for stat in stats: + row = ' {:<38} {:>5} {:>8.3f}s {:>8.3f}s {:>8.3f}s'.format( + stat.label[:38], stat.calls, stat.total, stat.mean, stat.max, + ) + if stat.retries: + row += f' retries={stat.retries}' + if stat.errors: + row += f' errors={stat.errors}' + lines.append(row) + + return '\n'.join(lines) + + def __str__(self) -> str: + return self.summary() + + def __repr__(self) -> str: + return ''.format( + len(self.events), self.elapsed, + ) + + +#: Traces collecting events in this context. A tuple so that nesting one trace +#: inside another feeds both, and so that resetting is a single assignment. +_active: contextvars.ContextVar[Tuple[Trace, ...]] = contextvars.ContextVar( + 'singlestoredb_management_traces', default=(), +) + + +@contextlib.contextmanager +def trace() -> Iterator[Trace]: + """ + Collect timing events for the duration of the block. + + Returns + ------- + :class:`Trace` + + """ + out = Trace().start() + try: + yield out + finally: + out.stop() + + +def logging_enabled() -> bool: + """Is every event being logged to stderr as it finishes?""" + return bool(config.get_option('management.trace')) + + +def _emit(event: Event) -> None: + """Hand an event to every trace collecting in this context.""" + for out in _active.get(): + out.add(event) + if logging_enabled(): + print(f'[singlestoredb.management] {event}', file=sys.stderr) + + +def recording() -> bool: + """ + Is anything recording? + + Checked before the bookkeeping in :func:`record_request` so that an + untraced call pays for one ``ContextVar.get`` and nothing else. + + """ + return bool(_active.get()) or logging_enabled() + + +def _response_size(res: Any) -> Optional[int]: + """Return the size of a response body in bytes, if it can be had cheaply.""" + length = res.headers.get('Content-Length') + if length is not None: + try: + return int(length) + except ValueError: + pass + try: + return len(res.content) + except Exception: + return None + + +def _retry_count(res: Any) -> int: + """ + Return the number of retries urllib3 made inside a request. + + Retries happen below ``requests``, so a request that took 30 seconds + because it was retried four times is otherwise indistinguishable from one + slow response. + + """ + try: + history = res.raw.retries.history + except Exception: + return 0 + return len(history or ()) + + +def record_request( + method: str, + path: str, + duration: float, + started_at: float, + response: Any = None, + error: Optional[BaseException] = None, +) -> None: + """ + Record one management HTTP request. + + Parameters + ---------- + method : str + HTTP method + path : str + Path of the resource, relative to the version root + duration : float + Seconds the request took + started_at : float + Value of :func:`time.monotonic` when the request started + response : requests.Response, optional + The response, if one arrived + error : Exception, optional + The transport failure, if none did + + """ + if not recording(): + return + request_bytes: Optional[int] = None + status: Optional[int] = None + retries = 0 + if response is not None: + status = response.status_code + retries = _retry_count(response) + body = getattr(response.request, 'body', None) + if body is not None: + request_bytes = len(body) + _emit( + Event( + REQUEST, route_of(method, path), duration, started_at, + status=status, retries=retries, request_bytes=request_bytes, + response_bytes=None if response is None else _response_size(response), + error=None if error is None else type(error).__name__, + ), + ) + + +@contextlib.contextmanager +def timed(label: str, kind: str = WAIT) -> Iterator[None]: + """ + Record the duration of a block of blocking work. + + For the parts of a wait that are not an HTTP request and not a sleep -- + a connection probe, say -- which would otherwise land in + :attr:`Trace.unaccounted` with no label on it. + + Parameters + ---------- + label : str + Aggregation key for the block + kind : str, optional + REQUEST or WAIT. Defaults to WAIT. + + """ + if not recording(): + yield + return + started_at = time.monotonic() + try: + yield + finally: + _emit(Event(kind, label, time.monotonic() - started_at, started_at)) + + +def now() -> float: + """Monotonic clock reading, for measuring what a poll iteration cost.""" + return time.monotonic() + + +def poll_cost(started_at: float, interval: float) -> float: + """ + Seconds to charge one poll iteration against its wait timeout. + + The ``wait_on_*`` loops used to charge every iteration a flat ``interval``, + which made ``wait_timeout`` a poll count rather than a duration: the + refetch between sleeps costs real time -- and, since the session gained + retries and a 180 second read timeout, can cost minutes of it -- that the + countdown never saw. A caller asking to wait 600 seconds could wait far + longer with no timeout raised. Charging the measured wall time instead + makes ``wait_timeout`` a genuine ceiling. + + The floor of ``interval`` keeps the loops bounded when :func:`time.sleep` + is patched out, as the offline tests do to poll without waiting. There the + measured time is ~0, and an iteration charged ~0 would never exhaust the + timeout. + + Parameters + ---------- + started_at : float + Reading from :func:`now` taken at the top of the iteration + interval : float + Nominal seconds between polls, used as the floor + + Returns + ------- + float + + """ + return max(interval, now() - started_at) + + +def sleep(seconds: float, label: str) -> None: + """ + Sleep between polls of a resource, recording the time as a wait. + + Every ``wait_on_*`` loop sleeps through here so that time spent waiting on + the server is reported separately from time spent talking to it. + + Parameters + ---------- + seconds : float + Seconds to sleep + label : str + What is being waited on, e.g. ``cluster state -> ACTIVE``. Used as the + aggregation key in :meth:`Trace.summary`. + + """ + if not recording(): + time.sleep(seconds) + return + started_at = time.monotonic() + time.sleep(seconds) + _emit(Event(WAIT, label, time.monotonic() - started_at, started_at)) diff --git a/singlestoredb/management/utils.py b/singlestoredb/management/utils.py index 5f1b5d522..bfdcc8658 100644 --- a/singlestoredb/management/utils.py +++ b/singlestoredb/management/utils.py @@ -1,7 +1,8 @@ #!/usr/bin/env python -"""SingleStoreDB Cluster Management.""" +"""Version-neutral helpers shared by the SingleStoreDB management API.""" import datetime import functools +import glob import itertools import os import re @@ -12,6 +13,7 @@ from typing import Dict from typing import List from typing import Optional +from typing import Set from typing import SupportsIndex from typing import Tuple from typing import TypeVar @@ -20,6 +22,7 @@ from .. import converters from ..config import get_option +from ..exceptions import ManagementError from ..utils import events JSON = Union[str, List[str], Dict[str, 'JSON']] @@ -35,36 +38,56 @@ PathLikeABC = os.PathLike[str] -class TTLProperty(object): - """Property with time limit.""" +class TTLProperty(property): + """ + Property with time limit. + + The value is cached on the instance, not on the descriptor. A descriptor is + shared by every instance of the class it is defined on, and what these + properties return is not: a manager's project list belongs to one + organization, so a descriptor-wide cache would hand one manager's list to a + manager holding a different token. + + Subclassing :class:`property` is what makes the decorated members read as + attributes rather than methods -- to :func:`isinstance` checks, and to + Sphinx, which documents anything else as a callable and so would tell + readers to write ``manager.projects()``. + """ def __init__(self, fget: Callable[[Any], Any], ttl: datetime.timedelta): - self.fget = fget + super().__init__(fget) + # ``property.fget`` is Optional to mypy and read-only at runtime, so the + # getter is kept here as well rather than narrowed at each call. + self._fget = fget self.ttl = ttl - self._last_executed = datetime.datetime(2000, 1, 1) - self._last_result = None self.__doc__ = fget.__doc__ self._name = '' - def reset(self) -> None: - self._last_executed = datetime.datetime(2000, 1, 1) - self._last_result = None - def __set_name__(self, owner: Any, name: str) -> None: self._name = name + @property + def _cache_key(self) -> str: + return f'_ttl_cache_{self._name or id(self)}' + + def reset(self, obj: Any) -> None: + """Discard the value cached for ``obj``, if any.""" + obj.__dict__.pop(self._cache_key, None) + def __get__(self, obj: Any, objtype: Any = None) -> Any: if obj is None: return self - if self._last_result is not None \ - and (datetime.datetime.now() - self._last_executed) < self.ttl: - return self._last_result + cached = obj.__dict__.get(self._cache_key) + if cached is not None: + value, fetched_at = cached + if (datetime.datetime.now() - fetched_at) < self.ttl: + return value - self._last_result = self.fget(obj) - self._last_executed = datetime.datetime.now() + value = self._fget(obj) + obj.__dict__[self._cache_key] = (value, datetime.datetime.now()) - return self._last_result + return value def ttl_property(ttl: datetime.timedelta) -> Callable[[Any], Any]: @@ -224,15 +247,49 @@ def get_token() -> Optional[str]: def get_cluster_id() -> Optional[str]: - """Return the cluster id for the current token or environment.""" - return os.environ.get('SINGLESTOREDB_CLUSTER') or None + """ + Return the cluster id for the current token or environment. + + The v2 spelling of :func:`get_workspace_id`, and the same value: there is + no ``SINGLESTOREDB_CLUSTER``, because the notebook environment publishes + the current deployment as ``SINGLESTOREDB_WORKSPACE`` whatever the API + version calls it. + """ + return get_workspace_id() def get_workspace_id() -> Optional[str]: - """Return the workspace id for the current token or environment.""" + """ + Return the deployment id for the current token or environment. + + ``SINGLESTOREDB_WORKSPACE`` is the notebook environment's name for the + current deployment at every API version: the workspace ID at v1, and the + cluster ID from v2 onward. + """ return os.environ.get('SINGLESTOREDB_WORKSPACE') or None +def get_project_id() -> Optional[str]: + """ + Return the inference API project id for the current environment. + + ``SINGLESTOREDB_PROJECT`` is *not* a project of the cluster management API, + despite the name. The notebook environment sets both, and they disagree: a + notebook attached to a cluster in one management project reports an + unrelated ID here, one that draws ``404 project not found`` from + ``GET /v2/projects/{id}``. It names a project of the inference API, which is + a separate service with its own namespace, and + :class:`singlestoredb.management.inference_api.InferenceAPIManager` is its + only legitimate consumer. + + To pick the management project a new deployment belongs in, use + :meth:`singlestoredb.management.v2.cluster.ClusterManager. + _resolve_project_id`, which reads the project off the current deployment + instead. + """ + return os.environ.get('SINGLESTOREDB_PROJECT') or None + + def get_virtual_workspace_id() -> Optional[str]: """Return the virtual workspace id for the current token or environment.""" return os.environ.get('SINGLESTOREDB_VIRTUAL_WORKSPACE') or None @@ -243,6 +300,102 @@ def get_database_name() -> Optional[str]: return os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE') or None +def normalize_remote_path(path: PathLike, *, strip_leading: bool = False) -> str: + """Normalize a caller-supplied remote path to POSIX form. + + Remote FileSpace / Stage paths always use ``/``. Callers may build a + path with :func:`os.path.join`, which uses the local separator, so + backslashes are converted to ``/`` before the path is used. Duplicate + separators are collapsed and the trailing separator is removed, so the + result can safely be concatenated with ``'/' + rel``. + + Parameters + ---------- + path : Path or str + Remote path to normalize + strip_leading : bool, optional + Also remove leading ``./`` and ``/`` segments, making the path + relative to the remote root + + Returns + ------- + str + + """ + out = str(path).replace('\\', '/') + if strip_leading: + out = re.sub(r'^(\./|/)+', r'', out) + out = re.sub(r'/{2,}', r'/', out) + return re.sub(r'/+$', r'', out) + + +def ensure_within(local_root: PathLike, target: PathLike) -> str: + """Verify ``target`` resolves inside ``local_root``. + + Returns the normalized (but unresolved) path on success. The + containment check uses :func:`os.path.realpath` so symlink trickery + can't escape ``local_root``. Raises :class:`ManagementError` if + ``target`` would escape ``local_root``, e.g. via ``..`` segments + coming from an untrusted remote listing. + """ + target_str = os.fspath(target) + normalized = os.path.normpath(target_str) + base = os.path.realpath(os.fspath(local_root)) + resolved = os.path.realpath(target_str) + if resolved != base and not resolved.startswith(base + os.sep): + raise ManagementError( + msg=f'Refusing to write outside destination: {target_str}', + ) + return normalized + + +def resolve_ignore_files( + local_root: PathLike, + ignore: Optional[Union[PathLike, List[PathLike]]], +) -> Set[str]: + """Expand ``ignore`` glob patterns into a set of local file paths. + + Relative patterns are resolved against ``local_root`` rather than the + process working directory, so patterns like ``**/*.pyc`` match the tree + actually being uploaded. Absolute patterns are used as given. Results are + normalized with :func:`os.path.normpath`, so callers must normalize the + paths they test for membership too — a raw ``os.walk`` result such as + ``./a.pyc`` will not compare equal to the normalized ``a.pyc``. + + Parameters + ---------- + local_root : Path or str + Local directory the patterns are relative to + ignore : Path or str or List[Path] or List[str], optional + Glob pattern(s) of files to ignore + + Returns + ------- + Set[str] + + """ + out: Set[str] = set() + + if not ignore: + return out + + root = os.path.normpath(os.fspath(local_root)) + patterns = ignore if isinstance(ignore, list) else [ignore] + + for item in patterns: + pattern = os.fspath(item) + if not os.path.isabs(pattern): + pattern = os.path.join(root, pattern) + # Always recursive so '**' works regardless of the caller's + # recursion setting. + out.update( + os.path.normpath(x) + for x in glob.glob(pattern, recursive=True) + ) + + return out + + def enable_http_tracing() -> None: """Enable tracing of HTTP requests.""" import logging @@ -319,10 +472,27 @@ def from_datetime( return out -def vars_to_str(obj: Any) -> str: - """Render a string representation of vars(obj).""" +def vars_to_str(obj: Any, extra: Optional[Dict[str, Any]] = None) -> str: + """ + Render a string representation of vars(obj). + + Parameters + ---------- + obj : Any + The object to render. Attributes whose name starts with ``_``, and + those with a falsy value, are left out. + extra : dict, optional + Attributes to report that ``vars(obj)`` does not hold. This is for a + lazily resolved property, whose value must not be fetched merely to + print the object: the owner passes what it already has, which is either + the resolved value or the ID it would resolve. Reported and sorted like + any other attribute, and left out on a falsy value the same way. + + """ attrs = [] - obj_vars = vars(obj) + obj_vars = dict(vars(obj)) + if extra: + obj_vars.update(extra) if 'name' in obj_vars: attrs.append('name={}'.format(repr(obj_vars['name']))) if 'id' in obj_vars: diff --git a/singlestoredb/management/v1/__init__.py b/singlestoredb/management/v1/__init__.py new file mode 100644 index 000000000..f35ec5a2f --- /dev/null +++ b/singlestoredb/management/v1/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python +""" +SingleStoreDB Management API v1 -- **deprecated**. + +.. deprecated:: + Use :mod:`singlestoredb.management.v2`, which the ``management.version`` + option (the ``SINGLESTOREDB_MANAGEMENT_VERSION`` environment variable) names + by default. This whole package is scheduled for removal; the public entry + points warn when they resolve to it. Drop ``version='v1'`` and leave the + option unset. + + The one exception is :mod:`.inference_api`, which has no replacement yet -- + see that module. +""" +# The version-neutral helpers in singlestoredb.management look these up here by +# name. A deployment is a workspace group, so they live in the workspace +# module. +from .workspace import get_organization as get_organization +from .workspace import get_secret as get_secret +from .workspace import get_stage as get_stage diff --git a/singlestoredb/management/v1/billing_usage.py b/singlestoredb/management/v1/billing_usage.py new file mode 100644 index 000000000..3a0c14580 --- /dev/null +++ b/singlestoredb/management/v1/billing_usage.py @@ -0,0 +1,14 @@ +#!/usr/bin/env python +""" +SingleStoreDB Billing Usage API v1 -- **deprecated**. + +.. deprecated:: + Only the module path is deprecated; the names re-exported below are the + shared implementations. Import them from + :mod:`singlestoredb.management.billing_usage`. + +``GET /v1/billing/usage`` is implemented in +:mod:`singlestoredb.management.billing_usage`, so this module only re-exports it. +""" +from ..billing_usage import BillingUsageItem as BillingUsageItem +from ..billing_usage import UsageItem as UsageItem diff --git a/singlestoredb/management/v1/export.py b/singlestoredb/management/v1/export.py new file mode 100644 index 000000000..9ab39e161 --- /dev/null +++ b/singlestoredb/management/v1/export.py @@ -0,0 +1,318 @@ +#!/usr/bin/env python +""" +SingleStoreDB export service (management API v1) -- **deprecated**. + +.. deprecated:: + Use :mod:`singlestoredb.management.export`, where an export is owned by a + :class:`~singlestoredb.management.cluster.Cluster`, and the ``CLUSTER``-based + Fusion ``EXPORT`` grammar. +""" +from __future__ import annotations + +import copy +import json +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +from ... import ManagementError +from ..utils import vars_to_str +from .workspace import WorkspaceGroup +from .workspace import WorkspaceManager + + +class ExportService(object): + """ + Export service (API v1). + + .. deprecated:: + Use :class:`singlestoredb.management.export.ExportService`, which takes a + :class:`Cluster`. + """ + + database: str + table: str + catalog_info: Dict[str, Any] + storage_info: Dict[str, Any] + columns: Optional[List[str]] + partition_by: Optional[List[Dict[str, str]]] + order_by: Optional[List[Dict[str, Dict[str, str]]]] + properties: Optional[Dict[str, Any]] + incremental: bool + refresh_interval: Optional[int] + export_id: Optional[str] + + def __init__( + self, + workspace_group: WorkspaceGroup, + database: str, + table: str, + catalog_info: Union[str, Dict[str, Any]], + storage_info: Union[str, Dict[str, Any]], + columns: Optional[List[str]] = None, + partition_by: Optional[List[Dict[str, str]]] = None, + order_by: Optional[List[Dict[str, Dict[str, str]]]] = None, + incremental: bool = False, + refresh_interval: Optional[int] = None, + properties: Optional[Dict[str, Any]] = None, + ): + #: Workspace group + self.workspace_group = workspace_group + + #: Name of SingleStoreDB database + self.database = database + + #: Name of SingleStoreDB table + self.table = table + + #: List of columns to export + self.columns = columns + + #: Catalog + if isinstance(catalog_info, str): + self.catalog_info = json.loads(catalog_info) + else: + self.catalog_info = copy.copy(catalog_info) + + #: Storage + if isinstance(storage_info, str): + self.storage_info = json.loads(storage_info) + else: + self.storage_info = copy.copy(storage_info) + + self.partition_by = partition_by or None + self.order_by = order_by or None + self.properties = properties or None + + self.incremental = incremental + self.refresh_interval = refresh_interval + + self.export_id = None + + self._manager: Optional[WorkspaceManager] = workspace_group._manager + + @classmethod + def from_export_id( + cls, + workspace_group: WorkspaceGroup, + export_id: str, + ) -> ExportService: + """Create export service from export ID.""" + out = cls( + workspace_group=workspace_group, + database='', + table='', + catalog_info={}, + storage_info={}, + ) + out.export_id = export_id + return out + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + def create_cluster_identity(self) -> Dict[str, Any]: + """Create a cluster identity.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + out = self._manager._post( + f'workspaceGroups/{self.workspace_group.id}/' + 'egress/createEgressClusterIdentity', + json=dict( + catalogInfo=self.catalog_info, + storageInfo=self.storage_info, + ), + ) + + return out.json() + + def start(self, tags: Optional[List[str]] = None) -> 'ExportStatus': + """Start the export process.""" + if not self.table or not self.database: + raise ManagementError( + msg='Database and table must be set before starting the export.', + ) + + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + partition_spec = None + if self.partition_by: + partition_spec = dict(partitions=self.partition_by) + + sort_order_spec = None + if self.order_by: + sort_order_spec = dict(keys=self.order_by) + + out = self._manager._post( + f'workspaceGroups/{self.workspace_group.id}/egress/startTableEgress', + json={ + k: v for k, v in dict( + databaseName=self.database, + tableName=self.table, + storageInfo=self.storage_info, + catalogInfo=self.catalog_info, + partitionSpec=partition_spec, + sortOrderSpec=sort_order_spec, + properties=self.properties, + incremental=self.incremental or None, + refreshInterval=self.refresh_interval + if self.refresh_interval is not None else None, + ).items() if v is not None + }, + ) + + self.export_id = str(out.json()['egressID']) + + return ExportStatus(self.export_id, self.workspace_group) + + def suspend(self) -> 'ExportStatus': + """Suspend the export process.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + if self.export_id is None: + raise ManagementError( + msg='Export ID is not set. You must start the export first.', + ) + + self._manager._post( + f'workspaceGroups/{self.workspace_group.id}/egress/suspendTableEgress', + json=dict(egressID=self.export_id), + ) + + return ExportStatus(self.export_id, self.workspace_group) + + def resume(self) -> 'ExportStatus': + """Resume the export process.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + if self.export_id is None: + raise ManagementError( + msg='Export ID is not set. You must start the export first.', + ) + + self._manager._post( + f'workspaceGroups/{self.workspace_group.id}/egress/resumeTableEgress', + json=dict(egressID=self.export_id), + ) + + return ExportStatus(self.export_id, self.workspace_group) + + def drop(self) -> None: + """Drop the export process.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + if self.export_id is None: + raise ManagementError( + msg='Export ID is not set. You must start the export first.', + ) + + self._manager._delete( + f'workspaceGroups/{self.workspace_group.id}/egress/dropTableEgress', + json=dict(egressID=self.export_id), + ) + + return None + + def status(self) -> ExportStatus: + """Get the status of the export process.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + if self.export_id is None: + raise ManagementError( + msg='Export ID is not set. You must start the export first.', + ) + + return ExportStatus(self.export_id, self.workspace_group) + + +class ExportStatus(object): + """ + Status of an export. + + .. deprecated:: + Use :class:`singlestoredb.management.export.ExportStatus`, which is keyed + by a :class:`Cluster`. + """ + + export_id: str + + def __init__(self, export_id: str, workspace_group: WorkspaceGroup): + self.export_id = export_id + self.workspace_group = workspace_group + self._manager: Optional[WorkspaceManager] = workspace_group._manager + + def _info(self) -> Dict[str, Any]: + """Return export status.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + out = self._manager._get( + f'workspaceGroups/{self.workspace_group.id}/egress/tableEgressStatus', + json=dict(egressID=self.export_id), + ) + + return out.json() + + @property + def status(self) -> str: + """Return export status.""" + return self._info().get('status', 'Unknown') + + @property + def message(self) -> str: + """Return export status message.""" + return self._info().get('statusMsg', '') + + def __str__(self) -> str: + return self.status + + def __repr__(self) -> str: + return self.status + + +def _get_exports( + workspace_group: WorkspaceGroup, + scope: str = 'all', +) -> List[ExportStatus]: + """Get all exports in the workspace group.""" + if workspace_group._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + out = workspace_group._manager._get( + f'workspaceGroups/{workspace_group.id}/egress/tableEgressStatus', + json=dict(scope=scope), + ) + + return [ + ExportStatus(item['egressID'], workspace_group) + for item in out.json() + ] diff --git a/singlestoredb/management/v1/files.py b/singlestoredb/management/v1/files.py new file mode 100644 index 000000000..074ffe720 --- /dev/null +++ b/singlestoredb/management/v1/files.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python +""" +SingleStoreDB Files Management API v1 -- **deprecated**. + +.. deprecated:: + Only the module path is deprecated; the names re-exported below are the + shared implementations. Import them from + :mod:`singlestoredb.management.files`, and call ``manage_files()`` without + ``version='v1'``. + +The ``files/fs/{space}/...`` routes are implemented in +:mod:`singlestoredb.management.files`, and this module re-exports them so that +``manage_files(version='v1')`` can resolve the implementation by name. +""" +from ..files import FileLocation as FileLocation +from ..files import FilesManager as FilesManager +from ..files import FilesObject as FilesObject +from ..files import FilesObjectBytesReader as FilesObjectBytesReader +from ..files import FilesObjectBytesWriter as FilesObjectBytesWriter +from ..files import FilesObjectTextReader as FilesObjectTextReader +from ..files import FilesObjectTextWriter as FilesObjectTextWriter +from ..files import FileSpace as FileSpace +from ..files import MODELS_SPACE as MODELS_SPACE +from ..files import PERSONAL_SPACE as PERSONAL_SPACE +from ..files import SHARED_SPACE as SHARED_SPACE diff --git a/singlestoredb/management/inference_api.py b/singlestoredb/management/v1/inference_api.py similarity index 93% rename from singlestoredb/management/inference_api.py rename to singlestoredb/management/v1/inference_api.py index be9568ff9..e44d5e93a 100644 --- a/singlestoredb/management/inference_api.py +++ b/singlestoredb/management/v1/inference_api.py @@ -1,14 +1,23 @@ #!/usr/bin/env python -"""SingleStoreDB Cloud Inference API.""" +""" +SingleStoreDB Cloud Inference API (v1 only). + +Deliberately **not** deprecated, unlike the rest of +:mod:`singlestoredb.management.v1`: the ``inference/*`` routes are served here +and nowhere else. Reached through :attr:`Organization.inference_apis`, which is +why ``fusion/handlers/utils.get_inference_api_manager`` and the +:mod:`singlestoredb.ai` helpers resolve this package directly rather than +following the ``management.version`` option, and do so without warning. +""" import os from typing import Any from typing import Dict from typing import List from typing import Optional -from .utils import vars_to_str -from singlestoredb.exceptions import ManagementError -from singlestoredb.management.manager import Manager +from ...exceptions import ManagementError +from ..manager import Manager +from ..utils import vars_to_str class ModelOperationResult(object): @@ -140,7 +149,7 @@ def __repr__(self) -> str: return str(self) -class InferenceAPIInfo(object): +class InferenceAPIInfo: """ Inference API definition. @@ -192,7 +201,7 @@ def from_dict( Returns ------- - :class:`Job` + :class:`InferenceAPIInfo` """ out = cls( @@ -257,7 +266,7 @@ def drop(self) -> ModelOperationResult: return self._manager.drop(self.name) -class InferenceAPIManager(object): +class InferenceAPIManager: """ SingleStoreDB Inference APIs manager. diff --git a/singlestoredb/management/v1/job.py b/singlestoredb/management/v1/job.py new file mode 100644 index 000000000..678d21ab6 --- /dev/null +++ b/singlestoredb/management/v1/job.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python +""" +SingleStoreDB Job Management API v1 -- **deprecated**. + +.. deprecated:: + Use :mod:`singlestoredb.management.job`, reached from + :attr:`Organization.jobs`. +""" +from ..job import Execution as Execution +from ..job import ExecutionConfig as ExecutionConfig +from ..job import ExecutionMetadata as ExecutionMetadata +from ..job import ExecutionsData as ExecutionsData +from ..job import Job as Job +from ..job import JobMetadata as JobMetadata +from ..job import JobsManager as _JobsManager +from ..job import Mode as Mode +from ..job import Parameter as Parameter +from ..job import Runtime as Runtime +from ..job import Schedule as Schedule +from ..job import Status as Status +from ..job import TargetConfig as TargetConfig +from ..job import TargetType as TargetType + + +class JobsManager(_JobsManager): + """ + SingleStoreDB scheduled notebook jobs manager (API v1). + + .. deprecated:: + Use :class:`singlestoredb.management.job.JobsManager`. + + The ``jobs`` routes are the shared ones; what is specific here is the + ``targetConfig.targetType`` vocabulary this manager schedules with, + ``'Workspace'`` and ``'VirtualWorkspace'``. + + ``'Cluster'`` is a third value these routes accept, naming a legacy + self-managed cluster. Only the read path ever sees it -- the deployment a + job is scheduled against comes from ``SINGLESTOREDB_WORKSPACE``, which + never names a legacy cluster. + """ + + _deployment_target_type = TargetType.WORKSPACE + _starter_target_type = TargetType.VIRTUAL_WORKSPACE diff --git a/singlestoredb/management/v1/organization.py b/singlestoredb/management/v1/organization.py new file mode 100644 index 000000000..78c79a126 --- /dev/null +++ b/singlestoredb/management/v1/organization.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python +""" +SingleStoreDB Organization API v1 -- **deprecated**. + +.. deprecated:: + Use :mod:`singlestoredb.management.organization`, reached from + :func:`singlestoredb.management.get_organization`. +""" +from ..organization import Organization as _Organization +from ..organization import Organizations as _Organizations +from ..organization import Secret as Secret +from .inference_api import InferenceAPIManager +from .job import JobsManager + + +class Organization(_Organization): + """ + Organization in SingleStoreDB Cloud portal (API v1). + + .. deprecated:: + Use :class:`singlestoredb.management.organization.Organization`. + + ``organizations/current`` and ``secrets`` are the shared routes; all this + subclass does is hand out the sub-managers that belong to this API version, + which is what keeps job schedules on the ``targetType`` vocabulary these + routes expect. + """ + + _jobs_manager_class = JobsManager + _inference_api_manager_class = InferenceAPIManager + + +class Organizations(_Organizations): + """ + Organizations (API v1). + + .. deprecated:: + Use :class:`singlestoredb.management.organization.Organizations`. + """ + + _organization_class = Organization + + @property + def current(self) -> Organization: + """Get current organization.""" + out = super().current + assert isinstance(out, Organization) + return out diff --git a/singlestoredb/management/v1/region.py b/singlestoredb/management/v1/region.py new file mode 100644 index 000000000..e0d2b40ba --- /dev/null +++ b/singlestoredb/management/v1/region.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python +""" +SingleStoreDB Region Management API v1 -- **deprecated**. + +.. deprecated:: + Only the module path is deprecated; the names re-exported below are the + shared implementations. Import them from + :mod:`singlestoredb.management.region`, and call ``manage_regions()`` + without ``version='v1'``. + +``GET /v1/regions`` and ``GET /v1/regions/sharedtier`` behave exactly as +:mod:`singlestoredb.management.region` implements them, so this module only +re-exports it. These routes report a ``regionID``, which +:meth:`Region.from_dict` treats as optional. +""" +from ..region import Region as Region +from ..region import RegionManager as RegionManager diff --git a/singlestoredb/management/v1/stage.py b/singlestoredb/management/v1/stage.py new file mode 100644 index 000000000..30d4df100 --- /dev/null +++ b/singlestoredb/management/v1/stage.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python +""" +SingleStoreDB Stage Management API v1 -- **deprecated**. + +.. deprecated:: + Use :mod:`singlestoredb.management.stage`. +""" +from ..stage import Stage as _Stage +from ..utils import PathLike + + +class Stage(_Stage): + """ + Stage file space for a workspace group or starter workspace. + + .. deprecated:: + Use :class:`singlestoredb.management.stage.Stage`, reached from + :attr:`Cluster.stage` or :func:`singlestoredb.management.get_stage`. + + Here Stage is a top-level resource keyed by deployment ID: + ``stage/{id}/fs/``, which is why this subclass overrides the path. + """ + + def _fs_path(self, path: PathLike = '') -> str: + return f'stage/{self._deployment_id}/fs/{path}' diff --git a/singlestoredb/management/v1/workspace.py b/singlestoredb/management/v1/workspace.py new file mode 100644 index 000000000..718b292f9 --- /dev/null +++ b/singlestoredb/management/v1/workspace.py @@ -0,0 +1,1559 @@ +#!/usr/bin/env python +""" +SingleStoreDB Workspace Management (management API v1) -- **deprecated**. + +.. deprecated:: + Use :mod:`singlestoredb.management.cluster`, where a deployment is the flat + :class:`~singlestoredb.management.cluster.Cluster` rather than a workspace + group containing workspaces. The replacements are: + + ============================================ ==================== + Deprecated Use instead + ============================================ ==================== + :func:`singlestoredb.manage_workspaces` ``manage_clusters`` + :class:`WorkspaceManager` ``ClusterManager`` + :class:`WorkspaceGroup` + :class:`Workspace` ``Cluster`` + :class:`StarterWorkspace` ``StarterCluster`` + :func:`get_workspace_group` ``get_cluster`` + :func:`get_workspace` ``get_cluster`` + :func:`get_stage` ``get_stage`` + ============================================ ==================== + + :func:`get_workspace_group` and :func:`get_workspace` both collapse onto + ``get_cluster``, which reads ``SINGLESTOREDB_WORKSPACE`` alone and finds a + cluster ID in it, where the two functions here read + ``SINGLESTOREDB_WORKSPACE_GROUP`` and ``SINGLESTOREDB_WORKSPACE`` + respectively. ``get_stage`` keeps its name and moves to the + version-neutral :mod:`singlestoredb.management`. +""" +from __future__ import annotations + +import datetime +import io +import os +import re +from typing import Any +from typing import cast +from typing import Dict +from typing import List +from typing import Literal +from typing import Optional +from typing import overload +from typing import Union + +from .. import timing +from ... import config +from ... import connection +from ...exceptions import ManagementError +from ..billing import Billing as Billing +from ..manager import Manager +from ..region import Region +from ..stage import StageObject as StageObject +from ..utils import camel_to_snake_dict +from ..utils import ensure_within +from ..utils import from_datetime +from ..utils import NamedList +from ..utils import normalize_remote_path +from ..utils import PathLike +from ..utils import resolve_ignore_files +from ..utils import snake_to_camel +from ..utils import snake_to_camel_dict +from ..utils import to_datetime +from ..utils import ttl_property +from ..utils import vars_to_str +from .organization import Organization +from .organization import Organizations as Organizations +from .stage import Stage as Stage + +#: Base management API path for the shared-tier resource. +SHAREDTIER_PATH = 'sharedtier/virtualWorkspaces' + + +def get_organization() -> Organization: + """ + Get the organization. + + .. deprecated:: + Call :func:`singlestoredb.management.get_organization`, which + dispatches on the ``management.version`` option. + """ + from ..workspace import _manage_workspaces_v1 + return _manage_workspaces_v1().organization + + +def get_secret(name: str) -> Optional[str]: + """ + Get a secret from the organization. + + .. deprecated:: + Call :func:`singlestoredb.management.get_secret`, which dispatches on + the ``management.version`` option. + """ + return get_organization().get_secret(name).value + + +def get_workspace_group( + workspace_group: Optional[Union[WorkspaceGroup, str]] = None, +) -> WorkspaceGroup: + """ + Get the workspace group. + + .. deprecated:: + Use :func:`singlestoredb.management.cluster.get_cluster`. + + Falls back to ``SINGLESTOREDB_WORKSPACE_GROUP``, the notebook environment's + group ID. :func:`singlestoredb.management.cluster.get_cluster` ignores that + variable and reads ``SINGLESTOREDB_WORKSPACE`` instead. + """ + from ..workspace import _manage_workspaces_v1 + if isinstance(workspace_group, WorkspaceGroup): + return workspace_group + elif workspace_group: + return _manage_workspaces_v1().workspace_groups[workspace_group] + elif 'SINGLESTOREDB_WORKSPACE_GROUP' in os.environ: + return _manage_workspaces_v1().workspace_groups[ + os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] + ] + raise RuntimeError('no workspace group specified') + + +def get_stage( + workspace_group: Optional[Union[WorkspaceGroup, str]] = None, +) -> Stage: + """ + Get the stage for the workspace group. + + .. deprecated:: + Call :func:`singlestoredb.management.get_stage`, which dispatches on + the ``management.version`` option. + """ + return get_workspace_group(workspace_group).stage + + +def get_workspace( + workspace_group: Optional[Union[WorkspaceGroup, str]] = None, + workspace: Optional[Union[Workspace, str]] = None, +) -> Workspace: + """ + Get a workspace within a workspace group. + + .. deprecated:: + Use :func:`singlestoredb.management.cluster.get_cluster`, which reads + the same ``SINGLESTOREDB_WORKSPACE`` variable but finds a cluster ID in + it. + + Falls back to ``SINGLESTOREDB_WORKSPACE``, the notebook environment's name + for the current deployment, whose value here is a workspace ID. + """ + if isinstance(workspace, Workspace): + return workspace + wg = get_workspace_group(workspace_group) + if workspace: + return wg.workspaces[workspace] + elif 'SINGLESTOREDB_WORKSPACE' in os.environ: + return wg.workspaces[ + os.environ['SINGLESTOREDB_WORKSPACE'] + ] + raise RuntimeError('no workspace group specified') + + +class Workspace: + """ + SingleStoreDB workspace definition. + + .. deprecated:: + Use :class:`singlestoredb.management.cluster.Cluster`, which is flat: + it carries the size and state this class holds together with the region + and Stage that :class:`WorkspaceGroup` held, so there is no separate + group object to look it up through. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`WorkspaceManager`. Workspaces are created using + :meth:`WorkspaceManager.create_workspace`, or existing workspaces are + accessed by either :attr:`WorkspaceManager.workspaces` or by calling + :meth:`WorkspaceManager.get_workspace`. + + See Also + -------- + :meth:`WorkspaceManager.create_workspace` + :meth:`WorkspaceManager.get_workspace` + :attr:`WorkspaceManager.workspaces` + + """ + + name: str + id: str + group_id: str + size: str + state: str + created_at: Optional[datetime.datetime] + terminated_at: Optional[datetime.datetime] + endpoint: Optional[str] + auto_suspend: Optional[Dict[str, Any]] + cache_config: Optional[float] + deployment_type: Optional[str] + 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[float] + + def __init__( + self, + name: str, + workspace_id: str, + workspace_group: Union[str, 'WorkspaceGroup'], + size: str, + state: str, + created_at: Union[str, datetime.datetime], + terminated_at: Optional[Union[str, datetime.datetime]] = None, + endpoint: Optional[str] = None, + auto_suspend: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + deployment_type: Optional[str] = None, + 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[float] = None, + ): + #: Name of the workspace + self.name = name + + #: Unique ID of the workspace + self.id = workspace_id + + #: Unique ID of the workspace group + if isinstance(workspace_group, WorkspaceGroup): + self.group_id = workspace_group.id + else: + self.group_id = workspace_group + + #: Size of the workspace in workspace size notation (S-00, S-1, etc.) + self.size = size + + #: State of the workspace: PendingCreation, Transitioning, Active, + #: Terminated, Suspended, Resuming, Failed + self.state = state.strip() + + #: Timestamp of when the workspace was created + self.created_at = to_datetime(created_at) + + #: Timestamp of when the workspace was terminated + self.terminated_at = to_datetime(terminated_at) + + #: Hostname (or IP address) of the workspace database server + self.endpoint = endpoint + + #: Current auto-suspend settings + self.auto_suspend = camel_to_snake_dict(auto_suspend) + + #: Multiplier for the persistent cache + self.cache_config = cache_config + + #: Deployment type of the workspace + self.deployment_type = deployment_type + + #: Database attachments + self.resume_attachments = [ + camel_to_snake_dict(x) # type: ignore + for x in resume_attachments or [] + if x is not None + ] + + #: Current progress percentage for scaling the workspace + self.scaling_progress = scaling_progress + + #: Timestamp when workspace was last resumed + self.last_resumed_at = to_datetime(last_resumed_at) + + #: Auto-scale settings for the workspace + self.auto_scale = camel_to_snake_dict(auto_scale) + + #: Whether Kai is enabled on this workspace + self.kai_enabled = kai_enabled + + #: Current scale factor for the workspace + self.scale_factor = scale_factor + + self._manager: Optional[WorkspaceManager] = 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: 'WorkspaceManager') -> 'Workspace': + """ + Construct a Workspace from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : WorkspaceManager, optional + The WorkspaceManager the Workspace belongs to + + Returns + ------- + :class:`Workspace` + + """ + out = cls( + name=obj['name'], + workspace_id=obj['workspaceID'], + workspace_group=obj['workspaceGroupID'], + size=obj.get('size', 'Unknown'), + state=obj['state'], + created_at=obj['createdAt'], + terminated_at=obj.get('terminatedAt'), + endpoint=obj.get('endpoint'), + auto_suspend=obj.get('autoSuspend'), + cache_config=obj.get('cacheConfig'), + deployment_type=obj.get('deploymentType'), + 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 + + def update( + self, + auto_suspend: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + deployment_type: Optional[str] = None, + size: Optional[str] = None, + auto_scale: Optional[Dict[str, Any]] = None, + enable_kai: Optional[bool] = None, + scale_factor: Optional[float] = None, + ) -> None: + """ + Update the workspace definition. + + Parameters + ---------- + auto_suspend : Dict[str, Any], optional + Auto-suspend mode for the workspace: IDLE, SCHEDULED, DISABLED + cache_config : float, optional + Specifies the multiplier for the persistent cache associated + with the workspace. If specified, it enables the cache configuration + multiplier. It can have one of the following values: 1, 2, or 4. + deployment_type : str, optional + The deployment type that will be applied to all the workspaces + within the group + size : str, optional + Size of the workspace (in workspace size notation), such as "S-1". + auto_scale : Dict[str, Any], optional + Auto-scale settings for the workspace. + enable_kai : bool, optional + Whether to enable SingleStore Kai on this workspace. + scale_factor : float, optional + Scale factor for the workspace. + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + data = { + k: v for k, v in dict( + autoSuspend=snake_to_camel_dict(auto_suspend), + cacheConfig=cache_config, + deploymentType=deployment_type, + size=size, + autoScale=snake_to_camel_dict(auto_scale), + enableKai=enable_kai, + scaleFactor=scale_factor, + ).items() if v is not None + } + self._manager._patch(f'workspaces/{self.id}', json=data) + self.refresh() + + def refresh(self) -> Workspace: + """Update the object to the current state.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + new_obj = self._manager.get_workspace(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + def terminate( + self, + wait_on_terminated: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + force: bool = False, + ) -> None: + """ + Terminate the workspace. + + Parameters + ---------- + wait_on_terminated : bool, optional + Wait for the workspace to go into 'Terminated' mode before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + force : bool, optional + Should the workspace group be terminated even if it has workspaces? + + Raises + ------ + ManagementError + If timeout is reached + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + force_str = 'true' if force else 'false' + self._manager._delete(f'workspaces/{self.id}?force={force_str}') + if wait_on_terminated: + self._manager._wait_on_state( + self._manager.get_workspace(self.id), + 'Terminated', interval=wait_interval, timeout=wait_timeout, + ) + self.refresh() + + def connect(self, **kwargs: Any) -> connection.Connection: + """ + Create a connection to the database server for this workspace. + + Parameters + ---------- + **kwargs : keyword-arguments, optional + Parameters to the SingleStoreDB `connect` function except host + and port which are supplied by the workspace object + + Returns + ------- + :class:`Connection` + + """ + if not self.endpoint: + raise ManagementError( + msg='An endpoint has not been set in this workspace configuration', + ) + kwargs['host'] = self.endpoint + return connection.connect(**kwargs) + + def suspend( + self, + wait_on_suspended: bool = False, + wait_interval: int = 20, + wait_timeout: int = 600, + ) -> None: + """ + Suspend the workspace. + + Parameters + ---------- + wait_on_suspended : bool, optional + Wait for the workspace to go into 'Suspended' mode before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + + Raises + ------ + ManagementError + If timeout is reached + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + self._manager._post(f'workspaces/{self.id}/suspend') + if wait_on_suspended: + self._manager._wait_on_state( + self._manager.get_workspace(self.id), + 'Suspended', interval=wait_interval, timeout=wait_timeout, + ) + self.refresh() + + def resume( + self, + disable_auto_suspend: bool = False, + wait_on_resumed: bool = False, + wait_interval: int = 20, + wait_timeout: int = 600, + ) -> None: + """ + Resume the workspace. + + Parameters + ---------- + disable_auto_suspend : bool, optional + Should auto-suspend be disabled? + wait_on_resumed : bool, optional + Wait for the workspace to go into 'Resumed' or 'Active' mode before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + + Raises + ------ + ManagementError + If timeout is reached + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + self._manager._post( + f'workspaces/{self.id}/resume', + json=dict(disableAutoSuspend=disable_auto_suspend), + ) + if wait_on_resumed: + self._manager._wait_on_state( + self._manager.get_workspace(self.id), + ['Resumed', 'Active'], interval=wait_interval, timeout=wait_timeout, + ) + self.refresh() + + +class WorkspaceGroup: + """ + SingleStoreDB workspace group definition. + + .. deprecated:: + Use :class:`singlestoredb.management.cluster.Cluster`. It has no + container resource: what this class held -- region, firewall ranges, + Stage, the workspaces inside it -- belongs to the cluster itself, and + :meth:`ClusterManager.create_cluster` replaces the two-step + create-group-then-create-workspace dance. Grouping is expressed by a + :class:`~singlestoredb.management.cluster.Project` instead, which is an + organizational unit rather than a deployment parent. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`WorkspaceManager`. Workspace groups are created using + :meth:`WorkspaceManager.create_workspace_group`, or existing workspace groups are + accessed by either :attr:`WorkspaceManager.workspace_groups` or by calling + :meth:`WorkspaceManager.get_workspace_group`. + + See Also + -------- + :meth:`WorkspaceManager.create_workspace_group` + :meth:`WorkspaceManager.get_workspace_group` + :attr:`WorkspaceManager.workspace_groups` + + """ + + name: str + id: str + created_at: Optional[datetime.datetime] + region: Optional[Region] + firewall_ranges: List[str] + terminated_at: Optional[datetime.datetime] + allow_all_traffic: bool + deployment_type: Optional[str] + expires_at: Optional[datetime.datetime] + high_availability_two_zones: Optional[bool] + opt_in_preview_feature: Optional[bool] + outbound_allow_list: Optional[str] + project_id: Optional[str] + project_name: Optional[str] + smart_dr_status: Optional[str] + state: Optional[str] + update_window: Optional[Dict[str, Any]] + provider: Optional[str] + region_name: Optional[str] + + def __init__( + self, + name: str, + id: str, + created_at: Union[str, datetime.datetime], + region: Optional[Region], + firewall_ranges: List[str], + terminated_at: Optional[Union[str, datetime.datetime]], + allow_all_traffic: Optional[bool], + deployment_type: Optional[str] = None, + expires_at: Optional[Union[str, datetime.datetime]] = None, + high_availability_two_zones: Optional[bool] = None, + opt_in_preview_feature: Optional[bool] = None, + outbound_allow_list: Optional[str] = None, + project_id: Optional[str] = None, + project_name: Optional[str] = None, + smart_dr_status: Optional[str] = None, + state: Optional[str] = None, + update_window: Optional[Dict[str, Any]] = None, + provider: Optional[str] = None, + region_name: Optional[str] = None, + ): + #: Name of the workspace group + self.name = name + + #: Unique ID of the workspace group + self.id = id + + #: Timestamp of when the workspace group was created + self.created_at = to_datetime(created_at) + + #: Region of the workspace group (see :class:`Region`) + self.region = region + + #: List of allowed incoming IP addresses / ranges + self.firewall_ranges = firewall_ranges + + #: Timestamp of when the workspace group was terminated + self.terminated_at = to_datetime(terminated_at) + + #: Should all traffic be allowed? + self.allow_all_traffic = allow_all_traffic or False + + #: Deployment type of the workspace group (PRODUCTION | NON-PRODUCTION) + self.deployment_type = deployment_type + + #: Timestamp of when the workspace group will expire + self.expires_at = to_datetime(expires_at) + + #: Whether high availability across two zones is enabled + self.high_availability_two_zones = high_availability_two_zones + + #: Whether preview features are opted in + self.opt_in_preview_feature = opt_in_preview_feature + + #: Account ID for outbound connections + self.outbound_allow_list = outbound_allow_list + + #: Project ID associated with the workspace group + self.project_id = project_id + + #: Project name associated with the workspace group + self.project_name = project_name + + #: SmartDR status of the workspace group (ACTIVE | STANDBY) + self.smart_dr_status = smart_dr_status + + #: State of the workspace group (ACTIVE | PENDING | FAILED | TERMINATED) + self.state = state + + #: Update window settings: dict(day=0-6, hour=0-23) + self.update_window = update_window + + #: Cloud provider as returned by the API (raw) + self.provider = provider + + #: Cloud provider region name as returned by the API (raw) + self.region_name = region_name + + self._manager: Optional[WorkspaceManager] = 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: 'WorkspaceManager', + ) -> 'WorkspaceGroup': + """ + Construct a WorkspaceGroup from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : WorkspaceManager, optional + The WorkspaceManager the WorkspaceGroup belongs to + + Returns + ------- + :class:`WorkspaceGroup` + + """ + region_id = obj.get('regionID') + region_name = obj.get('regionName') + provider = obj.get('provider') + region = None + if region_id is not None: + region = next( + (x for x in manager.regions if x.id == region_id), None, + ) + if region is None and region_name is not None: + region = next( + ( + x for x in manager.regions + if x.region_name == region_name and x.provider == provider + ), + None, + ) + if region is None: + region = Region( + name=region_name or '', + provider=provider or '', + id=region_id, + region_name=region_name, + ) + out = cls( + name=obj['name'], + id=obj['workspaceGroupID'], + created_at=obj['createdAt'], + region=region, + firewall_ranges=obj.get('firewallRanges', []), + terminated_at=obj.get('terminatedAt'), + allow_all_traffic=obj.get('allowAllTraffic'), + deployment_type=obj.get('deploymentType'), + expires_at=obj.get('expiresAt'), + high_availability_two_zones=obj.get('highAvailabilityTwoZones'), + opt_in_preview_feature=obj.get('optInPreviewFeature'), + outbound_allow_list=obj.get('outboundAllowList'), + project_id=obj.get('projectID'), + project_name=obj.get('projectName'), + smart_dr_status=obj.get('smartDRStatus'), + state=obj.get('state'), + update_window=obj.get('updateWindow'), + provider=obj.get('provider'), + region_name=obj.get('regionName'), + ) + out._manager = manager + return out + + @property + def organization(self) -> Organization: + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + return self._manager.organization + + @property + def stage(self) -> Stage: + """Stage manager.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + return Stage(self.id, self._manager) + + stages = stage + + def refresh(self) -> 'WorkspaceGroup': + """Update the object to the current state.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + new_obj = self._manager.get_workspace_group(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + def update( + self, + name: Optional[str] = None, + firewall_ranges: Optional[List[str]] = None, + admin_password: Optional[str] = None, + expires_at: Optional[str] = None, + allow_all_traffic: Optional[bool] = None, + update_window: Optional[Dict[str, int]] = None, + deployment_type: Optional[str] = None, + ) -> None: + """ + Update the workspace group definition. + + Parameters + ---------- + name : str, optional + Name of the workspace group + firewall_ranges : list[str], optional + List of allowed CIDR ranges. An empty list indicates that all + inbound requests are allowed. + admin_password : str, optional + Admin password for the workspace group. If no password is supplied, + a password will be generated and retured in the response. + expires_at : str, optional + The timestamp of when the workspace group will expire. + If the expiration time is not specified, + the workspace group will have no expiration time. + At expiration, the workspace group is terminated and all the data is lost. + Expiration time can be specified as a timestamp or duration. + Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m" + allow_all_traffic : bool, optional + Allow all traffic to the workspace group + update_window : Dict[str, int], optional + Specify the day and hour of an update window: dict(day=0-6, hour=0-23) + deployment_type : str, optional + The deployment type that will be applied to all the workspaces + within the group (PRODUCTION | NON-PRODUCTION) + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + data = { + k: v for k, v in dict( + name=name, + firewallRanges=firewall_ranges, + adminPassword=admin_password, + expiresAt=expires_at, + allowAllTraffic=allow_all_traffic, + updateWindow=snake_to_camel_dict(update_window), + deploymentType=deployment_type, + ).items() if v is not None + } + self._manager._patch(f'workspaceGroups/{self.id}', json=data) + self.refresh() + + def terminate( + self, force: bool = False, + wait_on_terminated: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + ) -> None: + """ + Terminate the workspace group. + + Parameters + ---------- + force : bool, optional + Terminate a workspace group even if it has active workspaces + wait_on_terminated : bool, optional + Wait for the workspace group to go into 'Terminated' mode before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + + Raises + ------ + ManagementError + If timeout is reached + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + self._manager._delete(f'workspaceGroups/{self.id}', params=dict(force=force)) + if wait_on_terminated: + remaining = float(wait_timeout) + while True: + started_at = timing.now() + self.refresh() + if self.terminated_at is not None: + break + if remaining <= 0: + raise ManagementError( + msg='Exceeded waiting time for WorkspaceGroup to terminate', + ) + timing.sleep(wait_interval, 'workspace group terminated') + # Charged by measured time, so the refresh above counts against + # the timeout too. See timing.poll_cost. + remaining -= timing.poll_cost(started_at, wait_interval) + + def create_workspace( + self, + name: str, + size: Optional[str] = None, + auto_suspend: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + enable_kai: Optional[bool] = None, + wait_on_active: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + auto_scale: Optional[Dict[str, Any]] = None, + scale_factor: Optional[float] = None, + ) -> Workspace: + """ + Create a new workspace. + + Parameters + ---------- + name : str + Name of the workspace + size : str, optional + Workspace size in workspace size notation (S-00, S-1, etc.) + auto_suspend : Dict[str, Any], optional + Auto suspend settings for the workspace. If this field is not + provided, no settings will be enabled. + cache_config : float, optional + Specifies the multiplier for the persistent cache associated + with the workspace. If specified, it enables the cache configuration + multiplier. It can have one of the following values: 1, 2, or 4. + enable_kai : bool, optional + Whether to create a SingleStore Kai-enabled workspace + wait_on_active : bool, optional + Wait for the workspace to be active before returning + wait_timeout : int, optional + Maximum number of seconds to wait before raising an exception + if wait=True + wait_interval : int, optional + Number of seconds between each polling interval + auto_scale : Dict[str, Any], optional + Auto-scale settings for the workspace. + scale_factor : float, optional + Scale factor for the workspace. + + Returns + ------- + :class:`Workspace` + + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + out = self._manager.create_workspace( + name=name, + workspace_group=self, + size=size, + auto_suspend=snake_to_camel_dict(auto_suspend), + cache_config=cache_config, + enable_kai=enable_kai, + wait_on_active=wait_on_active, + wait_interval=wait_interval, + wait_timeout=wait_timeout, + auto_scale=snake_to_camel_dict(auto_scale), + scale_factor=scale_factor, + ) + + return out + + @property + def workspaces(self) -> NamedList[Workspace]: + """Return a list of available workspaces.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + res = self._manager._get('workspaces', params=dict(workspaceGroupID=self.id)) + return NamedList( + [Workspace.from_dict(item, self._manager) for item in res.json()], + ) + + +class StarterWorkspace: + """ + SingleStoreDB starter workspace definition. + + .. deprecated:: + Use :class:`singlestoredb.management.cluster.StarterCluster`. + + This object is not instantiated directly. It is used in the results + of API calls on the :class:`WorkspaceManager`. Existing starter workspaces are + accessed by either :attr:`WorkspaceManager.starter_workspaces` or by calling + :meth:`WorkspaceManager.get_starter_workspace`. + + See Also + -------- + :meth:`WorkspaceManager.get_starter_workspace` + :meth:`WorkspaceManager.create_starter_workspace` + :meth:`WorkspaceManager.terminate_starter_workspace` + :meth:`WorkspaceManager.create_starter_workspace_user` + :attr:`WorkspaceManager.starter_workspaces` + + """ + + name: str + id: str + database_name: str + endpoint: Optional[str] + mysql_dml_port: Optional[int] + websocket_port: Optional[int] + project_id: Optional[str] + + def __init__( + self, + name: str, + id: str, + database_name: str, + endpoint: Optional[str] = None, + mysql_dml_port: Optional[int] = None, + websocket_port: Optional[int] = None, + project_id: Optional[str] = None, + ): + #: Name of the starter workspace + self.name = name + + #: Unique ID of the starter workspace + self.id = id + + #: Name of the database associated with the starter workspace + self.database_name = database_name + + #: Endpoint to connect to the starter workspace. The endpoint is in the form + #: of ``hostname:port`` + self.endpoint = endpoint + + #: MySQL DML port for the starter workspace + self.mysql_dml_port = mysql_dml_port + + #: WebSocket port for the starter workspace + self.websocket_port = websocket_port + + #: Project ID associated with the starter workspace + self.project_id = project_id + + self._manager: Optional[WorkspaceManager] = 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: 'WorkspaceManager', + ) -> 'StarterWorkspace': + """ + Construct a StarterWorkspace from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : WorkspaceManager, optional + The WorkspaceManager the StarterWorkspace belongs to + + Returns + ------- + :class:`StarterWorkspace` + + """ + out = cls( + name=obj['name'], + id=obj['virtualWorkspaceID'], + database_name=obj['databaseName'], + endpoint=obj.get('endpoint'), + mysql_dml_port=obj.get('mysqlDmlPort'), + websocket_port=obj.get('websocketPort'), + project_id=obj.get('projectID'), + ) + out._manager = manager + return out + + def connect(self, **kwargs: Any) -> connection.Connection: + """ + Create a connection to the database server for this starter workspace. + + Parameters + ---------- + **kwargs : keyword-arguments, optional + Parameters to the SingleStoreDB `connect` function except host + and port which are supplied by the starter workspace object + + Returns + ------- + :class:`Connection` + + """ + if not self.endpoint: + raise ManagementError( + msg='An endpoint has not been set in this ' + 'starter workspace configuration', + ) + + kwargs['host'] = self.endpoint + kwargs['database'] = self.database_name + + return connection.connect(**kwargs) + + def terminate(self) -> None: + """Terminate the starter workspace.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + self._manager._delete(f'{SHAREDTIER_PATH}/{self.id}') + + def refresh(self) -> StarterWorkspace: + """Update the object to the current state.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + new_obj = self._manager.get_starter_workspace(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + @property + def organization(self) -> Organization: + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + return self._manager.organization + + @property + def stage(self) -> Stage: + """Stage manager.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + return Stage(self.id, self._manager) + + stages = stage + + @property + def starter_workspaces(self) -> NamedList['StarterWorkspace']: + """Return a list of available starter workspaces.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + res = self._manager._get(SHAREDTIER_PATH) + return NamedList( + [type(self).from_dict(item, self._manager) for item in res.json()], + ) + + def create_user( + self, + username: str, + password: Optional[str] = None, + ) -> Dict[str, str]: + """ + Create a new user for this starter workspace. + + Parameters + ---------- + username : str + The starter workspace user name to connect the new user to the database + password : str, optional + Password for the new user. If not provided, a password will be + auto-generated by the system. + + Returns + ------- + Dict[str, str] + Dictionary containing 'userID' and 'password' of the created user + + Raises + ------ + ManagementError + If no workspace manager is associated with this object. + """ + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + + payload = { + 'userName': username, + } + if password is not None: + payload['password'] = password + + res = self._manager._post( + f'{SHAREDTIER_PATH}/{self.id}/users', + json=payload, + ) + + response_data = res.json() + user_id = response_data.get('userID') + if not user_id: + raise ManagementError(msg='No userID returned from API') + + # Return the password provided by user or generated by API + returned_password = password if password is not None \ + else response_data.get('password') + if not returned_password: + raise ManagementError(msg='No password available from API response') + + return { + 'user_id': user_id, + 'password': returned_password, + } + + +class WorkspaceManager(Manager): + """ + SingleStoreDB workspace manager. + + .. deprecated:: + Use :class:`singlestoredb.management.cluster.ClusterManager`, via + :func:`singlestoredb.manage_clusters`. ``manage_workspaces()`` warns + and requires ``version='v1'``, since that is not what the + ``management.version`` option defaults to. + + This class should be instantiated using :func:`singlestoredb.manage_workspaces`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the workspace management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the workspace management API + + See Also + -------- + :func:`singlestoredb.manage_workspaces` + + """ + + #: Workspace management API version if none is specified. Workspaces are + #: served by v1 alone, so this is a literal rather than a reading of the + #: ``management.version`` option. + default_version = 'v1' + + #: Base URL if none is specified. + default_base_url = config.get_option('management.base_url') \ + or 'https://api.singlestore.com' + + #: Object type + obj_type = 'workspace' + + @property + def workspace_groups(self) -> NamedList[WorkspaceGroup]: + """Return a list of available workspace groups.""" + res = self._get('workspaceGroups') + return NamedList([WorkspaceGroup.from_dict(item, self) for item in res.json()]) + + @property + def starter_workspaces(self) -> NamedList[StarterWorkspace]: + """Return a list of available starter workspaces.""" + res = self._get(SHAREDTIER_PATH) + return NamedList([StarterWorkspace.from_dict(item, self) for item in res.json()]) + + @property + def organizations(self) -> Organizations: + """Return the organizations.""" + return Organizations(self) + + @property + def organization(self) -> Organization: + """ Return the current organization.""" + return self.organizations.current + + @property + def billing(self) -> Billing: + """Return the current billing information.""" + return Billing(self) + + @ttl_property(datetime.timedelta(hours=1)) + def regions(self) -> NamedList[Region]: + """Return a list of available regions.""" + res = self._get('regions') + return NamedList([Region.from_dict(item, self) for item in res.json()]) + + @ttl_property(datetime.timedelta(hours=1)) + def shared_tier_regions(self) -> NamedList[Region]: + """Return a list of regions that support shared tier workspaces.""" + res = self._get('regions/sharedtier') + return NamedList( + [Region.from_dict(item, self) for item in res.json()], + ) + + def create_workspace_group( + self, + name: str, + region: Union[str, Region], + firewall_ranges: List[str], + admin_password: Optional[str] = None, + backup_bucket_kms_key_id: Optional[str] = None, + data_bucket_kms_key_id: Optional[str] = None, + expires_at: Optional[str] = None, + smart_dr: Optional[bool] = None, + allow_all_traffic: Optional[bool] = None, + update_window: Optional[Dict[str, int]] = None, + provider: Optional[str] = None, + region_name: Optional[str] = None, + deployment_type: Optional[str] = None, + high_availability_two_zones: Optional[bool] = None, + opt_in_preview_feature: Optional[bool] = None, + project_id: Optional[str] = None, + ) -> WorkspaceGroup: + """ + Create a new workspace group. + + Parameters + ---------- + name : str + Name of the workspace group + region : str or Region + ID of the region where the workspace group should be created + firewall_ranges : list[str] + List of allowed CIDR ranges. An empty list indicates that all + inbound requests are allowed. + admin_password : str, optional + Admin password for the workspace group. If no password is supplied, + a password will be generated and retured in the response. + backup_bucket_kms_key_id : str, optional + Specifies the KMS key ID associated with the backup bucket. + If specified, enables Customer-Managed Encryption Keys (CMEK) + encryption for the backup bucket of the workspace group. + This feature is only supported in workspace groups deployed in AWS. + data_bucket_kms_key_id : str, optional + Specifies the KMS key ID associated with the data bucket. + If specified, enables Customer-Managed Encryption Keys (CMEK) + encryption for the data bucket and Amazon Elastic Block Store + (EBS) volumes of the workspace group. This feature is only supported + in workspace groups deployed in AWS. + expires_at : str, optional + The timestamp of when the workspace group will expire. + If the expiration time is not specified, + the workspace group will have no expiration time. + At expiration, the workspace group is terminated and all the data is lost. + Expiration time can be specified as a timestamp or duration. + Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m" + smart_dr : bool, optional + Enables Smart Disaster Recovery (SmartDR) for the workspace group. + SmartDR is a disaster recovery solution that ensures seamless and + continuous replication of data from the primary region to a secondary region + allow_all_traffic : bool, optional + Allow all traffic to the workspace group + update_window : Dict[str, int], optional + Specify the day and hour of an update window: dict(day=0-6, hour=0-23) + provider : str, optional + Cloud provider for the workspace group (e.g., 'AWS', 'GCP', 'AZURE'). + Used together with ``region_name`` as an alternative to ``region``. + region_name : str, optional + Cloud provider region name for the workspace group. Used together + with ``provider`` as an alternative to ``region``. + deployment_type : str, optional + Deployment type for workspaces in this group (PRODUCTION | + NON-PRODUCTION). + high_availability_two_zones : bool, optional + Whether to enable high availability across two zones. + opt_in_preview_feature : bool, optional + Whether to opt in to preview features. + project_id : str, optional + Project ID to associate the workspace group with. + + Returns + ------- + :class:`WorkspaceGroup` + + """ + region_id: Optional[str] = None + if isinstance(region, Region): + if region.id: + region_id = region.id + else: + if provider is None: + provider = region.provider + if region_name is None: + region_name = region.region_name + else: + region_id = region + res = self._post( + 'workspaceGroups', json=dict( + name=name, regionID=region_id, + adminPassword=admin_password, + backupBucketKMSKeyID=backup_bucket_kms_key_id, + dataBucketKMSKeyID=data_bucket_kms_key_id, + firewallRanges=firewall_ranges or [], + expiresAt=expires_at, + smartDR=smart_dr, + allowAllTraffic=allow_all_traffic, + updateWindow=snake_to_camel_dict(update_window), + provider=provider, + regionName=region_name, + deploymentType=deployment_type, + highAvailabilityTwoZones=high_availability_two_zones, + optInPreviewFeature=opt_in_preview_feature, + projectID=project_id, + ), + ) + return self.get_workspace_group(res.json()['workspaceGroupID']) + + def create_workspace( + self, + name: str, + workspace_group: Union[str, WorkspaceGroup], + size: Optional[str] = None, + auto_suspend: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + enable_kai: Optional[bool] = None, + wait_on_active: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + auto_scale: Optional[Dict[str, Any]] = None, + scale_factor: Optional[float] = None, + ) -> Workspace: + """ + Create a new workspace. + + Parameters + ---------- + name : str + Name of the workspace + workspace_group : str or WorkspaceGroup + The workspace ID of the workspace + size : str, optional + Workspace size in workspace size notation (S-00, S-1, etc.) + auto_suspend : Dict[str, Any], optional + Auto suspend settings for the workspace. If this field is not + provided, no settings will be enabled. + cache_config : float, optional + Specifies the multiplier for the persistent cache associated + with the workspace. If specified, it enables the cache configuration + multiplier. It can have one of the following values: 1, 2, or 4. + enable_kai : bool, optional + Whether to create a SingleStore Kai-enabled workspace + wait_on_active : bool, optional + Wait for the workspace to be active before returning + wait_timeout : int, optional + Maximum number of seconds to wait before raising an exception + if wait=True + wait_interval : int, optional + Number of seconds between each polling interval + auto_scale : Dict[str, Any], optional + Auto-scale settings for the workspace. + scale_factor : float, optional + Scale factor for the workspace. + + Returns + ------- + :class:`Workspace` + + """ + if isinstance(workspace_group, WorkspaceGroup): + workspace_group = workspace_group.id + res = self._post( + 'workspaces', json=dict( + name=name, + workspaceGroupID=workspace_group, + size=size, + autoSuspend=snake_to_camel_dict(auto_suspend), + cacheConfig=cache_config, + enableKai=enable_kai, + autoScale=snake_to_camel_dict(auto_scale), + scaleFactor=scale_factor, + ), + ) + out = self.get_workspace(res.json()['workspaceID']) + if wait_on_active: + out = self._wait_on_state( + out, + 'Active', + interval=wait_interval, + timeout=wait_timeout, + ) + # After workspace is active, wait for endpoint to be ready + out = self._wait_on_endpoint( + out, + interval=wait_interval, + timeout=wait_timeout, + ) + return out + + def get_workspace_group(self, id: str) -> WorkspaceGroup: + """ + Retrieve a workspace group definition. + + Parameters + ---------- + id : str + ID of the workspace group + + Returns + ------- + :class:`WorkspaceGroup` + + """ + res = self._get(f'workspaceGroups/{id}') + return WorkspaceGroup.from_dict(res.json(), manager=self) + + def get_workspace(self, id: str) -> Workspace: + """ + Retrieve a workspace definition. + + Parameters + ---------- + id : str + ID of the workspace + + Returns + ------- + :class:`Workspace` + + """ + res = self._get(f'workspaces/{id}') + return Workspace.from_dict(res.json(), manager=self) + + def get_starter_workspace(self, id: str) -> StarterWorkspace: + """ + Retrieve a starter workspace definition. + + Parameters + ---------- + id : str + ID of the starter workspace + + Returns + ------- + :class:`StarterWorkspace` + + """ + res = self._get(f'{SHAREDTIER_PATH}/{id}') + return StarterWorkspace.from_dict(res.json(), manager=self) + + def create_starter_workspace( + self, + name: str, + database_name: str, + provider: str, + region_name: str, + project_id: Optional[str] = None, + ) -> 'StarterWorkspace': + """ + Create a new starter (shared tier) workspace. + + Parameters + ---------- + name : str + Name of the starter workspace + database_name : str + Name of the database for the starter workspace + provider : str + Cloud provider for the starter workspace (e.g., 'aws', 'gcp', 'azure') + region_name : str + Cloud provider region for the starter workspace (e.g., 'us-east-1') + project_id : str, optional + Project ID to associate the starter workspace with. + + Returns + ------- + :class:`StarterWorkspace` + """ + + payload: Dict[str, Any] = { + 'name': name, + 'databaseName': database_name, + 'provider': provider, + 'regionName': region_name, + } + if project_id is not None: + payload['projectID'] = project_id + + res = self._post(SHAREDTIER_PATH, json=payload) + virtual_workspace_id = res.json().get('virtualWorkspaceID') + if not virtual_workspace_id: + raise ManagementError(msg='No virtualWorkspaceID returned from API') + + res = self._get(f'{SHAREDTIER_PATH}/{virtual_workspace_id}') + return StarterWorkspace.from_dict(res.json(), self) diff --git a/singlestoredb/management/v2/__init__.py b/singlestoredb/management/v2/__init__.py new file mode 100644 index 000000000..612bcdaef --- /dev/null +++ b/singlestoredb/management/v2/__init__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python +"""SingleStoreDB Management API v2.""" +# The version-neutral helpers in singlestoredb.management look these up here by +# name. A deployment is a cluster, so they live in the cluster module. +from .cluster import get_organization as get_organization +from .cluster import get_secret as get_secret +from .cluster import get_stage as get_stage diff --git a/singlestoredb/management/v2/billing_usage.py b/singlestoredb/management/v2/billing_usage.py new file mode 100644 index 000000000..b141146b9 --- /dev/null +++ b/singlestoredb/management/v2/billing_usage.py @@ -0,0 +1,9 @@ +#!/usr/bin/env python +""" +SingleStoreDB Billing Usage API v2. + +``GET /v2/billing/usage`` is implemented in +:mod:`singlestoredb.management.billing_usage`, so this module only re-exports it. +""" +from ..billing_usage import BillingUsageItem as BillingUsageItem +from ..billing_usage import UsageItem as UsageItem diff --git a/singlestoredb/management/v2/cluster.py b/singlestoredb/management/v2/cluster.py new file mode 100644 index 000000000..511b7cace --- /dev/null +++ b/singlestoredb/management/v2/cluster.py @@ -0,0 +1,1694 @@ +#!/usr/bin/env python +""" +SingleStoreDB Cluster Management API v2. + +A deployment is a single flat ``clusters`` resource: one :class:`Cluster` +carries both the deployment's own settings -- size, state, connection endpoints +-- and the account-level settings around it, such as the firewall ranges and the +admin credentials. +""" +from __future__ import annotations + +import datetime +import re +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +from .. import timing +from ... import config +from ... import connection +from ...exceptions import ManagementError +from ..billing import Billing as Billing +from ..manager import Manager +from ..organization import Organization +from ..organization import Organizations as Organizations +from ..region import Region +from ..stage import Stage as Stage +from ..stage import StageObject as StageObject +from ..utils import camel_to_snake_dict +from ..utils import get_cluster_id +from ..utils import NamedList +from ..utils import PathLike +from ..utils import snake_to_camel_dict +from ..utils import to_datetime +from ..utils import ttl_property +from ..utils import vars_to_str +from .project import Project as Project + +#: Base management API path for the shared-tier resource. +SHAREDTIER_PATH = 'sharedtier/virtualClusters' + +#: Shape of a project ID. Anywhere a project can be named, a name is accepted +#: in place of an ID, and this is how the two are told apart. Sending a project +#: ID that is not a UUID comes back as ``400 uuid: incorrect UUID length``, so +#: a value that does not match this could never have been a valid ID and +#: nothing is lost by reading it as a name. +PROJECT_ID_RE = re.compile( + r'^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}' + r'-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$', +) + + +def _project_from_id( + manager: Optional['ClusterManager'], + project_id: Optional[str], +) -> Optional[Project]: + """ + Return the project with the given ID, as reported by ``manager``. + + A deployment reports only its ``projectID``, so the rest of the project is + recovered from :attr:`ClusterManager.projects` -- a cached list, so this + costs nothing per deployment after the first. An ID that matches no project + -- or no manager to match it against -- still yields a :class:`Project`, + carrying the ID and nothing else, so that ``cluster.project.id`` is always + readable. + """ + if project_id is None: + return None + if manager is None: + return Project(id=project_id, name='') + return next( + (x for x in manager.projects if x.id == project_id), + Project(id=project_id, name=''), + ) + + +def _region_from_name( + manager: Optional['ClusterManager'], + region_name: Optional[str], + provider: Optional[str], +) -> Optional[Region]: + """ + Return the region a deployment reported, as reported by ``manager``. + + No region ID is reported, so a region is identified by the + ``(provider, region_name)`` pair, and the display name lives only in + :attr:`ClusterManager.regions` -- a cached list, so this costs nothing per + deployment after the first. An unmatched pair -- or no manager to match it + against -- still yields a :class:`Region`, built from what the deployment + itself reports, so that ``cluster.region.region_name`` is always readable. + """ + if region_name is None: + return None + if manager is not None: + for region in manager.regions: + if region.region_name == region_name and region.provider == provider: + return region + return Region( + name=region_name, + provider=provider or '', + region_name=region_name, + ) + + +def _project_args( + project: Union[str, Project, None], +) -> Tuple[Optional[Project], Optional[str]]: + """ + Split a ``project`` constructor argument into a project and a project ID. + + A :class:`Project` is a resolved project and is kept as it stands; a string + is a project ID, which :func:`_lazy_project` resolves when it is asked for. + """ + if isinstance(project, Project): + return project, project.id + return None, project + + +def _region_args( + region: Union[str, Region, None], +) -> Tuple[Optional[Region], Optional[str]]: + """ + Split a ``region`` constructor argument into a region and a region name. + + A :class:`Region` is a resolved region and is kept as it stands; a string is + a provider region name, e.g. ``us-east-1``, which :func:`_lazy_region` + resolves when it is asked for. + """ + if isinstance(region, Region): + return region, region.region_name or region.name + return None, region + + +def _lazy_project(deployment: Any) -> Optional[Project]: + """ + Return the project of a deployment that reported only its project ID. + + Resolving the ID costs a ``GET /v2/projects`` for the first deployment a + manager resolves one for, so it happens on demand: a listing of N + deployments that nobody asks the project of costs nothing, and one that is + asked costs the one request, because :attr:`ClusterManager.projects` is + cached. + + Works on anything carrying the ``_project``, ``_project_id`` and + ``_manager`` attributes, which is :class:`Cluster` and + :class:`StarterCluster`. + """ + if deployment._project is None and deployment._project_id is not None: + deployment._project = _project_from_id( + deployment._manager, deployment._project_id, + ) + return deployment._project + + +def _lazy_region(deployment: Any) -> Optional[Region]: + """ + Return the region of a deployment that reported only a region name. + + On demand, and for the same reason as :func:`_lazy_project`: matching the + name costs a ``GET /v2/regions``, and a listing whose regions nobody reads + should not pay it. + """ + if deployment._region is None and deployment._region_name is not None: + deployment._region = _region_from_name( + deployment._manager, deployment._region_name, deployment.provider, + ) + return deployment._region + + +def get_organization() -> Organization: + """Get the organization.""" + from ..cluster import manage_clusters + # Pinned: this module's helpers are v2's own, so they must not follow the + # management.version option elsewhere. + return manage_clusters(version='v2').organization + + +def get_secret(name: str) -> Optional[str]: + """Get a secret from the organization.""" + return get_organization().get_secret(name).value + + +def get_cluster( + cluster: Optional[Union['Cluster', str]] = None, +) -> 'Cluster': + """ + Get a cluster. + + Parameters + ---------- + cluster : Cluster or str, optional + A cluster object, or the name or ID of a cluster. If not given, + ``SINGLESTOREDB_WORKSPACE`` is used: the notebook environment publishes + no ``SINGLESTOREDB_CLUSTER``, and that variable carries the ID of the + current cluster. ``SINGLESTOREDB_WORKSPACE_GROUP`` is *not* consulted -- + it holds a group ID, which is reported only as the read-only + :attr:`Cluster.group` and offers no route to look up. + + Returns + ------- + :class:`Cluster` + + """ + if isinstance(cluster, Cluster): + return cluster + from ..cluster import manage_clusters + mgr = manage_clusters(version='v2') + if cluster: + return mgr.clusters[cluster] + from_env = get_cluster_id() + if from_env: + return mgr.clusters[from_env] + raise RuntimeError('no cluster specified') + + +def get_stage( + cluster: Optional[Union['Cluster', str]] = None, +) -> Stage: + """Get the stage for a cluster.""" + return get_cluster(cluster).stage + + +class Cluster: + """ + SingleStoreDB cluster definition. + + This object is not instantiated directly. It is used in the results of API + calls on the :class:`ClusterManager`. Clusters are created using + :meth:`ClusterManager.create_cluster`, or existing clusters are accessed by + either :attr:`ClusterManager.clusters` or by calling + :meth:`ClusterManager.get_cluster`. + + A cluster is a single flat resource: the compute settings (size, + auto-suspend, cache) and the deployment-wide settings (firewall, update + window, expiration) all live on this object. + + See Also + -------- + :meth:`ClusterManager.create_cluster` + :meth:`ClusterManager.get_cluster` + :attr:`ClusterManager.clusters` + + """ + + name: str + id: str + group: Optional[str] + size: Optional[str] + scale_factor: Optional[float] + state: str + created_at: Optional[datetime.datetime] + terminated_at: Optional[datetime.datetime] + expires_at: Optional[datetime.datetime] + last_resumed_at: Optional[datetime.datetime] + endpoint: Optional[str] + provider: Optional[str] + deployment_type: Optional[str] + kai: Optional[bool] + multi_az: Optional[bool] + allow_all_traffic: bool + firewall_ranges: List[str] + outbound_allow_list: Optional[str] + opt_in_preview_feature: Optional[bool] + update_window: Optional[Dict[str, Any]] + auto_suspend: Optional[Dict[str, Any]] + auto_scale: Optional[Dict[str, Any]] + cache_config: Optional[float] + resume_attachments: List[Dict[str, Any]] + scaling_progress: Optional[int] + smart_dr_status: Optional[str] + + def __init__( + self, + name: str, + id: str, + state: str, + group: Optional[str] = None, + size: Optional[str] = None, + scale_factor: Optional[float] = None, + created_at: Optional[Union[str, datetime.datetime]] = None, + terminated_at: Optional[Union[str, datetime.datetime]] = None, + expires_at: Optional[Union[str, datetime.datetime]] = None, + last_resumed_at: Optional[Union[str, datetime.datetime]] = None, + endpoint: Optional[str] = None, + provider: Optional[str] = None, + region: Union[str, Region, None] = None, + project: Union[str, Project, None] = None, + deployment_type: Optional[str] = None, + kai: Optional[bool] = None, + multi_az: Optional[bool] = None, + allow_all_traffic: Optional[bool] = None, + firewall_ranges: Optional[List[str]] = None, + outbound_allow_list: Optional[str] = None, + opt_in_preview_feature: Optional[bool] = None, + update_window: Optional[Dict[str, Any]] = None, + auto_suspend: Optional[Dict[str, Any]] = None, + auto_scale: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + resume_attachments: Optional[List[Dict[str, Any]]] = None, + scaling_progress: Optional[int] = None, + smart_dr_status: Optional[str] = None, + ): + #: Name of the cluster + self.name = name + + #: Unique ID of the cluster + self.id = id + + #: State of the cluster: PENDING, ACTIVE, SUSPENDED, TERMINATED, + #: TRANSITIONING, RESUMING, FAILED + self.state = state.strip() + + #: Unique ID of the group the cluster belongs to. There is no group + #: route, so this is an opaque ID rather than a lookup key. + self.group = group + + #: Size of the cluster in cluster size notation (S-00, S-1, etc.) + self.size = size + + #: Current scale factor for the cluster + self.scale_factor = scale_factor + + #: Timestamp of when the cluster was created + self.created_at = to_datetime(created_at) + + #: Timestamp of when the cluster was terminated + self.terminated_at = to_datetime(terminated_at) + + #: Timestamp of when the cluster will expire + self.expires_at = to_datetime(expires_at) + + #: Timestamp of when the cluster was last resumed + self.last_resumed_at = to_datetime(last_resumed_at) + + #: Hostname (or IP address) of the cluster database server + self.endpoint = endpoint + + #: Cloud provider hosting the cluster (AWS | GCP | Azure) + self.provider = provider + + # Region the cluster is deployed in; see the region property. A string + # is taken as the provider region name, e.g. us-east-1, and is not + # resolved until it is asked for, so that listing clusters costs no + # GET /v2/regions. + self._region, self._region_name = _region_args(region) + + # Project the cluster belongs to; see the project property. A string + # is taken as the project ID and is not resolved until it is asked + # for, so that listing clusters costs no GET /v2/projects. + self._project, self._project_id = _project_args(project) + + #: Deployment type of the cluster (PRODUCTION | NON-PRODUCTION) + self.deployment_type = deployment_type + + #: Whether SingleStore Kai is enabled on this cluster + self.kai = kai + + #: Whether the cluster is deployed across multiple availability zones + self.multi_az = multi_az + + #: Should all inbound traffic be allowed? + self.allow_all_traffic = allow_all_traffic or False + + #: List of allowed incoming IP addresses / ranges + self.firewall_ranges = firewall_ranges or [] + + #: Account ID for outbound connections + self.outbound_allow_list = outbound_allow_list + + #: Whether preview features are opted in + self.opt_in_preview_feature = opt_in_preview_feature + + #: Update window settings: dict(day=0-6, hour=0-23) + self.update_window = update_window + + #: Current auto-suspend settings + self.auto_suspend = camel_to_snake_dict(auto_suspend) + + #: Auto-scale settings for the cluster + self.auto_scale = camel_to_snake_dict(auto_scale) + + #: Multiplier for the persistent cache + self.cache_config = cache_config + + #: Database attachments + self.resume_attachments = [ + camel_to_snake_dict(x) # type: ignore + for x in resume_attachments or [] + if x is not None + ] + + #: Current progress percentage for scaling the cluster + self.scaling_progress = scaling_progress + + #: SmartDR status of the cluster (ACTIVE | STANDBY) + self.smart_dr_status = smart_dr_status + + self._manager: Optional[ClusterManager] = None + + # Set by ClusterManager.create_cluster only; see the admin_password + # property. Private so it stays out of str() / repr(). + self._admin_password: Optional[str] = None + + @property + def region(self) -> Optional[Region]: + """ + Region the cluster is deployed in, or ``None`` if it reported none. + + No region ID is reported: a region is identified by the + ``(provider, region_name)`` pair, and the display name lives only in + :attr:`ClusterManager.regions` -- a request, and one that listing + clusters would otherwise pay for every row, so it is made the first time + this is read rather than when the cluster is built. An unmatched pair + still yields a :class:`Region` built from what the cluster itself + reports, so ``cluster.region.region_name`` is always readable. + """ + return _lazy_region(self) + + @property + def project(self) -> Optional[Project]: + """ + Project the cluster belongs to, or ``None`` if it reported no project. + + A cluster reports only its ``projectID``, so the rest of the project + comes from :attr:`ClusterManager.projects` -- a request, and one that + listing clusters would otherwise pay for every row, so it is made the + first time this is read rather than when the cluster is built. An ID + that matches no project still yields a :class:`Project` carrying the + ID, so ``cluster.project.id`` is always readable. + """ + return _lazy_project(self) + + @property + def admin_password(self) -> Optional[str]: + """ + Generated password for the ``admin`` database user. + + ``POST /v2/clusters`` generates the admin password itself and returns it + in the create response -- the ``admin_password`` passed to + :meth:`ClusterManager.create_cluster` is ignored -- and no other route + reports it. So this is set on the cluster returned by ``create_cluster`` + and is ``None`` everywhere else, including after :meth:`refresh`. Record + it when the cluster is created or it cannot be recovered. + + """ + return self._admin_password + + def __str__(self) -> str: + """Return string representation.""" + # project and region are resolved lazily, so they are not in vars(self). + # Report whatever is already in hand -- the resolved object if something + # has read the property, otherwise the ID / name the cluster itself + # reported -- so that printing a cluster never issues a request. + return vars_to_str( + self, extra=dict( + project=self._project or self._project_id, + region=self._region or self._region_name, + ), + ) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict(cls, obj: Dict[str, Any], manager: 'ClusterManager') -> 'Cluster': + """ + Construct a Cluster from a dictionary of values. + + Every field other than the name and ID is optional: the API omits null + fields entirely rather than returning them as ``null``. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : ClusterManager + The ClusterManager the Cluster belongs to + + Returns + ------- + :class:`Cluster` + + """ + # Size is reported as an object: dict(size='S-00', scaleFactor=1), + # keyed as ``sizeConfig``. The older ``size`` key is read as a + # fallback, so a response using either name still populates + # :attr:`Cluster.size`. The ``size`` argument and + # :attr:`Cluster.size` are wrapper-side names either way. + size_spec = obj.get('sizeConfig') or obj.get('size') or {} + + out = cls( + name=obj['name'], + id=obj['clusterID'], + state=obj.get('state', 'Unknown'), + group=obj.get('groupID'), + size=size_spec.get('size'), + scale_factor=size_spec.get('scaleFactor'), + created_at=obj.get('createdAt'), + terminated_at=obj.get('terminatedAt'), + expires_at=obj.get('expiresAt'), + last_resumed_at=obj.get('lastResumedAt'), + endpoint=obj.get('endpoint'), + provider=obj.get('provider'), + # The provider region name and the project ID are all the response + # carries; the region and project properties resolve them against + # the manager's cached listings when they are read. + region=obj.get('region'), + project=obj.get('projectID'), + deployment_type=obj.get('deploymentType'), + kai=obj.get('kai'), + multi_az=obj.get('multiAZ'), + allow_all_traffic=obj.get('allowAllTraffic'), + firewall_ranges=obj.get('firewallRanges'), + outbound_allow_list=obj.get('outboundAllowList'), + opt_in_preview_feature=obj.get('optInPreviewFeature'), + update_window=obj.get('updateWindow'), + auto_suspend=obj.get('autoSuspend'), + auto_scale=obj.get('autoScale'), + cache_config=obj.get('cacheConfig'), + resume_attachments=obj.get('resumeAttachments'), + scaling_progress=obj.get('scalingProgress'), + smart_dr_status=obj.get('smartDRStatus'), + ) + out._manager = manager + return out + + def _require_manager(self) -> 'ClusterManager': + if self._manager is None: + raise ManagementError( + msg='No cluster manager is associated with this object.', + ) + return self._manager + + @property + def organization(self) -> Organization: + """Return the organization the cluster belongs to.""" + return self._require_manager().organization + + @property + def stage(self) -> Stage: + """Stage manager.""" + return Stage(self.id, self._require_manager()) + + stages = stage + + def refresh(self) -> 'Cluster': + """Update the object to the current state.""" + manager = self._require_manager() + new_obj = manager.get_cluster(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + def update( + self, + name: Optional[str] = None, + size: Optional[str] = None, + scale_factor: Optional[float] = None, + auto_suspend: Optional[Dict[str, Any]] = None, + auto_scale: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + deployment_type: Optional[str] = None, + firewall_ranges: Optional[List[str]] = None, + allow_all_traffic: Optional[bool] = None, + admin_password: Optional[str] = None, + expires_at: Optional[str] = None, + update_window: Optional[Dict[str, int]] = None, + kai: Optional[bool] = None, + wait_on_active: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + ) -> None: + """ + Update the cluster definition. + + Both the compute settings (size, auto-suspend, cache) and the + deployment-wide settings (firewall, update window, expiration) are + changed through this one call. + + The API applies the ``PATCH`` asynchronously: the cluster cycles back + through PENDING and the trailing :meth:`refresh` still reports the + pre-PATCH values. Pass ``wait_on_active=True`` to wait the change out + so the object reflects it on return. + + Parameters + ---------- + name : str, optional + Name of the cluster + size : str, optional + Size of the cluster in cluster size notation, such as "S-1". + Resizing is done through this field; there is no ``resize`` route. + Sent nested in a ``size`` object alongside ``scale_factor``. + scale_factor : float, optional + Scale factor for the cluster + auto_suspend : Dict[str, Any], optional + Auto-suspend mode for the cluster: IDLE, SCHEDULED, DISABLED + auto_scale : Dict[str, Any], optional + Auto-scale settings for the cluster + cache_config : float, optional + Multiplier for the persistent cache associated with the cluster. + It can have one of the following values: 1, 2, or 4. + deployment_type : str, optional + Deployment type of the cluster (PRODUCTION | NON-PRODUCTION) + firewall_ranges : List[str], optional + List of allowed CIDR ranges. An empty list denies all inbound + traffic; omitting it leaves the current ranges alone. + allow_all_traffic : bool, optional + Allow all traffic to the cluster + admin_password : str, optional + Admin password for the cluster + expires_at : str, optional + Timestamp of when the cluster will expire. Expiration time can be + specified as a timestamp or a duration. + Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m" + update_window : Dict[str, int], optional + Day and hour of an update window: dict(day=0-6, hour=0-23) + kai : bool, optional + Whether SingleStore Kai is enabled on this cluster + wait_on_active : bool, optional + Wait for the cluster to be ACTIVE again -- and, if a firewall was + requested, for the new ranges to be reported -- before returning. + Defaults to ``False``, which returns as soon as the ``PATCH`` is + accepted and therefore reports pre-PATCH values. + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Maximum number of seconds to wait before raising an exception + + Raises + ------ + ManagementError + If ``wait_on_active`` is given and the timeout is reached + + """ + manager = self._require_manager() + size_spec: Optional[Dict[str, Any]] = None + if size is not None or scale_factor is not None: + size_spec = { + k: v for k, v in dict( + size=size, scaleFactor=scale_factor, + ).items() if v is not None + } + data = { + k: v for k, v in dict( + name=name, + # ``sizeConfig``, not ``size``; see Cluster.from_dict. + sizeConfig=size_spec, + autoSuspend=snake_to_camel_dict(auto_suspend), + autoScale=snake_to_camel_dict(auto_scale), + cacheConfig=cache_config, + deploymentType=deployment_type, + firewallRanges=firewall_ranges, + allowAllTraffic=allow_all_traffic, + adminPassword=admin_password, + expiresAt=expires_at, + updateWindow=snake_to_camel_dict(update_window), + kai=kai, + ).items() if v is not None + } + manager._patch(f'clusters/{self.id}', json=data) + + if wait_on_active: + out = manager._wait_on_state( + manager.get_cluster(self.id), 'ACTIVE', + interval=wait_interval, timeout=wait_timeout, + ) + if firewall_ranges or allow_all_traffic: + manager._wait_on_firewall( + out, interval=wait_interval, timeout=wait_timeout, + expected=firewall_ranges, + ) + + self.refresh() + + def terminate( + self, + wait_on_terminated: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + force: bool = False, + ) -> None: + """ + Terminate the cluster. + + Parameters + ---------- + wait_on_terminated : bool, optional + Wait for the cluster to be terminated before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + force : bool, optional + Should the cluster be terminated even if it is in use? + + Raises + ------ + ManagementError + If timeout is reached + + """ + manager = self._require_manager() + manager._delete(f'clusters/{self.id}', params=dict(force=force)) + if wait_on_terminated: + remaining = float(wait_timeout) + while True: + started_at = timing.now() + self.refresh() + if self.terminated_at is not None: + break + if remaining <= 0: + raise ManagementError( + msg='Exceeded waiting time for Cluster to terminate', + ) + timing.sleep(wait_interval, 'cluster terminated') + # Charged by measured time, so the refresh above counts against + # the timeout too. See timing.poll_cost. + remaining -= timing.poll_cost(started_at, wait_interval) + + def connect(self, **kwargs: Any) -> connection.Connection: + """ + Create a connection to the database server for this cluster. + + Parameters + ---------- + **kwargs : keyword-arguments, optional + Parameters to the SingleStoreDB `connect` function except host + and port which are supplied by the cluster object + + Returns + ------- + :class:`Connection` + + """ + if not self.endpoint: + raise ManagementError( + msg='An endpoint has not been set in this cluster configuration', + ) + kwargs['host'] = self.endpoint + return connection.connect(**kwargs) + + def suspend( + self, + wait_on_suspended: bool = False, + wait_interval: int = 20, + wait_timeout: int = 600, + ) -> None: + """ + Suspend the cluster. + + Parameters + ---------- + wait_on_suspended : bool, optional + Wait for the cluster to be suspended before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + + Raises + ------ + ManagementError + If timeout is reached + + """ + manager = self._require_manager() + manager._post(f'clusters/{self.id}/suspend') + if wait_on_suspended: + manager._wait_on_state( + manager.get_cluster(self.id), + 'SUSPENDED', interval=wait_interval, timeout=wait_timeout, + ) + self.refresh() + + def resume( + self, + disable_auto_suspend: bool = False, + wait_on_resumed: bool = False, + wait_interval: int = 20, + wait_timeout: int = 600, + ) -> None: + """ + Resume the cluster. + + Parameters + ---------- + disable_auto_suspend : bool, optional + Should auto-suspend be disabled? + wait_on_resumed : bool, optional + Wait for the cluster to be resumed or active before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + + Raises + ------ + ManagementError + If timeout is reached + + """ + manager = self._require_manager() + manager._post( + f'clusters/{self.id}/resume', + json=dict(disableAutoSuspend=disable_auto_suspend), + ) + if wait_on_resumed: + manager._wait_on_state( + manager.get_cluster(self.id), + ['RESUMED', 'ACTIVE'], interval=wait_interval, timeout=wait_timeout, + ) + self.refresh() + + +class StarterCluster: + """ + SingleStoreDB starter (shared tier) cluster definition. + + This object is not instantiated directly. Existing starter clusters are + accessed by either :attr:`ClusterManager.starter_clusters` or by calling + :meth:`ClusterManager.get_starter_cluster`. + + See Also + -------- + :meth:`ClusterManager.get_starter_cluster` + :meth:`ClusterManager.create_starter_cluster` + :attr:`ClusterManager.starter_clusters` + + """ + + name: str + id: str + database_name: str + endpoint: Optional[str] + mysql_dml_port: Optional[int] + websocket_port: Optional[int] + + def __init__( + self, + name: str, + id: str, + database_name: str, + endpoint: Optional[str] = None, + mysql_dml_port: Optional[int] = None, + websocket_port: Optional[int] = None, + project: Union[str, Project, None] = None, + ): + #: Name of the starter cluster + self.name = name + + #: Unique ID of the starter cluster + self.id = id + + #: Name of the database associated with the starter cluster + self.database_name = database_name + + #: Endpoint to connect to the starter cluster, in the form + #: ``hostname:port`` + self.endpoint = endpoint + + #: MySQL DML port for the starter cluster + self.mysql_dml_port = mysql_dml_port + + #: WebSocket port for the starter cluster + self.websocket_port = websocket_port + + # Project the starter cluster belongs to; see the project property. A + # string is taken as the project ID and is not resolved until it is + # asked for, so that listing starter clusters costs no + # GET /v2/projects. + self._project, self._project_id = _project_args(project) + + self._manager: Optional[ClusterManager] = None + + @property + def project(self) -> Optional[Project]: + """ + Project the starter cluster belongs to, or ``None`` if it reported + no project. + + Resolved on first read from :attr:`ClusterManager.projects`; see + :attr:`Cluster.project`. + """ + return _lazy_project(self) + + def __str__(self) -> str: + """Return string representation.""" + # See Cluster.__str__: project is lazy, so report what is in hand + # rather than resolving it just to print. + return vars_to_str( + self, extra=dict(project=self._project or self._project_id), + ) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + @classmethod + def from_dict( + cls, obj: Dict[str, Any], manager: 'ClusterManager', + ) -> 'StarterCluster': + """ + Construct a StarterCluster from a dictionary of values. + + Parameters + ---------- + obj : dict + Dictionary of values + manager : ClusterManager + The ClusterManager the StarterCluster belongs to + + Returns + ------- + :class:`StarterCluster` + + """ + out = cls( + name=obj['name'], + id=obj['virtualClusterID'], + database_name=obj['databaseName'], + endpoint=obj.get('endpoint'), + mysql_dml_port=obj.get('mysqlDmlPort'), + websocket_port=obj.get('websocketPort'), + project=obj.get('projectID'), + ) + out._manager = manager + return out + + def _require_manager(self) -> 'ClusterManager': + if self._manager is None: + raise ManagementError( + msg='No cluster manager is associated with this object.', + ) + return self._manager + + def connect(self, **kwargs: Any) -> connection.Connection: + """ + Create a connection to the database server for this starter cluster. + + Parameters + ---------- + **kwargs : keyword-arguments, optional + Parameters to the SingleStoreDB `connect` function except host + and port which are supplied by the starter cluster object + + Returns + ------- + :class:`Connection` + + """ + if not self.endpoint: + raise ManagementError( + msg='An endpoint has not been set in this ' + 'starter cluster configuration', + ) + kwargs['host'] = self.endpoint + kwargs['database'] = self.database_name + return connection.connect(**kwargs) + + def terminate(self) -> None: + """Terminate the starter cluster.""" + self._require_manager()._delete(f'{SHAREDTIER_PATH}/{self.id}') + + def refresh(self) -> 'StarterCluster': + """Update the object to the current state.""" + manager = self._require_manager() + new_obj = manager.get_starter_cluster(self.id) + for name, value in vars(new_obj).items(): + setattr(self, name, value) + return self + + @property + def organization(self) -> Organization: + """Return the organization the starter cluster belongs to.""" + return self._require_manager().organization + + @property + def stage(self) -> Stage: + """ + Stage manager. + + .. warning:: There is no Stage route for a starter (shared tier) + deployment at either API version -- ``clusters/{id}/stage/fs/`` + only resolves for a full cluster ID. This property is kept for + parity with :class:`Cluster`, but requests made through it will + fail. + + """ + return Stage(self.id, self._require_manager()) + + stages = stage + + @property + def starter_clusters(self) -> NamedList['StarterCluster']: + """Return a list of available starter clusters.""" + manager = self._require_manager() + res = manager._get(SHAREDTIER_PATH) + return NamedList( + [type(self).from_dict(item, manager) for item in res.json()], + ) + + def create_user( + self, + username: str, + password: Optional[str] = None, + ) -> Dict[str, str]: + """ + Create a new user for this starter cluster. + + Parameters + ---------- + username : str + The user name to connect the new user to the database + password : str, optional + Password for the new user. If not provided, a password will be + auto-generated by the system. + + Returns + ------- + Dict[str, str] + Dictionary containing 'user_id' and 'password' of the created user + + Raises + ------ + ManagementError + If no cluster manager is associated with this object + + """ + manager = self._require_manager() + + payload = {'userName': username} + if password is not None: + payload['password'] = password + + res = manager._post( + f'{SHAREDTIER_PATH}/{self.id}/users', + json=payload, + ) + + response_data = res.json() + user_id = response_data.get('userID') + if not user_id: + raise ManagementError(msg='No userID returned from API') + + # Return the password provided by user or generated by API + returned_password = password if password is not None \ + else response_data.get('password') + if not returned_password: + raise ManagementError(msg='No password available from API response') + + return { + 'user_id': user_id, + 'password': returned_password, + } + + +class ClusterManager(Manager): + """ + SingleStoreDB cluster manager. + + This class should be instantiated using + :func:`singlestoredb.manage_clusters`. + + Parameters + ---------- + access_token : str, optional + The API key or other access token for the cluster management API + version : str, optional + Version of the API to use + base_url : str, optional + Base URL of the cluster management API + + See Also + -------- + :func:`singlestoredb.manage_clusters` + + """ + + #: Cluster management API version if none is specified. A literal, because + #: this class implements the v2 routes; it does not follow whatever the + #: current default version is. + default_version = 'v2' + + #: Base URL if none is specified. + default_base_url = config.get_option('management.base_url') \ + or 'https://api.singlestore.com' + + #: Object type + obj_type = 'cluster' + + @property + def clusters(self) -> NamedList[Cluster]: + """Return a list of available clusters.""" + res = self._get('clusters') + return NamedList([Cluster.from_dict(item, self) for item in res.json()]) + + @property + def starter_clusters(self) -> NamedList[StarterCluster]: + """Return a list of available starter clusters.""" + res = self._get(SHAREDTIER_PATH) + return NamedList( + [StarterCluster.from_dict(item, self) for item in res.json()], + ) + + @property + def organizations(self) -> Organizations: + """Return the organizations.""" + return Organizations(self) + + @property + def organization(self) -> Organization: + """Return the current organization.""" + return self.organizations.current + + @property + def billing(self) -> Billing: + """Return the current billing information.""" + return Billing(self) + + @ttl_property(datetime.timedelta(hours=1)) + def regions(self) -> NamedList[Region]: + """ + Return a list of available regions. + + Cached for the same reason as :attr:`projects`: :attr:`Cluster.region` + resolves against this list, and a caller reading it per row of a + listing -- ``SHOW CLUSTERS`` does -- would otherwise cost a + ``GET /v2/regions`` per cluster. + """ + res = self._get('regions') + return NamedList([Region.from_dict(item, self) for item in res.json()]) + + @ttl_property(datetime.timedelta(hours=1)) + def projects(self) -> NamedList[Project]: + """ + Return a list of projects in the current organization. + + Cached like :attr:`regions`, because :attr:`Cluster.project` resolves + against this list and a caller reading it per row of a listing -- + ``SHOW CLUSTERS EXTENDED`` does -- would otherwise cost a + ``GET /v2/projects`` per cluster. + """ + res = self._get('projects') + return NamedList([Project.from_dict(item, self) for item in res.json()]) + + def get_project(self, id: str) -> Project: + """ + Retrieve a project definition. + + Parameters + ---------- + id : str + ID of the project + + Returns + ------- + :class:`Project` + + """ + res = self._get(f'projects/{id}') + return Project.from_dict(res.json(), manager=self) + + def _wait_on_firewall( + self, + out: Cluster, + interval: int = 10, + timeout: int = 600, + expected: Optional[List[str]] = None, + ) -> Cluster: + """ + Wait until the cluster reports the firewall that was asked for. + + ``POST /v2/clusters`` and ``PATCH /v2/clusters/{id}`` apply + ``firewallRanges`` asynchronously and outside the state machine: the + cluster reaches ACTIVE with a resolvable endpoint while + ``GET /v2/clusters/{id}`` still reports ``firewallRanges: []`` and + ``allowAllTraffic: null``. That combination denies all inbound traffic, + so a connection attempt in that window times out at the TCP level + rather than failing authentication. + + By default the wait is for the cluster to admit *anything* -- either + non-empty ``firewall_ranges`` or ``allow_all_traffic`` -- rather than + for set-equality with the ranges that were requested, because the + server normalizes: ``firewallRanges: ['0.0.0.0/0']`` comes back as + ``allowAllTraffic: True`` with ``firewallRanges: []``. Admitting + something is the property that matters on a fresh cluster -- it is the + difference between deny-all and reachable. + + On an *existing* cluster that already admits traffic, that says + nothing. Pass ``expected`` there to wait for the specific ranges + instead; a requested ``0.0.0.0/0`` is also satisfied by + ``allow_all_traffic``, which is how the server stores it. + + Parameters + ---------- + out : Cluster + Cluster to poll + interval : int, optional + Number of seconds between each server poll + timeout : int, optional + Maximum number of seconds to wait before raising an exception + expected : List[str], optional + Wait for exactly these ranges (compared as a set) rather than for + the firewall to admit anything at all + + Raises + ------ + ManagementError + If timeout is reached + + Returns + ------- + :class:`Cluster` + + """ + def done(cluster: Cluster) -> bool: + if expected is not None: + if set(cluster.firewall_ranges or []) == set(expected): + return True + # The server stores a requested 0.0.0.0/0 as allowAllTraffic + # and leaves firewallRanges empty. + return bool(cluster.allow_all_traffic) \ + and set(expected) == {'0.0.0.0/0'} + return bool(cluster.firewall_ranges) \ + or bool(cluster.allow_all_traffic) + + waited = 0.0 + remaining = float(timeout) + while not done(out): + if remaining <= 0: + wanted = 'to become {}'.format(expected) \ + if expected is not None else 'to be applied' + raise ManagementError( + msg=f'Exceeded waiting time for the firewall of cluster ' + f'{out.id} {wanted} ({waited:.0f}s); it reports ' + f'firewall_ranges={out.firewall_ranges!r}, ' + f'allow_all_traffic={out.allow_all_traffic!r}. While ' + 'the firewall admits nothing the endpoint refuses all ' + 'inbound connections.', + ) + started_at = timing.now() + timing.sleep(interval, 'cluster firewall') + out = self.get_cluster(out.id) + # Measured, and charged after the refetch, so a slow or retried GET + # counts against the timeout. See timing.poll_cost. + cost = timing.poll_cost(started_at, interval) + remaining -= cost + waited += cost + + return out + + def _project_id_for(self, name_or_id: Union[str, Project]) -> str: + """ + Return the ID of the project named by ``name_or_id``. + + A :class:`Project` is reduced to its ID. A UUID is taken as an ID and + returned untouched, which keeps an explicit ID free of a + ``GET /v2/projects`` round trip. Anything else is matched against the + project names in the current organization. The API does not promise + that names are unique, so an ambiguous name raises rather than picking + the first match. + + Parameters + ---------- + name_or_id : str or Project + Project, or project name or ID + + Returns + ------- + str + + Raises + ------ + ManagementError + If the name matches no project, or more than one + + """ + if isinstance(name_or_id, Project): + return name_or_id.id + + if PROJECT_ID_RE.match(name_or_id): + return name_or_id + + projects = self.projects + matches = [x for x in projects if x.name == name_or_id] + + if not matches: + raise ManagementError( + msg=f'No project named {name_or_id!r} exists in the current ' + 'organization. Its projects are: ' + + ( + ', '.join(f'{x.name} ({x.id})' for x in projects) + or 'none' + ) + '.', + ) + + if len(matches) > 1: + raise ManagementError( + msg=f'More than one project is named {name_or_id!r}; use an ID ' + 'instead. The matching IDs are: ' + + ', '.join(x.id for x in matches) + '.', + ) + + return matches[0].id + + def _current_deployment_project_id(self) -> Optional[str]: + """ + Return the project of the deployment this code is running in. + + A notebook publishes the deployment it is attached to as + ``SINGLESTOREDB_WORKSPACE``, and a deployment reports its own + ``projectID``, so the project a new cluster most likely belongs in is + the one the current cluster is already in. + + Returns ``None`` whenever that cannot be established, which covers + running outside a notebook, a deployment that is not a cluster -- a + starter cluster publishes the same variable -- and a stale ID. None of + those are errors here: the caller has further defaults to try. + """ + deployment_id = get_cluster_id() + if not deployment_id: + return None + + try: + project = self.get_cluster(deployment_id).project + except ManagementError: + return None + + return project.id if project is not None else None + + def _resolve_project_id( + self, + project: Union[str, Project, None] = None, + ) -> str: + """ + Return the project ID a new deployment should be created in. + + ``POST /v2/clusters`` requires ``projectID``. In priority order: the + project named by the caller, the project of the deployment this code is + running in, or the organization's only project. An organization with + more than one project and nothing else to go on has no default -- + naming the candidates is more useful than picking one. + + The caller may give a :class:`Project`, a project name or a project ID; + see :meth:`_project_id_for`. + + Note that ``SINGLESTOREDB_PROJECT`` is deliberately not consulted. The + notebook environment sets it, but not to a project of this API: it + names a project of the inference API, a separate namespace whose IDs do + not resolve here. See :func:`singlestoredb.management.utils. + get_project_id`. + + Parameters + ---------- + project : str or Project, optional + Project, or project name or ID, supplied by the caller + + Returns + ------- + str + + Raises + ------ + ManagementError + If no project ID can be determined + + """ + if project: + return self._project_id_for(project) + + from_deployment = self._current_deployment_project_id() + if from_deployment: + return from_deployment + + projects = self.projects + if len(projects) == 1: + return projects[0].id + + if not projects: + raise ManagementError( + msg='A project is required to create a cluster, but the ' + 'current organization reports no projects.', + ) + + raise ManagementError( + msg='A project is required to create a cluster and the current ' + 'organization has more than one. Pass project= naming one ' + 'of: ' + + ', '.join(f'{x.name} ({x.id})' for x in projects) + '.', + ) + + def create_cluster( + self, + name: str, + region: Union[str, Region, None] = None, + provider: Optional[str] = None, + size: Optional[str] = None, + scale_factor: Optional[float] = None, + firewall_ranges: Optional[List[str]] = None, + allow_all_traffic: Optional[bool] = None, + admin_password: Optional[str] = None, + auto_suspend: Optional[Dict[str, Any]] = None, + auto_scale: Optional[Dict[str, Any]] = None, + cache_config: Optional[float] = None, + deployment_type: Optional[str] = None, + expires_at: Optional[str] = None, + update_window: Optional[Dict[str, int]] = None, + kai: Optional[bool] = None, + multi_az: Optional[bool] = None, + opt_in_preview_feature: Optional[bool] = None, + project: Union[str, Project, None] = None, + wait_on_active: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + ) -> Cluster: + """ + Create a new cluster. + + A cluster is created in one call: the firewall, update window and + expiration settings are passed here alongside the compute settings. + + Parameters + ---------- + name : str + Name of the cluster + region : str or Region, optional + Region to create the cluster in. A :class:`Region` supplies both + halves of the ``(provider, region_name)`` pair that identifies a + region; a string is taken as the provider region name, e.g., + ``us-east-1``, and needs ``provider`` alongside it. There are no + region IDs. + provider : str, optional + Cloud provider for the cluster (AWS | GCP | Azure). Only needed + when ``region`` is a string; a :class:`Region` carries its own, + which this overrides if both are given. + size : str, optional + Cluster size in cluster size notation (S-00, S-1, etc.). Sent + nested in a ``size`` object alongside ``scale_factor``. + scale_factor : float, optional + Scale factor for the cluster + firewall_ranges : List[str], optional + List of allowed CIDR ranges. An empty list denies all inbound + traffic, which is also what is sent when this is not given: + ``POST /v2/clusters`` rejects a null ``firewallRanges`` outright, + so there is no way to leave the choice to the server. + allow_all_traffic : bool, optional + Allow all traffic to the cluster + admin_password : str, optional + Admin password for the cluster. + + .. warning:: This is ignored. ``POST /v2/clusters`` generates the + admin password regardless of what is sent and returns the + generated value, so read + :attr:`Cluster.admin_password` off the returned cluster instead + -- it is reported there and nowhere else. The field is still sent + in case the API starts honoring it. + auto_suspend : Dict[str, Any], optional + Auto-suspend settings for the cluster + auto_scale : Dict[str, Any], optional + Auto-scale settings for the cluster + cache_config : float, optional + Multiplier for the persistent cache: 1, 2, or 4 + deployment_type : str, optional + Deployment type of the cluster (PRODUCTION | NON-PRODUCTION) + expires_at : str, optional + Timestamp of when the cluster will expire + update_window : Dict[str, int], optional + Day and hour of an update window: dict(day=0-6, hour=0-23) + kai : bool, optional + Whether to enable SingleStore Kai on this cluster + multi_az : bool, optional + Whether to deploy across multiple availability zones + opt_in_preview_feature : bool, optional + Whether to opt in to preview features + project : str or Project, optional + Project to create the cluster in. A :class:`Project` is reduced to + its ID; a string that is not a UUID is looked up as a name. + Required by the API; if it is not + given it is resolved by :meth:`_resolve_project_id` from the + deployment this code is running in, or from the organization's only + project. + wait_on_active : bool, optional + Wait for the cluster to be usable before returning: first for the + state to become ACTIVE, then for the endpoint, then -- if a + firewall was requested -- for the firewall to be applied. The + firewall is included because the API applies it asynchronously and + outside the state machine, so an ACTIVE cluster with a resolvable + endpoint still refuses every inbound connection until the ranges + land. See :meth:`_wait_on_firewall`. + wait_interval : int, optional + Number of seconds between each polling interval + wait_timeout : int, optional + Maximum number of seconds to wait before raising an exception + + Returns + ------- + :class:`Cluster` + + """ + region_name: Optional[str] = None + if isinstance(region, Region): + provider = provider or region.provider + region_name = region.region_name or region.name + elif region is not None: + region_name = region + + project_id = self._resolve_project_id(project) + + # POST /v2/clusters rejects a null firewallRanges -- "indicate empty + # list [] to disallow all inbound traffic" -- so the field cannot be + # dropped the way every other unset field is. Deny-all is the only + # safe default for a cluster nobody asked to expose. + if firewall_ranges is None: + firewall_ranges = [] + + size_spec: Optional[Dict[str, Any]] = None + if size is not None or scale_factor is not None: + size_spec = { + k: v for k, v in dict( + size=size, scaleFactor=scale_factor, + ).items() if v is not None + } + + res = self._post( + 'clusters', json={ + k: v for k, v in dict( + name=name, + provider=provider, + region=region_name, + # ``sizeConfig``, not ``size``; see Cluster.from_dict. + sizeConfig=size_spec, + firewallRanges=firewall_ranges, + allowAllTraffic=allow_all_traffic, + adminPassword=admin_password, + autoSuspend=snake_to_camel_dict(auto_suspend), + autoScale=snake_to_camel_dict(auto_scale), + cacheConfig=cache_config, + deploymentType=deployment_type, + expiresAt=expires_at, + updateWindow=snake_to_camel_dict(update_window), + kai=kai, + multiAZ=multi_az, + optInPreviewFeature=opt_in_preview_feature, + projectID=project_id, + ).items() if v is not None + }, + ) + body = res.json() + out = self.get_cluster(body['clusterID']) + if wait_on_active: + out = self._wait_on_state( + out, 'ACTIVE', interval=wait_interval, timeout=wait_timeout, + ) + # After the cluster is active, wait for the endpoint to be ready + out = self._wait_on_endpoint( + out, interval=wait_interval, timeout=wait_timeout, + ) + # ...and the endpoint refuses everything until the firewall lands. + # Only when a firewall was actually asked for: firewall_ranges=[] + # is a legitimate deny-all request and must not hang waiting for a + # non-empty value that is never coming. + if firewall_ranges or allow_all_traffic: + out = self._wait_on_firewall( + out, interval=wait_interval, timeout=wait_timeout, + ) + # The API generates the admin password and reports it here and nowhere + # else, and every wait above re-fetches the cluster, so this assignment + # must stay after all of them: carry the password over onto whichever + # object is being returned. See Cluster.admin_password. + out._admin_password = body.get('adminPassword') + return out + + def get_cluster(self, id: str) -> Cluster: + """ + Retrieve a cluster definition. + + Parameters + ---------- + id : str + ID of the cluster + + Returns + ------- + :class:`Cluster` + + """ + res = self._get(f'clusters/{id}') + return Cluster.from_dict(res.json(), manager=self) + + def get_starter_cluster(self, id: str) -> StarterCluster: + """ + Retrieve a starter cluster definition. + + Parameters + ---------- + id : str + ID of the starter cluster + + Returns + ------- + :class:`StarterCluster` + + """ + res = self._get(f'{SHAREDTIER_PATH}/{id}') + return StarterCluster.from_dict(res.json(), manager=self) + + def create_starter_cluster( + self, + name: str, + database_name: str, + provider: Optional[str] = None, + region: Union[str, Region, None] = None, + project: Union[str, Project, None] = None, + ) -> StarterCluster: + """ + Create a new starter (shared tier) cluster. + + Parameters + ---------- + name : str + Name of the starter cluster + database_name : str + Name of the database for the starter cluster + provider : str, optional + Cloud provider for the starter cluster (AWS | GCP | Azure). Any + capitalization is accepted; see below. Only needed when ``region`` + is a string; a :class:`Region` carries its own, which this + overrides if both are given. + region : str or Region + Region to create the starter cluster in. A :class:`Region` supplies + both the provider and the provider region name; a string is taken + as the provider region name (e.g., 'us-east-1') and needs + ``provider`` alongside it. See :attr:`shared_tier_regions` for the + regions this route accepts. + project : str or Project, optional + Project to associate the starter cluster with. A :class:`Project` + is reduced to its ID; a string that is not a UUID is looked up as a + name. Unlike + :meth:`create_cluster` this route does not require one, so nothing + is resolved when it is omitted. + + Returns + ------- + :class:`StarterCluster` + + """ + region_name: Optional[str] = None + if isinstance(region, Region): + provider = provider or region.provider + region_name = region.region_name or region.name + elif region is not None: + region_name = region + + if not provider or not region_name: + raise ValueError( + 'a provider and a region name are required; pass a Region, ' + 'or a provider region name together with provider=', + ) + + payload: Dict[str, Any] = { + 'name': name, + 'databaseName': database_name, + # The shared-tier route accepts only the exact spellings AWS, + # AZURE and GCP: anything else, including the mixed-case 'Azure' + # that GET /v2/regions itself reports, fails with + # '500 Unspecified is not a valid CloudServiceProvider'. + # POST /v2/clusters is case-insensitive, so this is local to here. + 'provider': provider.upper(), + 'regionName': region_name, + } + if project is not None: + payload['projectID'] = self._project_id_for(project) + + res = self._post(SHAREDTIER_PATH, json=payload) + cluster_id = res.json().get('virtualClusterID') + if not cluster_id: + raise ManagementError(msg='No virtualClusterID returned from API') + + return self.get_starter_cluster(cluster_id) + + @ttl_property(datetime.timedelta(hours=1)) + def shared_tier_regions(self) -> NamedList[Region]: + """ + Return a list of regions that support starter clusters. + + Cached for one hour, like :attr:`regions`. + + """ + res = self._get('regions/sharedtier') + return NamedList([Region.from_dict(item, self) for item in res.json()]) diff --git a/singlestoredb/management/v2/export.py b/singlestoredb/management/v2/export.py new file mode 100644 index 000000000..bcfb90fbd --- /dev/null +++ b/singlestoredb/management/v2/export.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python +""" +SingleStoreDB export service (API v2). + +Table egress is driven through ``clusters/{id}/egress/...``, so an export is +owned by a :class:`~singlestoredb.management.v2.cluster.Cluster`. +""" +from __future__ import annotations + +import copy +import json +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 ..utils import vars_to_str +from .cluster import Cluster +from .cluster import ClusterManager + + +class ExportService(object): + """Export service.""" + + database: str + table: str + catalog_info: Dict[str, Any] + storage_info: Dict[str, Any] + columns: Optional[List[str]] + partition_by: Optional[List[Dict[str, str]]] + order_by: Optional[List[Dict[str, Dict[str, str]]]] + properties: Optional[Dict[str, Any]] + incremental: bool + refresh_interval: Optional[int] + export_id: Optional[str] + + def __init__( + self, + cluster: Cluster, + database: str, + table: str, + catalog_info: Union[str, Dict[str, Any]], + storage_info: Union[str, Dict[str, Any]], + columns: Optional[List[str]] = None, + partition_by: Optional[List[Dict[str, str]]] = None, + order_by: Optional[List[Dict[str, Dict[str, str]]]] = None, + incremental: bool = False, + refresh_interval: Optional[int] = None, + properties: Optional[Dict[str, Any]] = None, + ): + #: Cluster the export runs against + self.cluster = cluster + + #: Name of SingleStoreDB database + self.database = database + + #: Name of SingleStoreDB table + self.table = table + + #: List of columns to export + self.columns = columns + + #: Catalog + if isinstance(catalog_info, str): + self.catalog_info = json.loads(catalog_info) + else: + self.catalog_info = copy.copy(catalog_info) + + #: Storage + if isinstance(storage_info, str): + self.storage_info = json.loads(storage_info) + else: + self.storage_info = copy.copy(storage_info) + + self.partition_by = partition_by or None + self.order_by = order_by or None + self.properties = properties or None + + self.incremental = incremental + self.refresh_interval = refresh_interval + + self.export_id = None + + self._manager: Optional[ClusterManager] = cluster._manager + + @classmethod + def from_export_id( + cls, + cluster: Cluster, + export_id: str, + ) -> ExportService: + """Create export service from export ID.""" + out = cls( + cluster=cluster, + database='', + table='', + catalog_info={}, + storage_info={}, + ) + out.export_id = export_id + return out + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + def _require_manager(self) -> ClusterManager: + if self._manager is None: + raise ManagementError( + msg='No cluster manager is associated with this object.', + ) + return self._manager + + def _require_export_id(self) -> str: + if self.export_id is None: + raise ManagementError( + msg='Export ID is not set. You must start the export first.', + ) + return self.export_id + + def _egress_path(self, verb: str) -> str: + return f'clusters/{self.cluster.id}/egress/{verb}' + + def create_cluster_identity(self) -> Dict[str, Any]: + """Create a cluster identity.""" + out = self._require_manager()._post( + self._egress_path('createEgressClusterIdentity'), + json=dict( + catalogInfo=self.catalog_info, + storageInfo=self.storage_info, + ), + ) + return out.json() + + def start(self, tags: Optional[List[str]] = None) -> 'ExportStatus': + """Start the export process.""" + if not self.table or not self.database: + raise ManagementError( + msg='Database and table must be set before starting the export.', + ) + + manager = self._require_manager() + + partition_spec = None + if self.partition_by: + partition_spec = dict(partitions=self.partition_by) + + sort_order_spec = None + if self.order_by: + sort_order_spec = dict(keys=self.order_by) + + out = manager._post( + self._egress_path('startTableEgress'), + json={ + k: v for k, v in dict( + databaseName=self.database, + tableName=self.table, + storageInfo=self.storage_info, + catalogInfo=self.catalog_info, + partitionSpec=partition_spec, + sortOrderSpec=sort_order_spec, + properties=self.properties, + incremental=self.incremental or None, + refreshInterval=self.refresh_interval + if self.refresh_interval is not None else None, + ).items() if v is not None + }, + ) + + self.export_id = str(out.json()['egressID']) + + return ExportStatus(self.export_id, self.cluster) + + def suspend(self) -> 'ExportStatus': + """Suspend the export process.""" + manager = self._require_manager() + export_id = self._require_export_id() + manager._post( + self._egress_path('suspendTableEgress'), + json=dict(egressID=export_id), + ) + return ExportStatus(export_id, self.cluster) + + def resume(self) -> 'ExportStatus': + """Resume the export process.""" + manager = self._require_manager() + export_id = self._require_export_id() + manager._post( + self._egress_path('resumeTableEgress'), + json=dict(egressID=export_id), + ) + return ExportStatus(export_id, self.cluster) + + def drop(self) -> None: + """Drop the export process.""" + manager = self._require_manager() + export_id = self._require_export_id() + manager._delete( + self._egress_path('dropTableEgress'), + json=dict(egressID=export_id), + ) + return None + + def status(self) -> ExportStatus: + """Get the status of the export process.""" + self._require_manager() + return ExportStatus(self._require_export_id(), self.cluster) + + +class ExportStatus(object): + """Status of a table egress process.""" + + export_id: str + + def __init__(self, export_id: str, cluster: Cluster): + self.export_id = export_id + self.cluster = cluster + self._manager: Optional[ClusterManager] = cluster._manager + + def _info(self) -> Dict[str, Any]: + """Return export status.""" + if self._manager is None: + raise ManagementError( + msg='No cluster manager is associated with this object.', + ) + + out = self._manager._get( + f'clusters/{self.cluster.id}/egress/tableEgressStatus', + json=dict(egressID=self.export_id), + ) + + return out.json() + + @property + def status(self) -> str: + """Return export status.""" + return self._info().get('status', 'Unknown') + + @property + def message(self) -> str: + """Return export status message.""" + return self._info().get('statusMsg', '') + + def __str__(self) -> str: + return self.status + + def __repr__(self) -> str: + return self.status + + +def _get_exports( + cluster: Cluster, + scope: str = 'all', +) -> List[ExportStatus]: + """Get all exports in the cluster.""" + if cluster._manager is None: + raise ManagementError( + msg='No cluster manager is associated with this object.', + ) + + out = cluster._manager._get( + f'clusters/{cluster.id}/egress/tableEgressStatus', + json=dict(scope=scope), + ) + + return [ + ExportStatus(item['egressID'], cluster) + for item in out.json() + ] diff --git a/singlestoredb/management/v2/files.py b/singlestoredb/management/v2/files.py new file mode 100644 index 000000000..01d8e0d60 --- /dev/null +++ b/singlestoredb/management/v2/files.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python +""" +SingleStoreDB Files Management API v2. + +The ``files/fs/{space}/...`` routes are implemented in +:mod:`singlestoredb.management.files`, so this module only re-exports it. +""" +from ..files import FileLocation as FileLocation +from ..files import FilesManager as FilesManager +from ..files import FilesObject as FilesObject +from ..files import FilesObjectBytesReader as FilesObjectBytesReader +from ..files import FilesObjectBytesWriter as FilesObjectBytesWriter +from ..files import FilesObjectTextReader as FilesObjectTextReader +from ..files import FilesObjectTextWriter as FilesObjectTextWriter +from ..files import FileSpace as FileSpace +from ..files import MODELS_SPACE as MODELS_SPACE +from ..files import PERSONAL_SPACE as PERSONAL_SPACE +from ..files import SHARED_SPACE as SHARED_SPACE diff --git a/singlestoredb/management/v2/job.py b/singlestoredb/management/v2/job.py new file mode 100644 index 000000000..5e92b3384 --- /dev/null +++ b/singlestoredb/management/v2/job.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python +""" +SingleStoreDB Job Management API v2. + +The jobs routes and their ``targetConfig.targetType`` vocabulary are +implemented in :mod:`singlestoredb.management.job`, so this module only +re-exports it. +""" +from ..job import Execution as Execution +from ..job import ExecutionConfig as ExecutionConfig +from ..job import ExecutionMetadata as ExecutionMetadata +from ..job import ExecutionsData as ExecutionsData +from ..job import Job as Job +from ..job import JobMetadata as JobMetadata +from ..job import JobsManager as JobsManager +from ..job import Mode as Mode +from ..job import Parameter as Parameter +from ..job import Runtime as Runtime +from ..job import Schedule as Schedule +from ..job import Status as Status +from ..job import TargetConfig as TargetConfig +from ..job import TargetType as TargetType diff --git a/singlestoredb/management/v2/organization.py b/singlestoredb/management/v2/organization.py new file mode 100644 index 000000000..98a114149 --- /dev/null +++ b/singlestoredb/management/v2/organization.py @@ -0,0 +1,10 @@ +#!/usr/bin/env python +""" +SingleStoreDB Organization API v2. + +``GET /v2/organizations/current`` and ``GET /v2/secrets`` are implemented in +:mod:`singlestoredb.management.organization`, so this module only re-exports it. +""" +from ..organization import Organization as Organization +from ..organization import Organizations as Organizations +from ..organization import Secret as Secret diff --git a/singlestoredb/management/v2/project.py b/singlestoredb/management/v2/project.py new file mode 100644 index 000000000..10814911e --- /dev/null +++ b/singlestoredb/management/v2/project.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python +""" +SingleStoreDB Project API v2. + +``GET /v2/projects`` lists the projects in the current organization. A project +ID is required to create a cluster: ``POST /v2/clusters`` rejects a body without +``projectID`` (``400 projectID is required``). +""" +from __future__ import annotations + +import datetime +from typing import Any +from typing import Dict +from typing import Optional +from typing import Union + +from ..manager import Manager +from ..utils import to_datetime +from ..utils import vars_to_str + + +class Project: + """ + Project definition. + + This object is not directly instantiated. It is used in results of + ``ClusterManager`` API calls. + + See Also + -------- + :attr:`ClusterManager.projects` + + """ + + def __init__( + self, + id: str, + name: str, + edition: Optional[str] = None, + created_at: Optional[Union[str, datetime.datetime]] = None, + ) -> None: + """Use :attr:`ClusterManager.projects` instead.""" + #: Unique ID of the project + self.id = id + + #: Name of the project + self.name = name + + #: Edition of the project (SHARED | STANDARD | ENTERPRISE) + self.edition = edition + + #: Timestamp of when the project was created + self.created_at = to_datetime(created_at) + + self._manager: Optional[Manager] = 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: Manager) -> 'Project': + """ + Convert dictionary to a ``Project`` object. + + Parameters + ---------- + obj : dict + Key-value pairs to retrieve project information from + manager : ClusterManager + The ClusterManager the Project belongs to + + Returns + ------- + :class:`Project` + + """ + out = cls( + id=obj['projectID'], + name=obj['name'], + edition=obj.get('edition'), + created_at=obj.get('createdAt'), + ) + out._manager = manager + return out diff --git a/singlestoredb/management/v2/region.py b/singlestoredb/management/v2/region.py new file mode 100644 index 000000000..7a7c6c036 --- /dev/null +++ b/singlestoredb/management/v2/region.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python +""" +SingleStoreDB Region Management API v2. + +``GET /v2/regions`` returns entries containing ``provider``, ``region``, and +``regionName`` only -- no ``regionID``. :class:`Region` instances therefore +have ``id is None`` and ``region_name`` set, and a region is identified by +``(provider, region_name)``. That is what +:mod:`singlestoredb.management.region` implements, so this module only +re-exports it. +""" +from ..region import Region as Region +from ..region import RegionManager as RegionManager diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index 1b5d7c278..571a025a2 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -1,1962 +1,74 @@ #!/usr/bin/env python -"""SingleStoreDB Workspace Management.""" -from __future__ import annotations - -import datetime -import glob -import io -import os -import re -import time -from collections.abc import Mapping -from typing import Any -from typing import cast -from typing import Dict -from typing import List -from typing import Literal +""" +SingleStoreDB Workspace Management (management API v1) -- **deprecated**. + +.. deprecated:: + Every name below comes from :mod:`singlestoredb.management.v1.workspace`. + Workspaces and workspace groups are the management API v1 deployment + vocabulary; v2 replaced both with the flat + :class:`~singlestoredb.management.cluster.Cluster`. Use + :mod:`singlestoredb.management.cluster` and + :func:`singlestoredb.manage_clusters` instead. + + Deprecated, not removed: every name here still works against the live v1 + endpoints, and :func:`manage_workspaces` still hands back a working manager + without being asked for a version. Only the eventual removal of + :mod:`singlestoredb.management.v1` takes it away, and that has not happened. +""" +import warnings from typing import Optional -from typing import overload -from typing import Union - -from .. import config -from .. import connection -from ..exceptions import ManagementError -from .billing_usage import BillingUsageItem -from .files import FileLocation -from .files import FilesObject -from .files import FilesObjectBytesReader -from .files import FilesObjectBytesWriter -from .files import FilesObjectTextReader -from .files import FilesObjectTextWriter -from .manager import Manager -from .organization import Organization -from .region import Region -from .utils import camel_to_snake_dict -from .utils import from_datetime -from .utils import NamedList -from .utils import PathLike -from .utils import snake_to_camel -from .utils import snake_to_camel_dict -from .utils import to_datetime -from .utils import ttl_property -from .utils import vars_to_str - - -def get_organization() -> Organization: - """Get the organization.""" - return manage_workspaces().organization - - -def get_secret(name: str) -> Optional[str]: - """Get a secret from the organization.""" - return get_organization().get_secret(name).value - - -def get_workspace_group( - workspace_group: Optional[Union[WorkspaceGroup, str]] = None, -) -> WorkspaceGroup: - """Get the stage for the workspace group.""" - if isinstance(workspace_group, WorkspaceGroup): - return workspace_group - elif workspace_group: - return manage_workspaces().workspace_groups[workspace_group] - elif 'SINGLESTOREDB_WORKSPACE_GROUP' in os.environ: - return manage_workspaces().workspace_groups[ - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] - ] - raise RuntimeError('no workspace group specified') - - -def get_stage( - workspace_group: Optional[Union[WorkspaceGroup, str]] = None, -) -> Stage: - """Get the stage for the workspace group.""" - return get_workspace_group(workspace_group).stage - - -def get_workspace( - workspace_group: Optional[Union[WorkspaceGroup, str]] = None, - workspace: Optional[Union[Workspace, str]] = None, -) -> Workspace: - """Get the workspaces for a workspace_group.""" - if isinstance(workspace, Workspace): - return workspace - wg = get_workspace_group(workspace_group) - if workspace: - return wg.workspaces[workspace] - elif 'SINGLESTOREDB_WORKSPACE' in os.environ: - return wg.workspaces[ - os.environ['SINGLESTOREDB_WORKSPACE'] - ] - raise RuntimeError('no workspace group specified') - - -class Stage(FileLocation): - """ - Stage manager. - - This object is not instantiated directly. - It is returned by ``WorkspaceGroup.stage`` or ``StarterWorkspace.stage``. - - """ - - def __init__(self, deployment_id: str, manager: WorkspaceManager): - self._deployment_id = deployment_id - self._manager = manager - - def open( - self, - stage_path: PathLike, - mode: str = 'r', - encoding: Optional[str] = None, - ) -> Union[io.StringIO, io.BytesIO]: - """ - Open a Stage path for reading or writing. - - Parameters - ---------- - stage_path : Path or str - The stage path to read / write - mode : str, optional - The read / write mode. The following modes are supported: - * 'r' open for reading (default) - * 'w' open for writing, truncating the file first - * 'x' create a new file and open it for writing - The data type can be specified by adding one of the following: - * 'b' binary mode - * 't' text mode (default) - encoding : str, optional - The string encoding to use for text - - Returns - ------- - FilesObjectBytesReader - 'rb' or 'b' mode - FilesObjectBytesWriter - 'wb' or 'xb' mode - FilesObjectTextReader - 'r' or 'rt' mode - FilesObjectTextWriter - 'w', 'x', 'wt' or 'xt' mode - - """ - if '+' in mode or 'a' in mode: - raise ValueError('modifying an existing stage file is not supported') - - if 'w' in mode or 'x' in mode: - exists = self.exists(stage_path) - if exists: - if 'x' in mode: - raise FileExistsError(f'stage path already exists: {stage_path}') - self.remove(stage_path) - if 'b' in mode: - return FilesObjectBytesWriter(b'', self, stage_path) - return FilesObjectTextWriter('', self, stage_path) - - if 'r' in mode: - content = self.download_file(stage_path) - if isinstance(content, bytes): - if 'b' in mode: - return FilesObjectBytesReader(content) - encoding = 'utf-8' if encoding is None else encoding - return FilesObjectTextReader(content.decode(encoding)) - - if isinstance(content, str): - return FilesObjectTextReader(content) - - raise ValueError(f'unrecognized file content type: {type(content)}') - - raise ValueError(f'must have one of create/read/write mode specified: {mode}') - - def upload_file( - self, - local_path: Union[PathLike, io.IOBase], - stage_path: PathLike, - *, - overwrite: bool = False, - ) -> FilesObject: - """ - Upload a local file. - - Parameters - ---------- - local_path : Path or str or file-like - Path to the local file or an open file object - stage_path : Path or str - Path to the stage file - overwrite : bool, optional - Should the ``stage_path`` be overwritten if it exists already? - - """ - if isinstance(local_path, io.IOBase): - pass - elif not os.path.isfile(local_path): - raise IsADirectoryError(f'local path is not a file: {local_path}') - - if self.exists(stage_path): - if not overwrite: - raise OSError(f'stage path already exists: {stage_path}') - - self.remove(stage_path) - - if isinstance(local_path, io.IOBase): - return self._upload(local_path, stage_path, overwrite=overwrite) - - return self._upload(open(local_path, 'rb'), stage_path, overwrite=overwrite) - - def upload_folder( - self, - local_path: PathLike, - stage_path: PathLike, - *, - overwrite: bool = False, - recursive: bool = True, - include_root: bool = False, - ignore: Optional[Union[PathLike, List[PathLike]]] = None, - ) -> FilesObject: - """ - Upload a folder recursively. - - Only the contents of the folder are uploaded. To include the - folder name itself in the target path use ``include_root=True``. - - Parameters - ---------- - local_path : Path or str - Local directory to upload - stage_path : Path or str - Path of stage folder to upload to - overwrite : bool, optional - If a file already exists, should it be overwritten? - recursive : bool, optional - Should nested folders be uploaded? - include_root : bool, optional - Should the local root folder itself be uploaded as the top folder? - ignore : Path or str or List[Path] or List[str], optional - Glob patterns of files to ignore, for example, ``**/*.pyc`` will - ignore all ``*.pyc`` files in the directory tree - - """ - if not os.path.isdir(local_path): - raise NotADirectoryError(f'local path is not a directory: {local_path}') - if self.exists(stage_path) and not self.is_dir(stage_path): - raise NotADirectoryError(f'stage path is not a directory: {stage_path}') - - ignore_files = set() - if ignore: - if isinstance(ignore, list): - for item in ignore: - ignore_files.update(glob.glob(str(item), recursive=recursive)) - else: - ignore_files.update(glob.glob(str(ignore), recursive=recursive)) - - parent_dir = os.path.basename(os.getcwd()) - - files = glob.glob(os.path.join(local_path, '**'), recursive=recursive) - - for src in files: - if ignore_files and src in ignore_files: - continue - target = os.path.join(parent_dir, src) if include_root else src - self.upload_file(src, target, overwrite=overwrite) - - return self.info(stage_path) - - def _upload( - self, - content: Union[str, bytes, io.IOBase], - stage_path: PathLike, - *, - overwrite: bool = False, - ) -> FilesObject: - """ - Upload content to a stage file. - - Parameters - ---------- - content : str or bytes or file-like - Content to upload to stage - stage_path : Path or str - Path to the stage file - overwrite : bool, optional - Should the ``stage_path`` be overwritten if it exists already? - - """ - if self.exists(stage_path): - if not overwrite: - raise OSError(f'stage path already exists: {stage_path}') - self.remove(stage_path) - - self._manager._put( - f'stage/{self._deployment_id}/fs/{stage_path}', - files={'file': content}, - headers={'Content-Type': None}, - ) - - return self.info(stage_path) - - def mkdir(self, stage_path: PathLike, overwrite: bool = False) -> FilesObject: - """ - Make a directory in the stage. - - Parameters - ---------- - stage_path : Path or str - Path of the folder to create - overwrite : bool, optional - Should the stage path be overwritten if it exists already? - - Returns - ------- - FilesObject - - """ - stage_path = re.sub(r'/*$', r'', str(stage_path)) + '/' - - if self.exists(stage_path): - if not overwrite: - return self.info(stage_path) - - self.remove(stage_path) - - self._manager._put( - f'stage/{self._deployment_id}/fs/{stage_path}?isFile=false', - ) - - return self.info(stage_path) - - mkdirs = mkdir - - def rename( - self, - old_path: PathLike, - new_path: PathLike, - *, - overwrite: bool = False, - ) -> FilesObject: - """ - Move the stage file to a new location. - - Paraemeters - ----------- - old_path : Path or str - Original location of the path - new_path : Path or str - New location of the path - overwrite : bool, optional - Should the ``new_path`` be overwritten if it exists already? - - """ - if not self.exists(old_path): - raise OSError(f'stage path does not exist: {old_path}') - - if self.exists(new_path): - if not overwrite: - raise OSError(f'stage path already exists: {new_path}') - - if str(old_path).endswith('/') and not str(new_path).endswith('/'): - raise OSError('original and new paths are not the same type') - - if str(new_path).endswith('/'): - self.removedirs(new_path) - else: - self.remove(new_path) - - self._manager._patch( - f'stage/{self._deployment_id}/fs/{old_path}', - json=dict(newPath=new_path), - ) - - return self.info(new_path) - - def info(self, stage_path: PathLike) -> FilesObject: - """ - Return information about a stage location. - - Parameters - ---------- - stage_path : Path or str - Path to the stage location - - Returns - ------- - FilesObject - - """ - res = self._manager._get( - re.sub(r'/+$', r'/', f'stage/{self._deployment_id}/fs/{stage_path}'), - params=dict(metadata=1), - ).json() - - return FilesObject.from_dict(res, self) - - def exists(self, stage_path: PathLike) -> bool: - """ - Does the given stage path exist? - - Parameters - ---------- - stage_path : Path or str - Path to stage object - - Returns - ------- - bool - - """ - try: - self.info(stage_path) - return True - except ManagementError as exc: - if exc.errno == 404: - return False - raise - - def is_dir(self, stage_path: PathLike) -> bool: - """ - Is the given stage path a directory? - - Parameters - ---------- - stage_path : Path or str - Path to stage object - - Returns - ------- - bool - - """ - try: - return self.info(stage_path).type == 'directory' - except ManagementError as exc: - if exc.errno == 404: - return False - raise - - def is_file(self, stage_path: PathLike) -> bool: - """ - Is the given stage path a file? - - Parameters - ---------- - stage_path : Path or str - Path to stage object - - Returns - ------- - bool - - """ - try: - return self.info(stage_path).type != 'directory' - except ManagementError as exc: - if exc.errno == 404: - return False - raise - - def _listdir( - self, stage_path: PathLike, *, - recursive: bool = False, - return_objects: bool = False, - ) -> List[Union[str, 'FilesObject']]: - """ - Return the names (or FilesObject instances) of files in a directory. - - Parameters - ---------- - stage_path : Path or str - Path to the folder in Stage - recursive : bool, optional - Should folders be listed recursively? - return_objects : bool, optional - If True, return list of FilesObject instances. Otherwise just paths. - - """ - from .files import FilesObject - res = self._manager._get( - re.sub(r'/+$', r'/', f'stage/{self._deployment_id}/fs/{stage_path}'), - ).json() - if recursive: - out: List[Union[str, FilesObject]] = [] - for item in res['content'] or []: - if return_objects: - out.append(FilesObject.from_dict(item, self)) - else: - out.append(item['path']) - if item['type'] == 'directory': - out.extend( - self._listdir( - item['path'], - recursive=recursive, - return_objects=return_objects, - ), - ) - return out - if return_objects: - return [ - FilesObject.from_dict(x, self) - for x in res['content'] or [] - ] - return [x['path'] for x in res['content'] or []] - - @overload - def listdir( - self, - stage_path: PathLike = '/', - *, - recursive: bool = False, - return_objects: Literal[True], - ) -> List['FilesObject']: - ... - - @overload - def listdir( - self, - stage_path: PathLike = '/', - *, - recursive: bool = False, - return_objects: Literal[False] = False, - ) -> List[str]: - ... - - def listdir( - self, - stage_path: PathLike = '/', - *, - recursive: bool = False, - return_objects: bool = False, - ) -> Union[List[str], List['FilesObject']]: - """ - List the files / folders at the given path. - - Parameters - ---------- - stage_path : Path or str, optional - Path to the stage location - recursive : bool, optional - If True, recursively list all files and folders - return_objects : bool, optional - If True, return list of FilesObject instances. Otherwise just paths. - - Returns - ------- - List[str] or List[FilesObject] - - """ - from .files import FilesObject - stage_path = re.sub(r'^(\./|/)+', r'', str(stage_path)) - stage_path = re.sub(r'/+$', r'', stage_path) + '/' - - if self.is_dir(stage_path): - out = self._listdir( - stage_path, - recursive=recursive, - return_objects=return_objects, - ) - if stage_path != '/': - stage_path_n = len(stage_path.split('/')) - 1 - if return_objects: - result: List[FilesObject] = [] - for item in out: - if isinstance(item, FilesObject): - rel = '/'.join(item.path.split('/')[stage_path_n:]) - item.path = rel - result.append(item) - return result - out = ['/'.join(str(x).split('/')[stage_path_n:]) for x in out] - if return_objects: - return cast(List[FilesObject], out) - return cast(List[str], out) - - raise NotADirectoryError(f'stage path is not a directory: {stage_path}') - - def download_file( - self, - stage_path: PathLike, - local_path: Optional[PathLike] = None, - *, - overwrite: bool = False, - encoding: Optional[str] = None, - ) -> Optional[Union[bytes, str]]: - """ - Download the content of a stage path. - - Parameters - ---------- - stage_path : Path or str - Path to the stage file - local_path : Path or str - Path to local file target location - overwrite : bool, optional - Should an existing file be overwritten if it exists? - encoding : str, optional - Encoding used to convert the resulting data - - Returns - ------- - bytes or str - ``local_path`` is None - None - ``local_path`` is a Path or str - - """ - if local_path is not None and not overwrite and os.path.exists(local_path): - raise OSError('target file already exists; use overwrite=True to replace') - if self.is_dir(stage_path): - raise IsADirectoryError(f'stage path is a directory: {stage_path}') - - out = self._manager._get( - f'stage/{self._deployment_id}/fs/{stage_path}', - ).content - - if local_path is not None: - with open(local_path, 'wb') as outfile: - outfile.write(out) - return None - - if encoding: - return out.decode(encoding) - - return out - - def download_folder( - self, - stage_path: PathLike, - local_path: PathLike = '.', - *, - overwrite: bool = False, - ) -> None: - """ - Download a Stage folder to a local directory. - - Parameters - ---------- - stage_path : Path or str - Path to the stage file - local_path : Path or str - Path to local directory target location - overwrite : bool, optional - Should an existing directory / files be overwritten if they exist? - - """ - if local_path is not None and not overwrite and os.path.exists(local_path): - raise OSError( - 'target directory already exists; ' - 'use overwrite=True to replace', - ) - if not self.is_dir(stage_path): - raise NotADirectoryError(f'stage path is not a directory: {stage_path}') - - for f in self.listdir(stage_path, recursive=True, return_objects=False): - if self.is_dir(f): - continue - target = os.path.normpath(os.path.join(local_path, f)) - os.makedirs(os.path.dirname(target), exist_ok=True) - self.download_file(f, target, overwrite=overwrite) - - def remove(self, stage_path: PathLike) -> None: - """ - Delete a stage location. - - Parameters - ---------- - stage_path : Path or str - Path to the stage location - - """ - if self.is_dir(stage_path): - raise IsADirectoryError( - 'stage path is a directory, ' - f'use rmdir or removedirs: {stage_path}', - ) - - self._manager._delete(f'stage/{self._deployment_id}/fs/{stage_path}') - - def removedirs(self, stage_path: PathLike) -> None: - """ - Delete a stage folder recursively. - - Parameters - ---------- - stage_path : Path or str - Path to the stage location - - """ - stage_path = re.sub(r'/*$', r'', str(stage_path)) + '/' - self._manager._delete(f'stage/{self._deployment_id}/fs/{stage_path}') - - def rmdir(self, stage_path: PathLike) -> None: - """ - Delete a stage folder. - - Parameters - ---------- - stage_path : Path or str - Path to the stage location - - """ - stage_path = re.sub(r'/*$', r'', str(stage_path)) + '/' - - if self.listdir(stage_path): - raise OSError(f'stage folder is not empty, use removedirs: {stage_path}') - - self._manager._delete(f'stage/{self._deployment_id}/fs/{stage_path}') - - def __str__(self) -> str: - """Return string representation.""" - return vars_to_str(self) - - def __repr__(self) -> str: - """Return string representation.""" - return str(self) - - -StageObject = FilesObject # alias for backward compatibility - - -class Workspace(object): - """ - SingleStoreDB workspace definition. - - This object is not instantiated directly. It is used in the results - of API calls on the :class:`WorkspaceManager`. Workspaces are created using - :meth:`WorkspaceManager.create_workspace`, or existing workspaces are - accessed by either :attr:`WorkspaceManager.workspaces` or by calling - :meth:`WorkspaceManager.get_workspace`. - - See Also - -------- - :meth:`WorkspaceManager.create_workspace` - :meth:`WorkspaceManager.get_workspace` - :attr:`WorkspaceManager.workspaces` - - """ - - name: str - id: str - group_id: str - size: str - state: str - created_at: Optional[datetime.datetime] - terminated_at: Optional[datetime.datetime] - endpoint: Optional[str] - auto_suspend: Optional[Dict[str, Any]] - cache_config: Optional[int] - deployment_type: Optional[str] - resume_attachments: Optional[List[Dict[str, Any]]] - scaling_progress: Optional[int] - last_resumed_at: Optional[datetime.datetime] - - def __init__( - self, - name: str, - workspace_id: str, - workspace_group: Union[str, 'WorkspaceGroup'], - size: str, - state: str, - created_at: Union[str, datetime.datetime], - terminated_at: Optional[Union[str, datetime.datetime]] = None, - endpoint: Optional[str] = None, - auto_suspend: Optional[Dict[str, Any]] = None, - cache_config: Optional[int] = None, - deployment_type: Optional[str] = None, - resume_attachments: Optional[List[Dict[str, Any]]] = None, - scaling_progress: Optional[int] = None, - last_resumed_at: Optional[Union[str, datetime.datetime]] = None, - ): - #: Name of the workspace - self.name = name - - #: Unique ID of the workspace - self.id = workspace_id - - #: Unique ID of the workspace group - if isinstance(workspace_group, WorkspaceGroup): - self.group_id = workspace_group.id - else: - self.group_id = workspace_group - - #: Size of the workspace in workspace size notation (S-00, S-1, etc.) - self.size = size - - #: State of the workspace: PendingCreation, Transitioning, Active, - #: Terminated, Suspended, Resuming, Failed - self.state = state.strip() - - #: Timestamp of when the workspace was created - self.created_at = to_datetime(created_at) - - #: Timestamp of when the workspace was terminated - self.terminated_at = to_datetime(terminated_at) - - #: Hostname (or IP address) of the workspace database server - self.endpoint = endpoint - - #: Current auto-suspend settings - self.auto_suspend = camel_to_snake_dict(auto_suspend) - - #: Multiplier for the persistent cache - self.cache_config = cache_config - - #: Deployment type of the workspace - self.deployment_type = deployment_type - - #: Database attachments - self.resume_attachments = [ - camel_to_snake_dict(x) # type: ignore - for x in resume_attachments or [] - if x is not None - ] - - #: Current progress percentage for scaling the workspace - self.scaling_progress = scaling_progress - - #: Timestamp when workspace was last resumed - self.last_resumed_at = to_datetime(last_resumed_at) - - self._manager: Optional[WorkspaceManager] = 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: 'WorkspaceManager') -> 'Workspace': - """ - Construct a Workspace from a dictionary of values. - - Parameters - ---------- - obj : dict - Dictionary of values - manager : WorkspaceManager, optional - The WorkspaceManager the Workspace belongs to - - Returns - ------- - :class:`Workspace` - - """ - out = cls( - name=obj['name'], - workspace_id=obj['workspaceID'], - workspace_group=obj['workspaceGroupID'], - size=obj.get('size', 'Unknown'), - state=obj['state'], - created_at=obj['createdAt'], - terminated_at=obj.get('terminatedAt'), - endpoint=obj.get('endpoint'), - auto_suspend=obj.get('autoSuspend'), - cache_config=obj.get('cacheConfig'), - deployment_type=obj.get('deploymentType'), - last_resumed_at=obj.get('lastResumedAt'), - resume_attachments=obj.get('resumeAttachments'), - scaling_progress=obj.get('scalingProgress'), - ) - out._manager = manager - return out - - def update( - self, - auto_suspend: Optional[Dict[str, Any]] = None, - cache_config: Optional[int] = None, - deployment_type: Optional[str] = None, - size: Optional[str] = None, - ) -> None: - """ - Update the workspace definition. - - Parameters - ---------- - auto_suspend : Dict[str, Any], optional - Auto-suspend mode for the workspace: IDLE, SCHEDULED, DISABLED - cache_config : int, optional - Specifies the multiplier for the persistent cache associated - with the workspace. If specified, it enables the cache configuration - multiplier. It can have one of the following values: 1, 2, or 4. - deployment_type : str, optional - The deployment type that will be applied to all the workspaces - within the group - size : str, optional - Size of the workspace (in workspace size notation), such as "S-1". - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - data = { - k: v for k, v in dict( - autoSuspend=snake_to_camel_dict(auto_suspend), - cacheConfig=cache_config, - deploymentType=deployment_type, - size=size, - ).items() if v is not None - } - self._manager._patch(f'workspaces/{self.id}', json=data) - self.refresh() - - def refresh(self) -> Workspace: - """Update the object to the current state.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - new_obj = self._manager.get_workspace(self.id) - for name, value in vars(new_obj).items(): - if isinstance(value, Mapping): - setattr(self, name, snake_to_camel_dict(value)) - else: - setattr(self, name, value) - return self - - def terminate( - self, - wait_on_terminated: bool = False, - wait_interval: int = 10, - wait_timeout: int = 600, - force: bool = False, - ) -> None: - """ - Terminate the workspace. - - Parameters - ---------- - wait_on_terminated : bool, optional - Wait for the workspace to go into 'Terminated' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - force : bool, optional - Should the workspace group be terminated even if it has workspaces? - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - force_str = 'true' if force else 'false' - self._manager._delete(f'workspaces/{self.id}?force={force_str}') - if wait_on_terminated: - self._manager._wait_on_state( - self._manager.get_workspace(self.id), - 'Terminated', interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() - - def connect(self, **kwargs: Any) -> connection.Connection: - """ - Create a connection to the database server for this workspace. - - Parameters - ---------- - **kwargs : keyword-arguments, optional - Parameters to the SingleStoreDB `connect` function except host - and port which are supplied by the workspace object - - Returns - ------- - :class:`Connection` - - """ - if not self.endpoint: - raise ManagementError( - msg='An endpoint has not been set in this workspace configuration', - ) - kwargs['host'] = self.endpoint - return connection.connect(**kwargs) - - def suspend( - self, - wait_on_suspended: bool = False, - wait_interval: int = 20, - wait_timeout: int = 600, - ) -> None: - """ - Suspend the workspace. - Parameters - ---------- - wait_on_suspended : bool, optional - Wait for the workspace to go into 'Suspended' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - self._manager._post(f'workspaces/{self.id}/suspend') - if wait_on_suspended: - self._manager._wait_on_state( - self._manager.get_workspace(self.id), - 'Suspended', interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() - - def resume( - self, - disable_auto_suspend: bool = False, - wait_on_resumed: bool = False, - wait_interval: int = 20, - wait_timeout: int = 600, - ) -> None: - """ - Resume the workspace. - - Parameters - ---------- - disable_auto_suspend : bool, optional - Should auto-suspend be disabled? - wait_on_resumed : bool, optional - Wait for the workspace to go into 'Resumed' or 'Active' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - self._manager._post( - f'workspaces/{self.id}/resume', - json=dict(disableAutoSuspend=disable_auto_suspend), - ) - if wait_on_resumed: - self._manager._wait_on_state( - self._manager.get_workspace(self.id), - ['Resumed', 'Active'], interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() - - -class WorkspaceGroup(object): - """ - SingleStoreDB workspace group definition. - - This object is not instantiated directly. It is used in the results - of API calls on the :class:`WorkspaceManager`. Workspace groups are created using - :meth:`WorkspaceManager.create_workspace_group`, or existing workspace groups are - accessed by either :attr:`WorkspaceManager.workspace_groups` or by calling - :meth:`WorkspaceManager.get_workspace_group`. - - See Also - -------- - :meth:`WorkspaceManager.create_workspace_group` - :meth:`WorkspaceManager.get_workspace_group` - :attr:`WorkspaceManager.workspace_groups` - - """ - - name: str - id: str - created_at: Optional[datetime.datetime] - region: Optional[Region] - firewall_ranges: List[str] - terminated_at: Optional[datetime.datetime] - allow_all_traffic: bool - - def __init__( - self, - name: str, - id: str, - created_at: Union[str, datetime.datetime], - region: Optional[Region], - firewall_ranges: List[str], - terminated_at: Optional[Union[str, datetime.datetime]], - allow_all_traffic: Optional[bool], - ): - #: Name of the workspace group - self.name = name - - #: Unique ID of the workspace group - self.id = id - - #: Timestamp of when the workspace group was created - self.created_at = to_datetime(created_at) - - #: Region of the workspace group (see :class:`Region`) - self.region = region - - #: List of allowed incoming IP addresses / ranges - self.firewall_ranges = firewall_ranges - - #: Timestamp of when the workspace group was terminated - self.terminated_at = to_datetime(terminated_at) - - #: Should all traffic be allowed? - self.allow_all_traffic = allow_all_traffic or False - - self._manager: Optional[WorkspaceManager] = 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: 'WorkspaceManager', - ) -> 'WorkspaceGroup': - """ - Construct a WorkspaceGroup from a dictionary of values. - - Parameters - ---------- - obj : dict - Dictionary of values - manager : WorkspaceManager, optional - The WorkspaceManager the WorkspaceGroup belongs to - - Returns - ------- - :class:`WorkspaceGroup` - - """ - try: - region = [x for x in manager.regions if x.id == obj['regionID']][0] - except IndexError: - region = Region('', '', obj.get('regionID', '')) - out = cls( - name=obj['name'], - id=obj['workspaceGroupID'], - created_at=obj['createdAt'], - region=region, - firewall_ranges=obj.get('firewallRanges', []), - terminated_at=obj.get('terminatedAt'), - allow_all_traffic=obj.get('allowAllTraffic'), - ) - out._manager = manager - return out - - @property - def organization(self) -> Organization: - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - return self._manager.organization - - @property - def stage(self) -> Stage: - """Stage manager.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - return Stage(self.id, self._manager) - - stages = stage - - def refresh(self) -> 'WorkspaceGroup': - """Update the object to the current state.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - new_obj = self._manager.get_workspace_group(self.id) - for name, value in vars(new_obj).items(): - if isinstance(value, Mapping): - setattr(self, name, camel_to_snake_dict(value)) - else: - setattr(self, name, value) - return self - - def update( - self, - name: Optional[str] = None, - firewall_ranges: Optional[List[str]] = None, - admin_password: Optional[str] = None, - expires_at: Optional[str] = None, - allow_all_traffic: Optional[bool] = None, - update_window: Optional[Dict[str, int]] = None, - ) -> None: - """ - Update the workspace group definition. - - Parameters - ---------- - name : str, optional - Name of the workspace group - firewall_ranges : list[str], optional - List of allowed CIDR ranges. An empty list indicates that all - inbound requests are allowed. - admin_password : str, optional - Admin password for the workspace group. If no password is supplied, - a password will be generated and retured in the response. - expires_at : str, optional - The timestamp of when the workspace group will expire. - If the expiration time is not specified, - the workspace group will have no expiration time. - At expiration, the workspace group is terminated and all the data is lost. - Expiration time can be specified as a timestamp or duration. - Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m" - allow_all_traffic : bool, optional - Allow all traffic to the workspace group - update_window : Dict[str, int], optional - Specify the day and hour of an update window: dict(day=0-6, hour=0-23) - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - data = { - k: v for k, v in dict( - name=name, - firewallRanges=firewall_ranges, - adminPassword=admin_password, - expiresAt=expires_at, - allowAllTraffic=allow_all_traffic, - updateWindow=snake_to_camel_dict(update_window), - ).items() if v is not None - } - self._manager._patch(f'workspaceGroups/{self.id}', json=data) - self.refresh() - - def terminate( - self, force: bool = False, - wait_on_terminated: bool = False, - wait_interval: int = 10, - wait_timeout: int = 600, - ) -> None: - """ - Terminate the workspace group. - - Parameters - ---------- - force : bool, optional - Terminate a workspace group even if it has active workspaces - wait_on_terminated : bool, optional - Wait for the workspace group to go into 'Terminated' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - self._manager._delete(f'workspaceGroups/{self.id}', params=dict(force=force)) - if wait_on_terminated: - while True: - self.refresh() - if self.terminated_at is not None: - break - if wait_timeout <= 0: - raise ManagementError( - msg='Exceeded waiting time for WorkspaceGroup to terminate', - ) - time.sleep(wait_interval) - wait_timeout -= wait_interval - - def create_workspace( - self, - name: str, - size: Optional[str] = None, - auto_suspend: Optional[Dict[str, Any]] = None, - cache_config: Optional[int] = None, - enable_kai: Optional[bool] = None, - wait_on_active: bool = False, - wait_interval: int = 10, - wait_timeout: int = 600, - ) -> Workspace: - """ - Create a new workspace. - - Parameters - ---------- - name : str - Name of the workspace - size : str, optional - Workspace size in workspace size notation (S-00, S-1, etc.) - auto_suspend : Dict[str, Any], optional - Auto suspend settings for the workspace. If this field is not - provided, no settings will be enabled. - cache_config : int, optional - Specifies the multiplier for the persistent cache associated - with the workspace. If specified, it enables the cache configuration - multiplier. It can have one of the following values: 1, 2, or 4. - enable_kai : bool, optional - Whether to create a SingleStore Kai-enabled workspace - wait_on_active : bool, optional - Wait for the workspace to be active before returning - wait_timeout : int, optional - Maximum number of seconds to wait before raising an exception - if wait=True - wait_interval : int, optional - Number of seconds between each polling interval - - Returns - ------- - :class:`Workspace` - - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - out = self._manager.create_workspace( - name=name, - workspace_group=self, - size=size, - auto_suspend=snake_to_camel_dict(auto_suspend), - cache_config=cache_config, - enable_kai=enable_kai, - wait_on_active=wait_on_active, - wait_interval=wait_interval, - wait_timeout=wait_timeout, - ) - - return out - - @property - def workspaces(self) -> NamedList[Workspace]: - """Return a list of available workspaces.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - res = self._manager._get('workspaces', params=dict(workspaceGroupID=self.id)) - return NamedList( - [Workspace.from_dict(item, self._manager) for item in res.json()], - ) - - -class StarterWorkspace(object): - """ - SingleStoreDB starter workspace definition. - - This object is not instantiated directly. It is used in the results - of API calls on the :class:`WorkspaceManager`. Existing starter workspaces are - accessed by either :attr:`WorkspaceManager.starter_workspaces` or by calling - :meth:`WorkspaceManager.get_starter_workspace`. - - See Also - -------- - :meth:`WorkspaceManager.get_starter_workspace` - :meth:`WorkspaceManager.create_starter_workspace` - :meth:`WorkspaceManager.terminate_starter_workspace` - :meth:`WorkspaceManager.create_starter_workspace_user` - :attr:`WorkspaceManager.starter_workspaces` - - """ - - name: str - id: str - database_name: str - endpoint: Optional[str] - - def __init__( - self, - name: str, - id: str, - database_name: str, - endpoint: Optional[str] = None, - ): - #: Name of the starter workspace - self.name = name - - #: Unique ID of the starter workspace - self.id = id - - #: Name of the database associated with the starter workspace - self.database_name = database_name - - #: Endpoint to connect to the starter workspace. The endpoint is in the form - #: of ``hostname:port`` - self.endpoint = endpoint - - self._manager: Optional[WorkspaceManager] = 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: 'WorkspaceManager', - ) -> 'StarterWorkspace': - """ - Construct a StarterWorkspace from a dictionary of values. - - Parameters - ---------- - obj : dict - Dictionary of values - manager : WorkspaceManager, optional - The WorkspaceManager the StarterWorkspace belongs to - - Returns - ------- - :class:`StarterWorkspace` - - """ - out = cls( - name=obj['name'], - id=obj['virtualWorkspaceID'], - database_name=obj['databaseName'], - endpoint=obj.get('endpoint'), - ) - out._manager = manager - return out - - def connect(self, **kwargs: Any) -> connection.Connection: - """ - Create a connection to the database server for this starter workspace. - - Parameters - ---------- - **kwargs : keyword-arguments, optional - Parameters to the SingleStoreDB `connect` function except host - and port which are supplied by the starter workspace object - - Returns - ------- - :class:`Connection` - - """ - if not self.endpoint: - raise ManagementError( - msg='An endpoint has not been set in this ' - 'starter workspace configuration', - ) - - kwargs['host'] = self.endpoint - kwargs['database'] = self.database_name - - return connection.connect(**kwargs) - - def terminate(self) -> None: - """Terminate the starter workspace.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - self._manager._delete(f'sharedtier/virtualWorkspaces/{self.id}') - - def refresh(self) -> StarterWorkspace: - """Update the object to the current state.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - new_obj = self._manager.get_starter_workspace(self.id) - for name, value in vars(new_obj).items(): - if isinstance(value, Mapping): - setattr(self, name, snake_to_camel_dict(value)) - else: - setattr(self, name, value) - return self - - @property - def organization(self) -> Organization: - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - return self._manager.organization - - @property - def stage(self) -> Stage: - """Stage manager.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - return Stage(self.id, self._manager) - - stages = stage - - @property - def starter_workspaces(self) -> NamedList['StarterWorkspace']: - """Return a list of available starter workspaces.""" - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - res = self._manager._get('sharedtier/virtualWorkspaces') - return NamedList( - [StarterWorkspace.from_dict(item, self._manager) for item in res.json()], - ) - - def create_user( - self, - username: str, - password: Optional[str] = None, - ) -> Dict[str, str]: - """ - Create a new user for this starter workspace. - - Parameters - ---------- - username : str - The starter workspace user name to connect the new user to the database - password : str, optional - Password for the new user. If not provided, a password will be - auto-generated by the system. - - Returns - ------- - Dict[str, str] - Dictionary containing 'userID' and 'password' of the created user - - Raises - ------ - ManagementError - If no workspace manager is associated with this object. - """ - if self._manager is None: - raise ManagementError( - msg='No workspace manager is associated with this object.', - ) - - payload = { - 'userName': username, - } - if password is not None: - payload['password'] = password - - res = self._manager._post( - f'sharedtier/virtualWorkspaces/{self.id}/users', - json=payload, - ) - - response_data = res.json() - user_id = response_data.get('userID') - if not user_id: - raise ManagementError(msg='No userID returned from API') - - # Return the password provided by user or generated by API - returned_password = password if password is not None \ - else response_data.get('password') - if not returned_password: - raise ManagementError(msg='No password available from API response') - - return { - 'user_id': user_id, - 'password': returned_password, - } - - -class Billing(object): - """Billing information.""" - - COMPUTE_CREDIT = 'compute_credit' - STORAGE_AVG_BYTE = 'storage_avg_byte' - - HOUR = 'hour' - DAY = 'day' - MONTH = 'month' - - def __init__(self, manager: Manager): - self._manager = manager - - def usage( - self, - start_time: datetime.datetime, - end_time: datetime.datetime, - metric: Optional[str] = None, - aggregate_by: Optional[str] = None, - ) -> List[BillingUsageItem]: - """ - Get usage information. - - Parameters - ---------- - start_time : datetime.datetime - Start time for usage interval - end_time : datetime.datetime - End time for usage interval - metric : str, optional - Possible metrics are ``mgr.billing.COMPUTE_CREDIT`` and - ``mgr.billing.STORAGE_AVG_BYTE`` (default is all) - aggregate_by : str, optional - Aggregate type used to group usage: ``mgr.billing.HOUR``, - ``mgr.billing.DAY``, or ``mgr.billing.MONTH`` - - Returns - ------- - List[BillingUsage] - - """ - res = self._manager._get( - 'billing/usage', - params={ - k: v for k, v in dict( - 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, - ).items() if v is not None - }, - ) - return [ - BillingUsageItem.from_dict(x, self._manager) - for x in res.json()['billingUsage'] - ] - - -class Organizations(object): - """Organizations.""" - - def __init__(self, manager: Manager): - self._manager = manager - - @property - def current(self) -> Organization: - """Get current organization.""" - res = self._manager._get('organizations/current').json() - return Organization.from_dict(res, self._manager) - - -class WorkspaceManager(Manager): +from ._version_import import _import_versioned_module +from .v1.organization import Organization as Organization +from .v1.workspace import Billing as Billing +from .v1.workspace import get_organization as get_organization +from .v1.workspace import get_secret as get_secret +from .v1.workspace import get_stage as get_stage +from .v1.workspace import get_workspace as get_workspace +from .v1.workspace import get_workspace_group as get_workspace_group +from .v1.workspace import Organizations as Organizations +from .v1.workspace import Stage as Stage +from .v1.workspace import StarterWorkspace as StarterWorkspace +from .v1.workspace import Workspace as Workspace +from .v1.workspace import WorkspaceGroup as WorkspaceGroup +from .v1.workspace import WorkspaceManager as WorkspaceManager +# Re-export from default version for backward compatibility + + +def _manage_workspaces_v1( + access_token: Optional[str] = None, + version: Optional[str] = None, + base_url: Optional[str] = None, + *, + organization_id: Optional[str] = None, +) -> 'WorkspaceManager': """ - SingleStoreDB workspace manager. + Retrieve a SingleStoreDB workspace manager without warning. - This class should be instantiated using :func:`singlestoredb.manage_workspaces`. - - Parameters - ---------- - access_token : str, optional - The API key or other access token for the workspace management API - version : str, optional - Version of the API to use - base_url : str, optional - Base URL of the workspace management API - - See Also - -------- - :func:`singlestoredb.manage_workspaces` + This is the body of :func:`manage_workspaces` minus the deprecation warning. + Internal callers that are v1-only by design -- Fusion, the UDF ``stage://`` + handling, the AI inference helpers -- go through here so they do not emit a + warning the caller can do nothing about. They are asking for a workspace + manager specifically, not for whatever the environment prefers. + Neither function consults the ``management.version`` option: workspaces + exist only at v1. See :func:`manage_workspaces` for why. """ - - #: Workspace management API version if none is specified. - default_version = config.get_option('management.version') or 'v1' - - #: Base URL if none is specified. - default_base_url = config.get_option('management.base_url') \ - or 'https://api.singlestore.com' - - #: Object type - obj_type = 'workspace' - - @property - def workspace_groups(self) -> NamedList[WorkspaceGroup]: - """Return a list of available workspace groups.""" - res = self._get('workspaceGroups') - return NamedList([WorkspaceGroup.from_dict(item, self) for item in res.json()]) - - @property - def starter_workspaces(self) -> NamedList[StarterWorkspace]: - """Return a list of available starter workspaces.""" - res = self._get('sharedtier/virtualWorkspaces') - return NamedList([StarterWorkspace.from_dict(item, self) for item in res.json()]) - - @property - def organizations(self) -> Organizations: - """Return the organizations.""" - return Organizations(self) - - @property - def organization(self) -> Organization: - """ Return the current organization.""" - return self.organizations.current - - @property - def billing(self) -> Billing: - """Return the current billing information.""" - return Billing(self) - - @ttl_property(datetime.timedelta(hours=1)) - def regions(self) -> NamedList[Region]: - """Return a list of available regions.""" - res = self._get('regions') - return NamedList([Region.from_dict(item, self) for item in res.json()]) - - @ttl_property(datetime.timedelta(hours=1)) - def shared_tier_regions(self) -> NamedList[Region]: - """Return a list of regions that support shared tier workspaces.""" - res = self._get('regions/sharedtier') - return NamedList( - [Region.from_dict(item, self) for item in res.json()], + from ..exceptions import ManagementError + ver = version or 'v1' + if ver != 'v1': + raise ManagementError( + msg=f'workspaces do not exist in management API {ver}; they were ' + 'replaced by clusters. Use manage_clusters() instead, or pass ' + 'version="v1" here. Note that the management.version option ' + 'does not reach this function: workspaces are v1-only, so it ' + 'has nothing to select.', ) - - def create_workspace_group( - self, - name: str, - region: Union[str, Region], - firewall_ranges: List[str], - admin_password: Optional[str] = None, - backup_bucket_kms_key_id: Optional[str] = None, - data_bucket_kms_key_id: Optional[str] = None, - expires_at: Optional[str] = None, - smart_dr: Optional[bool] = None, - allow_all_traffic: Optional[bool] = None, - update_window: Optional[Dict[str, int]] = None, - ) -> WorkspaceGroup: - """ - Create a new workspace group. - - Parameters - ---------- - name : str - Name of the workspace group - region : str or Region - ID of the region where the workspace group should be created - firewall_ranges : list[str] - List of allowed CIDR ranges. An empty list indicates that all - inbound requests are allowed. - admin_password : str, optional - Admin password for the workspace group. If no password is supplied, - a password will be generated and retured in the response. - backup_bucket_kms_key_id : str, optional - Specifies the KMS key ID associated with the backup bucket. - If specified, enables Customer-Managed Encryption Keys (CMEK) - encryption for the backup bucket of the workspace group. - This feature is only supported in workspace groups deployed in AWS. - data_bucket_kms_key_id : str, optional - Specifies the KMS key ID associated with the data bucket. - If specified, enables Customer-Managed Encryption Keys (CMEK) - encryption for the data bucket and Amazon Elastic Block Store - (EBS) volumes of the workspace group. This feature is only supported - in workspace groups deployed in AWS. - expires_at : str, optional - The timestamp of when the workspace group will expire. - If the expiration time is not specified, - the workspace group will have no expiration time. - At expiration, the workspace group is terminated and all the data is lost. - Expiration time can be specified as a timestamp or duration. - Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m" - smart_dr : bool, optional - Enables Smart Disaster Recovery (SmartDR) for the workspace group. - SmartDR is a disaster recovery solution that ensures seamless and - continuous replication of data from the primary region to a secondary region - allow_all_traffic : bool, optional - Allow all traffic to the workspace group - update_window : Dict[str, int], optional - Specify the day and hour of an update window: dict(day=0-6, hour=0-23) - - Returns - ------- - :class:`WorkspaceGroup` - - """ - if isinstance(region, Region) and region.id: - region = region.id - res = self._post( - 'workspaceGroups', json=dict( - name=name, regionID=region, - adminPassword=admin_password, - backupBucketKMSKeyID=backup_bucket_kms_key_id, - dataBucketKMSKeyID=data_bucket_kms_key_id, - firewallRanges=firewall_ranges or [], - expiresAt=expires_at, - smartDR=smart_dr, - allowAllTraffic=allow_all_traffic, - updateWindow=snake_to_camel_dict(update_window), - ), - ) - return self.get_workspace_group(res.json()['workspaceGroupID']) - - def create_workspace( - self, - name: str, - workspace_group: Union[str, WorkspaceGroup], - size: Optional[str] = None, - auto_suspend: Optional[Dict[str, Any]] = None, - cache_config: Optional[int] = None, - enable_kai: Optional[bool] = None, - wait_on_active: bool = False, - wait_interval: int = 10, - wait_timeout: int = 600, - ) -> Workspace: - """ - Create a new workspace. - - Parameters - ---------- - name : str - Name of the workspace - workspace_group : str or WorkspaceGroup - The workspace ID of the workspace - size : str, optional - Workspace size in workspace size notation (S-00, S-1, etc.) - auto_suspend : Dict[str, Any], optional - Auto suspend settings for the workspace. If this field is not - provided, no settings will be enabled. - cache_config : int, optional - Specifies the multiplier for the persistent cache associated - with the workspace. If specified, it enables the cache configuration - multiplier. It can have one of the following values: 1, 2, or 4. - enable_kai : bool, optional - Whether to create a SingleStore Kai-enabled workspace - wait_on_active : bool, optional - Wait for the workspace to be active before returning - wait_timeout : int, optional - Maximum number of seconds to wait before raising an exception - if wait=True - wait_interval : int, optional - Number of seconds between each polling interval - - Returns - ------- - :class:`Workspace` - - """ - if isinstance(workspace_group, WorkspaceGroup): - workspace_group = workspace_group.id - res = self._post( - 'workspaces', json=dict( - name=name, - workspaceGroupID=workspace_group, - size=size, - autoSuspend=snake_to_camel_dict(auto_suspend), - cacheConfig=cache_config, - enableKai=enable_kai, - ), - ) - out = self.get_workspace(res.json()['workspaceID']) - if wait_on_active: - out = self._wait_on_state( - out, - 'Active', - interval=wait_interval, - timeout=wait_timeout, - ) - # After workspace is active, wait for endpoint to be ready - out = self._wait_on_endpoint( - out, - interval=wait_interval, - timeout=wait_timeout, - ) - return out - - def get_workspace_group(self, id: str) -> WorkspaceGroup: - """ - Retrieve a workspace group definition. - - Parameters - ---------- - id : str - ID of the workspace group - - Returns - ------- - :class:`WorkspaceGroup` - - """ - res = self._get(f'workspaceGroups/{id}') - return WorkspaceGroup.from_dict(res.json(), manager=self) - - def get_workspace(self, id: str) -> Workspace: - """ - Retrieve a workspace definition. - - Parameters - ---------- - id : str - ID of the workspace - - Returns - ------- - :class:`Workspace` - - """ - res = self._get(f'workspaces/{id}') - return Workspace.from_dict(res.json(), manager=self) - - def get_starter_workspace(self, id: str) -> StarterWorkspace: - """ - Retrieve a starter workspace definition. - - Parameters - ---------- - id : str - ID of the starter workspace - - Returns - ------- - :class:`StarterWorkspace` - - """ - res = self._get(f'sharedtier/virtualWorkspaces/{id}') - return StarterWorkspace.from_dict(res.json(), manager=self) - - def create_starter_workspace( - self, - name: str, - database_name: str, - provider: str, - region_name: str, - ) -> 'StarterWorkspace': - """ - Create a new starter (shared tier) workspace. - - Parameters - ---------- - name : str - Name of the starter workspace - database_name : str - Name of the database for the starter workspace - provider : str - Cloud provider for the starter workspace (e.g., 'aws', 'gcp', 'azure') - region_name : str - Cloud provider region for the starter workspace (e.g., 'us-east-1') - - Returns - ------- - :class:`StarterWorkspace` - """ - - payload = { - 'name': name, - 'databaseName': database_name, - 'provider': provider, - 'regionName': region_name, - } - - res = self._post('sharedtier/virtualWorkspaces', json=payload) - virtual_workspace_id = res.json().get('virtualWorkspaceID') - if not virtual_workspace_id: - raise ManagementError(msg='No virtualWorkspaceID returned from API') - - res = self._get(f'sharedtier/virtualWorkspaces/{virtual_workspace_id}') - return StarterWorkspace.from_dict(res.json(), self) + mod = _import_versioned_module(ver, 'workspace') + return mod.WorkspaceManager( + access_token=access_token, base_url=base_url, + version=ver, organization_id=organization_id, + ) def manage_workspaces( @@ -1965,16 +77,24 @@ def manage_workspaces( base_url: Optional[str] = None, *, organization_id: Optional[str] = None, -) -> WorkspaceManager: +) -> 'WorkspaceManager': """ Retrieve a SingleStoreDB workspace manager. + .. deprecated:: + Workspaces and workspace groups were replaced by the flat ``Cluster`` + resource in management API v2. Use + :func:`singlestoredb.manage_clusters` instead. + Parameters ---------- access_token : str, optional The API key or other access token for the workspace management API version : str, optional - Version of the API to use + Version of the API to use. Defaults to ``'v1'``, **not** to the + ``management.version`` option: workspaces exist only at v1, so there is + no version for this function to dispatch on. Passing anything else + raises. This is the one public entry point the option does not steer. base_url : str, optional Base URL of the workspace management API organization_id : str, optional @@ -1984,8 +104,35 @@ def manage_workspaces( ------- :class:`WorkspaceManager` + Raises + ------ + :class:`ManagementError` + If the caller explicitly asks for a version other than ``v1``. + Workspaces and workspace groups were replaced by clusters in v2; use + :func:`singlestoredb.manage_clusters` instead. + """ - return WorkspaceManager( - access_token=access_token, base_url=base_url, - version=version, organization_id=organization_id, + warnings.warn( + 'manage_workspaces() is deprecated: workspaces and workspace groups ' + 'were replaced by the flat Cluster resource in management API v2. ' + 'Use manage_clusters() instead. This still returns a working v1 ' + 'manager.', + DeprecationWarning, + stacklevel=2, + ) + # Pinned to v1 rather than resolved through the management.version option. + # The option selects between implementations of a resource that exists at + # more than one version; workspaces exist only at v1, so there is nothing + # here for it to select, and resolving it would make a bare + # manage_workspaces() raise as soon as the default moved past v1 -- v1 + # ceasing to work rather than v1 being deprecated. Callers are steered to + # clusters by the deprecation warning above, not by an exception. + # + # Deliberately *not* symmetrical with manage_clusters(), which does consult + # the option and raises when it names v1: an option reading v1 is a + # deliberate request that manage_clusters() cannot satisfy, whereas an + # option sitting at its default says nothing about workspaces. + return _manage_workspaces_v1( + access_token, version, base_url, + organization_id=organization_id, ) diff --git a/singlestoredb/notebook/_objects.py b/singlestoredb/notebook/_objects.py index a3c16cd61..814b583ac 100644 --- a/singlestoredb/notebook/_objects.py +++ b/singlestoredb/notebook/_objects.py @@ -3,7 +3,23 @@ from typing import Any from typing import Optional +from .. import management as _mgmt from ..management import workspace as _ws +from ..management.organization import Organization as _OrganizationBase +from ..management.stage import Stage as _StageBase +# Still the v1 shim, and only for the two globals below that are v1 vocabulary +# outright: ``workspace`` and ``workspacegroup``. Management API v2 replaced +# both with the flat ``Cluster``, and there is no ``cluster`` notebook global +# to proxy to yet, so these cannot simply be repointed -- giving the notebook +# environment a cluster global is a port, not a version bump. The other three +# globals (``secrets``, ``stage``, ``organization``) go through ``_mgmt``, +# whose helpers dispatch on the ``management.version`` option; taking them from +# this shim pinned them to v1 no matter what the option said. +# +# Consequence for a v1 notebook environment: those three now follow the option, +# which defaults to v2, so such an environment has to set +# SINGLESTOREDB_MANAGEMENT_VERSION=v1. That is the cost of them being neutral at +# all -- pinned, v1 worked and v2 was simply broken. class Secrets(object): @@ -12,10 +28,10 @@ class Secrets(object): def __getattr__(self, name: str) -> Optional[str]: if name.startswith('_ipython') or name.startswith('_repr_'): raise AttributeError(name) - return _ws.get_secret(name) + return _mgmt.get_secret(name) def __getitem__(self, name: str) -> Optional[str]: - return _ws.get_secret(name) + return _mgmt.get_secret(name) class Stage(object): @@ -25,36 +41,36 @@ def __new__(cls, *args: Any, **kwargs: Any) -> Any: # autocomplete still works in Jupyter / IPython, but we # bypass the real method / attribute calls and apply them # to the currently selected stage. - for name in [x for x in dir(_ws.Stage) if not x.startswith('_')]: + for name in [x for x in dir(_StageBase) if not x.startswith('_')]: if name in ['from_dict', 'refresh', 'update']: continue - attr = getattr(_ws.Stage, name) + attr = getattr(_StageBase, name) def make_wrapper(m: str, is_method: bool = False) -> Any: if is_method: def wrap(self: Stage, *a: Any, **kw: Any) -> Any: - return getattr(_ws.get_stage(), m)(*a, **kw) + return getattr(_mgmt.get_stage(), m)(*a, **kw) return functools.update_wrapper(wrap, attr) else: def wrap(self: Stage, *a: Any, **kw: Any) -> Any: - return getattr(_ws.get_stage(), m) + return getattr(_mgmt.get_stage(), m) return property(functools.update_wrapper(wrap, attr)) setattr(cls, name, make_wrapper(m=name, is_method=callable(attr))) for name in [ - x for x in _ws.Stage.__annotations__.keys() + x for x in _StageBase.__annotations__.keys() if not x.startswith('_') ]: - def make_wrapper(m: str, is_method: bool = False) -> Any: + def make_annotation_wrapper(m: str) -> Any: def wrap(self: Stage) -> Any: - return getattr(_ws.get_stage(), m) - return property(functools.update_wrapper(wrap, attr)) + return getattr(_mgmt.get_stage(), m) + return property(wrap) - setattr(cls, name, make_wrapper(m=name)) + setattr(cls, name, make_annotation_wrapper(m=name)) - cls.__doc__ = _ws.Stage.__doc__ + cls.__doc__ = _StageBase.__doc__ return super().__new__(cls, *args, **kwargs) @@ -89,12 +105,12 @@ def wrap(self: WorkspaceGroup, *a: Any, **kw: Any) -> Any: if not x.startswith('_') ]: - def make_wrapper(m: str, is_method: bool = False) -> Any: + def make_annotation_wrapper(m: str) -> Any: def wrap(self: WorkspaceGroup) -> Any: return getattr(_ws.get_workspace_group(), m) - return property(functools.update_wrapper(wrap, attr)) + return property(wrap) - setattr(cls, name, make_wrapper(m=name)) + setattr(cls, name, make_annotation_wrapper(m=name)) cls.__doc__ = _ws.WorkspaceGroup.__doc__ @@ -137,12 +153,12 @@ def wrap(self: Workspace, *a: Any, **kw: Any) -> Any: if not x.startswith('_') ]: - def make_wrapper(m: str, is_method: bool = False) -> Any: + def make_annotation_wrapper(m: str) -> Any: def wrap(self: Workspace) -> Any: return getattr(_ws.get_workspace(), m) - return property(functools.update_wrapper(wrap, attr)) + return property(wrap) - setattr(cls, name, make_wrapper(m=name)) + setattr(cls, name, make_annotation_wrapper(m=name)) cls.__doc__ = _ws.Workspace.__doc__ @@ -162,45 +178,45 @@ def __new__(cls, *args: Any, **kwargs: Any) -> Any: # autocomplete still works in Jupyter / IPython, but we # bypass the real method / attribute calls and apply them # to the currently selected organization. - for name in [x for x in dir(_ws.Organization) if not x.startswith('_')]: + for name in [x for x in dir(_OrganizationBase) if not x.startswith('_')]: if name in ['from_dict', 'refresh', 'update']: continue - attr = getattr(_ws.Organization, name) + attr = getattr(_OrganizationBase, name) def make_wrapper(m: str, is_method: bool = False) -> Any: if is_method: def wrap(self: Organization, *a: Any, **kw: Any) -> Any: - return getattr(_ws.get_organization(), m)(*a, **kw) + return getattr(_mgmt.get_organization(), m)(*a, **kw) return functools.update_wrapper(wrap, attr) else: def wrap(self: Organization, *a: Any, **kw: Any) -> Any: - return getattr(_ws.get_organization(), m) + return getattr(_mgmt.get_organization(), m) return property(functools.update_wrapper(wrap, attr)) setattr(cls, name, make_wrapper(m=name, is_method=callable(attr))) for name in [ - x for x in _ws.Organization.__annotations__.keys() + x for x in _OrganizationBase.__annotations__.keys() if not x.startswith('_') ]: - def make_wrapper(m: str, is_method: bool = False) -> Any: + def make_annotation_wrapper(m: str) -> Any: def wrap(self: Organization) -> Any: - return getattr(_ws.get_organization(), m) - return property(functools.update_wrapper(wrap, attr)) + return getattr(_mgmt.get_organization(), m) + return property(wrap) - setattr(cls, name, make_wrapper(m=name)) + setattr(cls, name, make_annotation_wrapper(m=name)) - cls.__doc__ = _ws.Organization.__doc__ + cls.__doc__ = _OrganizationBase.__doc__ return super().__new__(cls, *args, **kwargs) def __str__(self) -> str: - return _ws.get_organization().__str__() + return _mgmt.get_organization().__str__() def __repr__(self) -> str: - return _ws.get_organization().__repr__() + return _mgmt.get_organization().__repr__() secrets = Secrets() diff --git a/singlestoredb/notebook/_portal.py b/singlestoredb/notebook/_portal.py index 861737280..4b9462126 100644 --- a/singlestoredb/notebook/_portal.py +++ b/singlestoredb/notebook/_portal.py @@ -135,7 +135,13 @@ def secrets(self) -> obj.Secrets: @property def workspace_group_id(self) -> Optional[str]: - """Workspace Group ID.""" + """ + Workspace Group ID. + + The deployment's group ID. A group is an addressable resource only at + management API v1; at v2 the same ID is reported back as + ``Cluster.group`` and cannot be looked up. + """ try: return self._connection_info['workspace_group'] except KeyError: @@ -156,7 +162,12 @@ def workspace_group(self) -> None: @property def workspace_id(self) -> Optional[str]: - """Workspace ID.""" + """ + Workspace ID. + + The current deployment: a workspace ID at v1, a cluster ID from v2 + onward. See :attr:`cluster_id`, which is the same value. + """ try: return self._connection_info['workspace'] except KeyError: @@ -260,11 +271,31 @@ def connection( @property def cluster_id(self) -> Optional[str]: - """Cluster ID.""" + """ + Cluster ID. + + The same value as :attr:`workspace_id`: management API v2 calls the + deployment a cluster where v1 called it a workspace, and the notebook + environment publishes it under its original name -- + ``SINGLESTOREDB_WORKSPACE`` -- rather than adding a second variable. + """ + return self.workspace_id + + @property + def project_id(self) -> Optional[str]: + """ + Project ID. + + The inference API project, which is not a project of the cluster + management API: the two are separate namespaces and the notebook + environment reports different IDs for them. Do not pass this where a + management project is wanted -- see + :func:`singlestoredb.management.utils.get_project_id`. + """ try: - return self._connection_info['cluster'] + return self._connection_info['project'] except KeyError: - return os.environ.get('SINGLESTOREDB_CLUSTER') + return os.environ.get('SINGLESTOREDB_PROJECT') def _parse_url(self) -> Dict[str, Any]: url = urllib.parse.urlparse( diff --git a/singlestoredb/tests/cleanup_deployments.py b/singlestoredb/tests/cleanup_deployments.py new file mode 100644 index 000000000..7b8d745c2 --- /dev/null +++ b/singlestoredb/tests/cleanup_deployments.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python +# type: ignore +r""" +Terminate deployments left behind by earlier test runs. + +The test suite now sweeps what it creates (see ``utils.track()`` and the +hooks in ``conftest.py``), but a run that was killed -- or one from before +that sweep existed -- leaves live workspace groups, workspaces, clusters and +starter clusters behind, and they are billed until someone removes them. + +Only names the test suite generates are considered, and the default is a dry +run:: + + python -m singlestoredb.tests.cleanup_deployments + python -m singlestoredb.tests.cleanup_deployments --yes + +If the organization is visibly full of strays and this reports none, the names +are not in ``PATTERNS``. ``--show-unmatched`` lists every live deployment the +tool does not recognize, which is how an unconventionally named one gets found. + +This tool is organization-wide, not run-scoped: it matches on names, and a +name says which suite made a deployment but not which run. A concurrent run's +fixtures look exactly like stranded ones. Age is the only thing separating +them, so ``--older-than`` defaults to a span longer than a full suite rather +than to zero -- raise it if your runs can take longer than that, and only +pass ``--older-than 0`` when you know nothing else is running. + +The other direction -- clearing out what a recent session made, rather than +what an old one stranded -- is ``--since``, which replaces the age guard with +a calendar cutoff, and ``--any-name``, which drops the name gate. Clearing +every workspace group created yesterday or today:: + + python -m singlestoredb.tests.cleanup_deployments \ + --kind workspace-group --any-name --since yesterday --yes + +Those two flags together remove both of the guards that keep this tool off +deployments it did not create, so ``--kind`` matters: without it the same +cutoff sweeps every cluster of that age as well, the shared pool included. +Read the dry run before adding ``--yes``. + +For deployments the current process created, nothing here is needed: those +are tracked as they are created and swept per test class by ``conftest.py``, +which cannot see -- or touch -- another run's deployments. + +Why strays keep appearing: that tracking, the per-class sweep and this script +all live on the ``versioned-management-api`` branch and nowhere else. A run +from ``main`` has only ``tearDownClass``, so a killed run or a ``setUpClass`` +that raises leaks a workspace group permanently, and ``main`` still uses names +this script only knows through :data:`LEGACY_PATTERNS`. Until the sweep is on +the default branch, expect to run this by hand. +""" +import argparse +import datetime +import re +import sys +import warnings +from collections.abc import Container +from typing import Any +from typing import List +from typing import Optional +from typing import Tuple + +import singlestoredb as s2 + + +#: The kinds of deployment this tool can sweep, in the order it lists them. +#: ``--kind`` selects from these; the default is all of them. +KINDS = ( + 'cluster', + 'starter-cluster', + 'workspace-group', + 'starter-workspace', +) + +#: Hours a deployment must have existed before it is treated as stranded. +#: The slowest class creates three clusters with a 1200s wait each and then +#: terminates them the same way, so a full suite is comfortably inside this; +#: anything younger could belong to a run in progress. +DEFAULT_MIN_AGE_HOURS = 6.0 + +#: Names the suite generates. Anchored, because these run against a real +#: organization: a pattern that matched a name someone chose by hand would +#: terminate a deployment that is not ours. +PATTERNS = [ + # test_management_v1.py / test_management_v2.py fixtures + re.compile(r'^(wg|ws|cl)-test-[A-Za-z0-9_-]+$'), + re.compile(r'^starter-(ws|cl)-test-[A-Za-z0-9_-]+$'), + # test_fusion.py fixtures + re.compile(r'^[A-C] Fusion Testing [0-9a-f]+$'), + re.compile(r'^[a-z]-fusion-cluster-[0-9a-f]+$'), + re.compile(r'^jobs-fusion-[0-9a-f]+$'), + re.compile(r'^stage-fusion-\d-[0-9a-f]+$'), + # test_create_drop_workspace_group's subject. Hex covers the decimal + # id(self) the test used to name it with, so groups stranded by older + # runs -- which this pattern did not match, and which therefore piled up + # invisibly -- are reaped too. + re.compile(r'^Create WG Test [0-9a-f]+$'), +] + +#: Names the suite used to generate. Kept separate so it is obvious what is +#: only here for cleanup, and matched all the same: a stranded deployment is +#: billed regardless of which revision made it, and ``main`` still creates +#: these -- it carries none of ``utils.track()``, the per-class sweep or this +#: script, so a run there leaks with nothing to reap it. Retire an entry once +#: no branch produces the name and the organization is clean of it. +LEGACY_PATTERNS = [ + # TestStageFusion's two workspace groups, before it moved to v2 clusters + # named stage-fusion-- and then to the shared cluster pool + re.compile(r'^Stage Fusion Testing \d [0-9a-f]+$'), + # TestFilesFusion's workspace group, which nothing in the class ever + # read; it creates no deployment at all now + re.compile(r'^Files Fusion Testing [0-9a-f]+$'), + # 'Group '. No revision of this repo generates this, so it is here + # on the owner's say-so rather than by attribution. Eight hex characters + # minimum, which is what the ones in the organization have: the bare + # 'Group 1' / 'Group 2' that a person or the portal produces is a real + # deployment someone is using, and a plain [0-9a-f]+ would match it. + re.compile(r'^Group [0-9a-f]{8,}$'), +] + + +def is_test_deployment(name: Optional[str]) -> bool: + """Was this name generated by the test suite, now or in the past?""" + if not name: + return False + return any(x.match(name) for x in PATTERNS + LEGACY_PATTERNS) + + +def _created_at(obj: Any) -> Optional[datetime.datetime]: + """When this deployment was created, or None if the API did not say.""" + created = getattr(obj, 'created_at', None) + if not isinstance(created, datetime.datetime): + return None + if created.tzinfo is None: + # A naive timestamp from the API is UTC. Reading it as local time + # would overstate the age by the offset, which is the direction that + # sweeps a deployment a live run still owns. + created = created.replace(tzinfo=datetime.timezone.utc) + return created + + +def _age_hours(obj: Any) -> Optional[float]: + """Hours since creation, or None if the API did not report it.""" + created = _created_at(obj) + if created is None: + return None + now = datetime.datetime.now(tz=datetime.timezone.utc) + return (now - created).total_seconds() / 3600.0 + + +def parse_since(text: str) -> datetime.datetime: + """ + Read a ``--since`` value as the local midnight starting that day. + + ``today``, ``yesterday`` or an ISO date. The cutoff is local midnight + rather than a UTC one because the caller is thinking in their own + calendar days -- "created yesterday" means yesterday where they are. + """ + today = datetime.date.today() + if text == 'today': + day = today + elif text == 'yesterday': + day = today - datetime.timedelta(days=1) + else: + try: + day = datetime.date.fromisoformat(text) + except ValueError: + raise argparse.ArgumentTypeError( + f'{text!r} is not a date; expected YYYY-MM-DD, ' + "'today' or 'yesterday'", + ) + # A naive datetime's astimezone() reads it as local time, which is what + # gives midnight the caller's offset rather than UTC's. + return datetime.datetime.combine(day, datetime.time.min).astimezone() + + +def find_leftovers( + older_than: float = DEFAULT_MIN_AGE_HOURS, + include_unknown_age: bool = False, + since: Optional[datetime.datetime] = None, + any_name: bool = False, + kinds: Container[str] = KINDS, +) -> Tuple[List[Tuple[str, Any]], List[str], List[str]]: + """ + List the live, test-named deployments in the current organization. + + Both API versions are asked: v1 owns workspace groups and workspaces, + v2 owns clusters, and a suite that has run under either may have left + something behind. + + ``since`` replaces the ``older_than`` guard with the opposite test -- + created at or after that moment, rather than old enough to be stranded -- + and ``any_name`` drops the name gate, which makes every live deployment a + candidate. Between them they turn this from "sweep what the suite + stranded" into "clear out this organization", so ``kinds`` is what keeps + such a run off deployments the caller did not mean. + + Returns + ------- + (List[Tuple[str, Any]], List[str], List[str]) + The deployments to sweep, labels for the ones held back by the age + guard so the caller can say what it did not touch, and labels for the + live deployments whose names :data:`PATTERNS` does not recognize. + + That third list is the answer to "the organization is full of strays + and this tool says there are none". A test that names a deployment + outside the conventions above is invisible here, so it accumulates + silently -- which is exactly what ``Create WG Test `` did. + Reporting the unrecognized names makes the next one findable. + + """ + found: List[Tuple[str, Any]] = [] + spared: List[str] = [] + unmatched: List[str] = [] + + def keep(obj: Any) -> bool: + name = getattr(obj, 'name', None) + if getattr(obj, 'terminated_at', None) is not None: + return False + if not any_name and not is_test_deployment(name): + age = _age_hours(obj) + unmatched.append( + '{}{}'.format( + name or '', + '' if age is None else f' ({age:.1f}h old)', + ), + ) + return False + + # Age is the only thing separating a stranded deployment from one a + # concurrent run is using right now: names carry a per-class random + # id, not a per-run one, and a cluster name is capped at 32 + # characters, so there is no room to stamp a run id into it. + created = _created_at(obj) + if created is None: + if not include_unknown_age: + spared.append(f'{name} (creation time not reported)') + return False + return True + if since is not None: + if created < since: + spared.append( + f'{name} (created {created.astimezone():%Y-%m-%d %H:%M}, ' + 'before the cutoff)', + ) + return False + return True + now = datetime.datetime.now(tz=datetime.timezone.utc) + age = (now - created).total_seconds() / 3600.0 + if older_than > 0 and age < older_than: + spared.append(f'{name} ({age:.1f}h old, too new)') + return False + return True + + if 'cluster' in kinds or 'starter-cluster' in kinds: + try: + clusters = s2.manage_clusters(version='v2') + except Exception as exc: + print(f'! Could not reach management API v2: {exc}', file=sys.stderr) + else: + if 'cluster' in kinds: + for cluster in clusters.clusters: + if keep(cluster): + found.append(( + f'cluster {cluster.name} ({cluster.id})', cluster, + )) + if 'starter-cluster' in kinds: + for starter in clusters.starter_clusters: + if keep(starter): + found.append(( + f'starter cluster {starter.name} ({starter.id})', + starter, + )) + + if 'workspace-group' in kinds or 'starter-workspace' in kinds: + try: + # v1 is deprecated, and asking for it here is the point: workspace + # groups exist nowhere else, so the warning is noise on every run. + with warnings.catch_warnings(): + warnings.filterwarnings( + 'ignore', category=DeprecationWarning, + message='.*manage_workspaces.*', + ) + workspaces = s2.manage_workspaces(version='v1') + except Exception as exc: + print(f'! Could not reach management API v1: {exc}', file=sys.stderr) + else: + if 'workspace-group' in kinds: + for group in workspaces.workspace_groups: + if keep(group): + # The group takes its workspaces with it, so they are + # not listed separately. + found.append(( + f'workspace group {group.name} ({group.id})', group, + )) + if 'starter-workspace' in kinds: + for starter in workspaces.starter_workspaces: + if keep(starter): + found.append(( + f'starter workspace {starter.name} ({starter.id})', + starter, + )) + + return found, spared, unmatched + + +def main(argv: Optional[List[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split('\n\n')[1]) + parser.add_argument( + '--yes', action='store_true', + help='actually terminate; without this the run only reports', + ) + parser.add_argument( + '--older-than', type=float, default=DEFAULT_MIN_AGE_HOURS, + metavar='HOURS', + help='only sweep deployments at least this old ' + f'(default: {DEFAULT_MIN_AGE_HOURS}). Pass 0 to sweep every ' + 'match, which will terminate deployments a concurrent test run ' + 'is still using', + ) + parser.add_argument( + '--since', type=parse_since, metavar='DATE', + help="sweep what was created on or after DATE -- 'today', " + "'yesterday' or YYYY-MM-DD, counted from local midnight -- " + 'instead of what is older than --older-than. This is for ' + 'clearing out a recent session rather than reaping strays, so ' + 'it removes the guard against terminating a deployment a live ' + 'run owns: pair it with --kind', + ) + parser.add_argument( + '--any-name', action='store_true', + help='consider every live deployment, not only the ones named like ' + "the test suite's. This will terminate deployments nothing in " + 'this repo created, including ones a colleague is using, so ' + 'read the dry run first', + ) + parser.add_argument( + '--kind', action='append', choices=KINDS, dest='kinds', + metavar='KIND', + help='restrict the sweep to this kind of deployment; repeatable. ' + f'One of: {", ".join(KINDS)}. Defaults to all of them, which is ' + 'rarely what you want alongside --any-name', + ) + parser.add_argument( + '--include-unknown-age', action='store_true', + help='also sweep matches whose creation time the API did not report ' + '(skipped by default, since an unknown age can be shown neither ' + 'to be old enough nor to fall after --since)', + ) + parser.add_argument( + '--show-unmatched', action='store_true', + help='also list the live deployments this tool does not recognize as ' + "the suite's, without touching them. Run this when the " + 'organization looks full of strays but the sweep finds none: a ' + 'test that names a deployment outside the conventions in ' + 'PATTERNS is invisible here until its name is added', + ) + args = parser.parse_args(argv) + + kinds = args.kinds or list(KINDS) + + leftovers, spared, unmatched = find_leftovers( + args.older_than, args.include_unknown_age, + since=args.since, any_name=args.any_name, kinds=kinds, + ) + + if args.since is not None: + print( + 'Selecting {} created on or after {:%Y-%m-%d %H:%M %Z}, {}.\n' + .format( + '/'.join(kinds), + args.since, + 'any name' if args.any_name + else 'named like the test suite', + ), + ) + + if args.show_unmatched: + if unmatched: + print( + f'{len(unmatched)} live deployment(s) not recognized as the ' + "suite's, and so never swept:", + ) + for label in sorted(unmatched): + print(f' ? {label}') + print( + '\nIf one of these was made by a test, add its name to ' + 'PATTERNS in this module.\n', + ) + else: + print('Every live deployment is recognized by PATTERNS.\n') + + if spared: + print(f'{len(spared)} match(es) left alone by the age filter:') + for label in spared: + print(f' - {label}') + print() + + subject = 'deployment' if args.any_name else 'test deployment' + + if not leftovers: + print(f'No matching {subject}s found.') + return 0 + + print(f'{len(leftovers)} matching {subject}(s):') + for label, _ in leftovers: + print(f' - {label}') + + if not args.yes: + print('\nDry run; pass --yes to terminate these.') + return 0 + + from singlestoredb.tests import utils + + failed = 0 + for label, obj in leftovers: + try: + utils.terminate(obj) + except Exception as exc: + failed += 1 + print(f'✗ {label}: {exc}') + else: + print(f'✓ terminated {label}') + + return 1 if failed else 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/singlestoredb/tests/conftest.py b/singlestoredb/tests/conftest.py index 681c2fe54..9d426a647 100644 --- a/singlestoredb/tests/conftest.py +++ b/singlestoredb/tests/conftest.py @@ -29,7 +29,10 @@ import logging import os from collections.abc import Iterator +from typing import Any +from typing import List from typing import Optional +from typing import Tuple import pytest @@ -48,6 +51,18 @@ _container_manager: Optional[_TestContainerManager] = None +def _test_utils() -> Any: + """ + Return the test helper module. + + Imported through a function returning ``Any`` because ``tests/utils.py`` + carries a module-level ``# type: ignore``, which leaves mypy with no + attributes to check against. + """ + from singlestoredb.tests import utils + return utils + + def pytest_configure(config: pytest.Config) -> None: """ Pytest hook that runs before test collection. @@ -58,6 +73,11 @@ def pytest_configure(config: pytest.Config) -> None: """ global _container_manager + # Before any test module is imported: a setUpClass can create clusters, + # and they have to be tracked from the first one. + _test_utils().install_deployment_tracking() + _install_sweep_fallbacks() + # Prevent double initialization - pytest_configure can be called multiple times if _container_manager is not None: logger.debug('pytest_configure already called, skipping') @@ -134,14 +154,160 @@ def pytest_configure(config: pytest.Config) -> None: logger.debug(f'Using existing SINGLESTOREDB_URL={url}') +def pytest_runtest_setup(item: pytest.Item) -> None: + """ + Sweep the previous test class's deployments before the next one starts. + + This hook runs before pytest triggers ``setUpClass``, so by the time a + class begins the one before it is finished -- ``tearDownClass`` included, + or skipped because ``setUpClass`` raised. Sweeping here rather than only + at the end of the session means a leaked cluster is billed for one class, + not for the rest of the run. + """ + utils = _test_utils() + + try: + cls = getattr(item, 'cls', None) + module = getattr(item, 'module', None) + owner = '{}.{}'.format( + module.__name__ if module is not None else '', + cls.__name__ if cls is not None else '', + ) + except Exception: # pragma: no cover - non-python items + return + + if owner == utils.get_owner(): + return + + previous = utils.get_owner() + utils.set_owner(owner) + if previous: + _sweep_live_deployments(previous) + + +def _sweep_live_deployments(owner: Optional[str] = None) -> None: + """ + Terminate workspace groups, workspaces and clusters tests left behind. + + A class whose ``setUpClass`` raises never gets its ``tearDownClass``, so + the deployments it had already created would otherwise stay live -- and + billed -- indefinitely. Anything a test created is swept here whether or + not the test that made it ran to completion. + + A whole-session sweep -- ``owner is None``, so ``pytest_unconfigure``, + ``atexit`` or SIGTERM -- first recovers the creations still in progress. + Those have POSTed but are blocked waiting for the deployment to come up, so + nothing has tracked them yet, and killing the process here would leak them. + The per-class sweep skips that step: it runs between tests, where no + creation is in flight. + """ + try: + if owner is None: + _test_utils().recover_in_flight() + removed = _test_utils().cleanup_tracked(owner) + except Exception as exc: # pragma: no cover - shutdown path + print(f'\n✗ Failed to sweep leftover deployments: {exc}') + logger.error(f'Failed to sweep leftover deployments: {exc}') + return + + if removed: + print('\n' + '=' * 70) + print('Terminated deployments left behind by tests:') + for label in removed: + print(f' - {label}') + print('=' * 70) + logger.info(f'Swept {len(removed)} leftover deployment(s)') + + # A deployment still tracked after a full sweep is one the sweep could not + # terminate -- it is live and billing. Say so loudly rather than letting + # the run end quietly; `python -m singlestoredb.tests.cleanup_deployments` + # is the way to reap it. + if owner is None: + try: + stranded = _test_utils().tracked_labels() + except Exception: # pragma: no cover - shutdown path + return + if stranded: + print('\n' + '!' * 70) + print( + 'STILL LIVE -- these deployments could not be terminated and ' + 'are costing money:', + ) + for label in stranded: + print(f' - {label}') + print( + 'Reap them with: python -m singlestoredb.tests.' + 'cleanup_deployments --yes', + ) + print('!' * 70) + logger.error(f'{len(stranded)} deployment(s) left live') + + +#: Set once the atexit/signal fallbacks are in place, so a repeated +#: ``pytest_configure`` does not stack handlers. +_sweep_fallbacks_installed = False + + +def _install_sweep_fallbacks() -> None: + """ + Sweep leftover deployments even when ``pytest_unconfigure`` never runs. + + ``pytest_unconfigure`` is the normal path, but it is skipped whenever the + process does not shut down through pytest: a cancelled CI job, a killed + xdist worker holding the shared cluster pool, or an interpreter crash. The + pool is the expensive case -- it is attributed to owner ``''``, so no + per-class sweep ever touches it, and ``pytest_unconfigure`` is its only + scheduled cleanup. + + ``atexit`` covers ``sys.exit`` and an unhandled exception; a SIGTERM + handler covers the cancellation case, since Python does not run ``atexit`` + for a signal-terminated process. SIGKILL is unreachable by design -- that + is what ``cleanup_deployments.py`` is for. + + Both paths go through ``_sweep_live_deployments()`` with no owner, which + recovers the creations still waiting on their deployment before sweeping -- + a job cancelled mid ``wait_on_active`` is otherwise the one leak these + handlers cannot see. + """ + global _sweep_fallbacks_installed + if _sweep_fallbacks_installed: + return + _sweep_fallbacks_installed = True + + import atexit + import signal + + # Idempotent: a successful sweep empties the tracking list, so the normal + # path leaves these with nothing to do. + atexit.register(_sweep_live_deployments) + + previous = signal.getsignal(signal.SIGTERM) + + def on_sigterm(signum: int, frame: Any) -> None: + _sweep_live_deployments() + if callable(previous): + previous(signum, frame) + elif previous == signal.SIG_DFL: + signal.signal(signal.SIGTERM, signal.SIG_DFL) + os.kill(os.getpid(), signum) + + try: + signal.signal(signal.SIGTERM, on_sigterm) + except ValueError: # pragma: no cover - not the main thread + logger.debug('Not the main thread; no SIGTERM sweep installed') + + def pytest_unconfigure(config: pytest.Config) -> None: """ Pytest hook that runs after all tests complete. - Cleans up the Docker container if one was started. + Terminates any live deployment a test left behind, then cleans up the + Docker container if one was started. """ global _container_manager + _sweep_live_deployments() + if _container_manager is not None and not _container_manager.use_existing: print('\n' + '=' * 70) print('Cleaning up Docker container...') @@ -157,6 +323,126 @@ def pytest_unconfigure(config: pytest.Config) -> None: logger.error(f'Failed to stop Docker container: {e}') +#: Management API timings per test, collected when SINGLESTOREDB_MANAGEMENT_TRACE +#: is set. Kept here rather than on the config object so that +#: ``pytest_terminal_summary`` can read it without a fixture. +#: +#: Every test that ran under an active trace is in here, including the ones that +#: made no management call at all: ``trace_management_api_class`` subtracts this +#: list from the class total to get the fixture share, so a test missing from it +#: has its wall clock charged to ``setUpClass``. The event-less ones are +#: filtered out at report time by :func:`_traced` instead. +_management_traces: List[Tuple[str, Any]] = [] + +#: The same, for the class fixtures rather than the tests. Separate because the +#: two overlap: a class trace spans its tests as well as its fixtures, so the +#: fixture share is what is left after the tests' events are taken out of it. +_management_fixture_traces: List[Tuple[str, Any]] = [] + + +@pytest.fixture(autouse=True) +def trace_management_api(request: pytest.FixtureRequest) -> Iterator[None]: + """ + Record where each test's management API time went. + + Only active when ``SINGLESTOREDB_MANAGEMENT_TRACE`` is set, and reported by + :func:`pytest_terminal_summary`. The per-event stderr log the same variable + turns on is swallowed by pytest's capturing unless ``-s`` is given, so the + summary is written through the terminal reporter instead, which is always + shown. + """ + from singlestoredb.management import timing + + if not timing.logging_enabled(): + yield + return + + with timing.trace() as trace: + yield + + _management_traces.append((request.node.nodeid, trace)) + + +@pytest.fixture(scope='class', autouse=True) +def trace_management_api_class(request: pytest.FixtureRequest) -> Iterator[None]: + """ + Record what a class's ``setUpClass``/``tearDownClass`` cost. + + :func:`trace_management_api` is function-scoped, so it opens after + ``setUpClass`` has already run and closes before ``tearDownClass`` -- which + made the most expensive management calls in the suite invisible. The + fixtures here deploy the clusters and workspace groups the tests share, so + a run could report 5592 traced seconds out of 10125 and only three + ``POST clusters``. + + This trace spans the whole class, tests included, and + :func:`timing.Trace.of` subtracts the tests back out: the events are the + same objects in both traces, since :func:`timing._emit` hands each one to + every trace active in the context, so identity separates them exactly. + """ + from singlestoredb.management import timing + + if not timing.logging_enabled(): + yield + return + + # The tests of this class are the ones appended from here on. + first_test = len(_management_traces) + with timing.trace() as trace: + yield + + tests = [x[1] for x in _management_traces[first_test:]] + in_a_test = {id(x) for one in tests for x in one.events} + fixtures = timing.Trace.of( + [x for x in trace.events if id(x) not in in_a_test], + trace.elapsed - sum(x.elapsed for x in tests), + ) + if fixtures.events: + _management_fixture_traces.append((request.node.nodeid, fixtures)) + + +def _traced(traces: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]: + """ + Drop the traces that recorded no management call. + + ``_management_traces`` holds every test so that the fixture arithmetic is + right, but a test that made no management call has nothing to report and its + wall clock would inflate the combined total. + """ + return [x for x in traces if x[1].events] + + +def pytest_terminal_summary(terminalreporter: Any) -> None: + """Report the management API time the run spent, if it was traced.""" + tests = _traced(_management_traces) + fixtures = _traced(_management_fixture_traces) + if not tests and not fixtures: + return + + from singlestoredb.management import timing + + terminalreporter.write_sep('=', 'management API timing') + combined = timing.Trace.combine(x[1] for x in tests + fixtures) + terminalreporter.write_line(combined.summary()) + + for heading, traces in ( + ('slowest traced tests', tests), + ('slowest traced class fixtures', fixtures), + ): + if not traces: + continue + slowest = sorted(traces, key=lambda x: x[1].elapsed, reverse=True)[:10] + terminalreporter.write_line('') + terminalreporter.write_line(f' {heading}') + for nodeid, trace in slowest: + terminalreporter.write_line( + ' {:>8.3f}s requests={:>7.3f}s waiting={:>8.3f}s {}'.format( + trace.elapsed, trace.total(timing.REQUEST), + trace.total(timing.WAIT), nodeid, + ), + ) + + @pytest.fixture(scope='session', autouse=True) def setup_test_environment() -> Iterator[None]: """ diff --git a/singlestoredb/tests/test_fusion.py b/singlestoredb/tests/test_fusion.py index 21a15beaa..248259dc0 100644 --- a/singlestoredb/tests/test_fusion.py +++ b/singlestoredb/tests/test_fusion.py @@ -9,10 +9,12 @@ import unittest from typing import Any from typing import List +from typing import Tuple import pytest import singlestoredb as s2 +from singlestoredb.management import timing from singlestoredb.tests import utils @@ -90,9 +92,867 @@ def test_show_grammar(self): assert cmds assert [x for x in cmds if x.strip().startswith('CREATE WORKSPACE')], cmds + def test_cluster_commands_registered(self): + from singlestoredb.fusion import registry + + want = { + 'SHOW CLUSTERS', 'SHOW CLUSTER REGIONS', 'SHOW PROJECTS', + 'CREATE CLUSTER', 'DROP CLUSTER', 'SUSPEND CLUSTER', + 'RESUME CLUSTER', 'USE CLUSTER', 'SHOW STARTER CLUSTERS', + 'SHOW STARTER CLUSTER REGIONS', + 'CREATE STARTER CLUSTER', 'DROP STARTER CLUSTER', + } + missing = want - set(registry._handlers) + assert not missing, missing + + def test_show_cluster_status_is_not_shadowed(self): + """ + ``SHOW CLUSTER STATUS`` must reach the engine, not Fusion. + + The registry matches the longest key first, so registering a bare + two-word ``SHOW CLUSTER`` would swallow the engine's own + ``SHOW CLUSTER STATUS``. That is why the region command is spelled + ``SHOW CLUSTER REGIONS``. + """ + from singlestoredb.fusion import registry + + assert registry.get_handler('SHOW CLUSTER STATUS') is None + assert registry.get_handler('SHOW CLUSTERS') is not None + assert registry.get_handler('SHOW CLUSTER REGIONS') is not None + + def test_in_region_is_matched_without_regard_to_case(self): + """ + ``IN REGION`` must match a region whatever case it is written in. + + A miss is not an error -- the literal is passed through for the API to + rule on -- but the provider is only ever recovered *from* a match, so a + case-sensitive comparison would quietly post a region with no provider + alongside it. Both spellings of the name, and either case, must match + and must come back in the API's own spelling. + """ + from unittest import mock + + from singlestoredb.fusion.handlers import cluster as handlers + from singlestoredb.management.region import Region + + regions = [ + Region( + name='US East 1 (N. Virginia)', + provider='AWS', region_name='us-east-1', + ), + ] + manager = mock.MagicMock() + type(manager).regions = mock.PropertyMock(return_value=regions) + + want = dict(provider='AWS', region='us-east-1') + for written in ( + 'us-east-1', 'US-EAST-1', 'Us-East-1', + 'US East 1 (N. Virginia)', 'us east 1 (n. virginia)', + ): + params = {'in_region': {'region_name': written}} + got = handlers._resolve_region(params, manager) + assert got == want, (written, got) + + # A provider is matched without regard to case too, and narrows an + # otherwise ambiguous name. + for written in ('AWS', 'aws', 'Aws'): + params = { + 'in_region': {'region_name': 'US-EAST-1'}, + 'using_provider': written, + } + got = handlers._resolve_region(params, manager) + assert got == want, (written, got) + + # An unknown region is still passed through untouched. + params = {'in_region': {'region_name': 'mars-north-1'}} + assert handlers._resolve_region(params, manager) == \ + dict(provider=None, region='mars-north-1') + + # An ambiguous name reports every candidate rather than picking one. + regions.append( + Region( + name='US East 1 (N. Virginia)', + provider='GCP', region_name='us-east1', + ), + ) + params = {'in_region': {'region_name': 'us east 1 (n. virginia)'}} + with self.assertRaises(ValueError) as cm: + handlers._resolve_region(params, manager) + assert 'more than one region matches' in str(cm.exception), cm.exception + + def test_starter_cluster_regions_uses_the_shared_tier_list(self): + """ + ``SHOW STARTER CLUSTER REGIONS`` must not report every region. + + The shared-tier route accepts only the regions + ``ClusterManager.shared_tier_regions`` reports, which are a subset of + ``ClusterManager.regions``. Listing the latter would offer regions + that ``CREATE STARTER CLUSTER`` then rejects, which is the mistake + this command exists to prevent. + """ + from unittest import mock + + from singlestoredb.fusion.handlers import cluster as handlers + from singlestoredb.management.region import Region + + shared = [ + Region( + name='US East 1 (N. Virginia)', + provider='AWS', region_name='us-east-1', + ), + ] + every = shared + [ + Region( + name='US East 2 (Ohio)', + provider='AWS', region_name='us-east-2', + ), + ] + + manager = mock.MagicMock() + type(manager).shared_tier_regions = mock.PropertyMock( + return_value=shared, + ) + type(manager).regions = mock.PropertyMock(return_value=every) + + with mock.patch.object( + handlers, 'get_cluster_manager', return_value=manager, + ): + handler = handlers.ShowStarterClusterRegionsHandler(self.conn) + handler.compile() + res = handler.execute('SHOW STARTER CLUSTER REGIONS;') + + assert [x[0] for x in res.description] == \ + ['Name', 'Provider', 'RegionName'], res.description + assert [tuple(x) for x in res.rows] == \ + [('US East 1 (N. Virginia)', 'AWS', 'us-east-1')], res.rows + + def test_create_cluster_grammar(self): + from singlestoredb.fusion import registry + + self.cur.execute('show fusion grammar for "create cluster"') + cmds = [x[0] for x in self.cur.fetchall()] + assert cmds + assert [x for x in cmds if x.strip().startswith('CREATE CLUSTER')], cmds + + # Assert against the rendered clause list rather than the output of + # SHOW FUSION GRAMMAR, which also carries the prose remarks -- and + # those *mention* the absent clauses in order to explain the absence. + handler = registry._handlers['CREATE CLUSTER'] + handler.compile() + syntax = handler.syntax + + # v2 assigns no region IDs, so there is no ID alternate to offer. + assert '' not in syntax, syntax + assert '' in syntax, syntax + + # Dropped at v2 (audit item 14). A clause for any of these would + # parse, be sent, and be silently discarded. + assert 'KMS' not in syntax.upper(), syntax + assert 'SMART DR' not in syntax.upper(), syntax + assert 'PASSWORD' not in syntax.upper(), syntax + + def test_create_workspace_group_grammar_still_has_region_id(self): + """The v1 command keeps its region-ID alternate; v2 never had one.""" + from singlestoredb.fusion import registry + + handler = registry._handlers['CREATE WORKSPACE GROUP'] + handler.compile() + syntax = handler.syntax + assert '' in syntax, syntax + assert 'KMS' in syntax.upper(), syntax + + def test_v1_workspace_commands_are_deprecated(self): + """ + Every v1 WORKSPACE command points at its v2 CLUSTER replacement. + + No exceptions: every command in the module reads the v1 API, so every + one of them warns. ``SHOW REGIONS`` is the loosest pairing -- v2 assigns + no region IDs, so ``SHOW CLUSTER REGIONS`` reports ``RegionName`` where + it reports ``ID`` -- but it is still where a caller has to go. Asserted + so that adding a v1 command without a pointer fails here. + """ + from singlestoredb.fusion import registry + + undeprecated = set() + for key, handler in registry._handlers.items(): + if not handler.__module__.endswith('.workspace'): + continue + if handler._deprecated_by: + # The replacement must be a real command, not a typo. + assert handler._deprecated_by in registry._handlers, \ + (key, handler._deprecated_by) + else: + undeprecated.add(key) + + assert not undeprecated, undeprecated + + def test_v2_cluster_commands_are_not_deprecated(self): + """The replacements must not themselves warn.""" + from singlestoredb.fusion import registry + + for key, handler in registry._handlers.items(): + if handler.__module__.endswith('.cluster'): + assert not handler._deprecated_by, key + + def test_deprecation_warning_fires_on_execute(self): + """ + ``_deprecated_by`` warns, names the command, and still runs. + + Driven through a probe handler rather than a real ``WORKSPACE`` command + so the assertion needs no management API token: what is under test is + the mechanism in ``SQLHandler.execute``, not any one command's body. + """ + from singlestoredb.fusion.handler import SQLHandler + from singlestoredb.warnings import DeprecatedFeatureWarning + + class _DeprecatedProbeHandler(SQLHandler): + """ + SHOW FUSION DEPRECATION PROBE; + + """ + + _deprecated_by = 'SHOW CLUSTERS' + + def run(self, params): + return None + + # Deliberately not registered -- execute() only needs the class. + handler = _DeprecatedProbeHandler(self.conn) + + with self.assertWarns(DeprecatedFeatureWarning) as caught: + res = handler.execute('SHOW FUSION DEPRECATION PROBE') + + msg = str(caught.warning) + assert 'SHOW FUSION DEPRECATION PROBE' in msg, msg + assert 'SHOW CLUSTERS' in msg, msg + # Deprecated, not removed: the command still returns a result. + assert res is not None + + def test_no_deprecation_warning_by_default(self): + """A command without ``_deprecated_by`` stays silent.""" + import warnings + + from singlestoredb.warnings import DeprecatedFeatureWarning + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + self.cur.execute('show fusion commands') + self.cur.fetchall() + + assert not [ + x for x in caught + if issubclass(x.category, DeprecatedFeatureWarning) + ], [str(x.message) for x in caught] + + def test_maximal_create_cluster_parses(self): + from singlestoredb.fusion import registry + + handler = registry._handlers['CREATE CLUSTER'] + handler.compile() + + sql = ( + "CREATE CLUSTER IF NOT EXISTS 'fusion-parse-test' " + "IN REGION 'us-east-1' USING PROVIDER 'AWS' " + "IN PROJECT 'Some Project' " + # The /* ... */ is matched by the `ws*` tail of the `number` rule, + # so it lands inside the number node -- visit_number must read the + # regex match, not the whole node's text. + "WITH SIZE 'S-00' USING SCALE FACTOR 1 /* scale comment */ " + 'AUTO SUSPEND AFTER 30 MINUTES WITH TYPE IDLE ' + 'ENABLE KAI WITH CACHE CONFIG 2 ' + "WITH FIREWALL RANGES '0.0.0.0/0' ALLOW ALL TRAFFIC " + "WITH UPDATE WINDOW '3:5' EXPIRES AT '1h' " + 'WAIT ON ACTIVE' + ) + + inst = handler.__new__(handler) + inst.connection = None + inst._handled = set() + params = inst.visit(handler.grammar.parse(sql)) + for key, value in list(params.items()): + params[key] = inst.validate_rule(key, value) + + assert params['cluster_name'] == 'fusion-parse-test' + assert params['in_region'] == {'region_name': 'us-east-1'} + assert params['using_provider'] == 'AWS' + assert params['in_project'] == {'project_name': 'Some Project'} + # must accept a bare integer, not only 1.0 + assert params['using_scale_factor'] == 1.0 + # The clause is one flat dict, not a list of one dict per sub-rule. + assert params['auto_suspend'] == dict( + suspend_after_value=30, + suspend_after_units='MINUTES', + suspend_type='IDLE', + ) + assert params['with_update_window'] == '3:5' + assert params['wait_on_active'] is True + + def test_create_cluster_has_no_v2_only_clauses(self): + """ + The grammar stops at what the v1 pair exposes. + + ``deploymentType`` and ``multiAZ`` have no ``CREATE WORKSPACE`` or + ``CREATE WORKSPACE GROUP`` counterpart, so they are reachable only + through ``ClusterManager.create_cluster``. Asserted rather than left + implicit: re-adding a clause is a deliberate widening of the SQL + surface, not a detail of the handler. + """ + from singlestoredb.fusion import registry + + handler = registry._handlers['CREATE CLUSTER'] + handler.compile() + syntax = handler.syntax.upper() + assert 'DEPLOYMENT TYPE' not in syntax, syntax + assert 'MULTI AZ' not in syntax, syntax + # ... while the clauses the v1 commands do have are still here. + for clause in ('ENABLE KAI', 'UPDATE WINDOW', 'CACHE CONFIG'): + assert clause in syntax, (clause, syntax) + + def test_create_cluster_rejects_region_id(self): + from singlestoredb.fusion import registry + + handler = registry._handlers['CREATE CLUSTER'] + handler.compile() + inst = handler.__new__(handler) + inst.connection = None + inst._handled = set() + + with self.assertRaises(Exception): + inst.visit( + handler.grammar.parse( + "CREATE CLUSTER 'c' IN REGION ID 'some-region-id'", + ), + ) + + def test_stage_handlers_name_a_deployment_with_a_bare_in(self): + """ + All six Stage handlers take a bare ``IN``, and no ``IN CLUSTER``. + + A deployment is named the same way whatever kind it is, so there is + nothing for a qualified spelling to disambiguate -- ``IN CLUSTER`` would + resolve exactly where the bare ``IN`` already does. ``IN GROUP`` stays + because it names something else: a v1 workspace group, which is why it + carries its own ``group_id``/``group_name`` placeholders rather than the + ``deployment_*`` ones. + """ + from singlestoredb.fusion import registry + from singlestoredb.fusion.handler import SQLHandler + from singlestoredb.fusion.handlers import stage + + handlers = [ + x for x in vars(stage).values() + if isinstance(x, type) + and issubclass(x, SQLHandler) and x is not SQLHandler + ] + assert len(handlers) == 6, [x.__name__ for x in handlers] + + for cls in handlers: + cls.compile() + grammar = cls._grammar + assert 'IN CLUSTER' not in grammar, cls.__name__ + assert 'in_group = IN GROUP' in grammar, cls.__name__ + # IN GROUP carries the group placeholders, so a group name never + # reaches the deployment lookup. + assert 'in_group = IN GROUP { group_id | group_name }' in grammar, \ + cls.__name__ + # in_group must precede the bare in_deployment in the alternation, + # or IN would win before GROUP is considered and IN GROUP 'x' would + # parse as a deployment named GROUP. + alternation = 'in = { in_group | in_deployment }' + assert alternation in grammar, cls.__name__ + + # SHOW STAGE FILES is representative; the clause is identical on all six. + cls = registry._handlers['SHOW STAGE FILES'] + cls.compile() + for sql, key in [ + ("SHOW STAGE FILES IN GROUP 'g1'", 'in_group'), + ("SHOW STAGE FILES IN GROUP ID 'abc'", 'in_group'), + ("SHOW STAGE FILES IN 'd1'", 'in_deployment'), + ("SHOW STAGE FILES IN ID 'abc'", 'in_deployment'), + ]: + inst = cls.__new__(cls) + inst.connection = None + inst._handled = set() + params = inst.visit(cls.grammar.parse(sql)) + assert key in params['in'], (sql, params['in']) + + # IN CLUSTER no longer parses at all. It must not quietly become a + # deployment named CLUSTER, which is what dropping in_cluster from the + # alternation would do if CLUSTER were a valid . + with self.assertRaises(Exception): + cls.grammar.parse("SHOW STAGE FILES IN CLUSTER 'c1'") + + def test_fusion_managers_are_version_pinned(self): + """ + Each Fusion manager names its version rather than following the option. + + The option is an org-wide preference; a handler that *is* one version's + vocabulary has nothing to learn from it. + """ + import inspect + + from singlestoredb.fusion.handlers import utils + + assert '_manage_workspaces_v1()' in inspect.getsource( + utils.get_workspace_manager, + ) + for func in (utils.get_cluster_manager, utils.get_files_manager): + assert "version='v2'" in inspect.getsource(func), func.__name__ + + def test_get_deployment_resolves_a_bare_in_against_v2(self): + """ + The deployment lookup is v2 only; the group lookup is v1 only. + + Keeping them in separate functions is what makes the order of the two + enforceable -- clusters first, group second -- so the split is asserted + here rather than only through behaviour. + """ + import inspect + + from singlestoredb.fusion.handlers import utils + + src = inspect.getsource(utils.get_deployment) + assert 'workspace_groups' not in src + assert 'clusters' in src + + src = inspect.getsource(utils._workspace_group) + assert 'workspace_groups' in src + assert 'clusters' not in src + + def test_job_commands_use_the_cluster_manager(self): + """ + JOB commands are not v1 vocabulary. + + A job runs against a deployment, and at v2 a deployment is a cluster, + so routing them through the v1 manager gave every scheduled job a v1 + ``targetType``. + """ + import inspect + + from singlestoredb.fusion.handlers import job + + src = inspect.getsource(job) + assert 'get_workspace_manager' not in src + assert src.count('get_cluster_manager().organizations.current.jobs') == 8 + + #: Minimal ``POST jobs`` response, enough for ``Job.from_dict``. The + #: handlers read nothing but ``jobID`` off it. + _JOB_RESPONSE = { + 'completedExecutionsCount': 0, + 'createdAt': '2026-09-14T00:00:00Z', + 'enqueuedBy': 'someone@example.com', + 'executionConfig': { + 'createSnapshot': False, + 'maxAllowedExecutionDurationInMinutes': 0, + 'notebookPath': 'nb.ipynb', + }, + 'jobID': 'job-1', + 'jobMetadata': [], + 'schedule': {'mode': 'Once'}, + } + + def _job_request_body(self, sql, **env): + """ + Return the ``POST jobs`` body a JOB statement produces. + + The manager is mocked at the request layer rather than replaced, so the + body is the one ``JobsManager`` really assembles -- which is the whole + point: the target is not in the statement, so this is the only place it + can be observed. + """ + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion import registry + from singlestoredb.fusion.handlers import job as job_handlers + from singlestoredb.management.job import JobsManager + + api = MagicMock() + api._post.return_value.json.return_value = self._JOB_RESPONSE + + clusters = MagicMock() + clusters.organizations.current.jobs = JobsManager(api) + + handler = registry.get_handler(sql) + assert handler is not None, sql + inst = handler.__new__(handler) + inst.connection = None + inst._handled = set() + params = inst.visit(handler.grammar.parse(sql)) + for key, value in list(params.items()): + params[key] = inst.validate_rule(key, value) + + with patch.object( + job_handlers, 'get_cluster_manager', return_value=clusters, + ): + self._fusion_env(**env) + inst.run(params) + + route, kwargs = api._post.call_args + assert route == ('jobs',), route + return kwargs['json'] + + def test_job_commands_target_the_environment_deployment(self): + """ + Nothing in the JOB grammar names a deployment, so pin what does. + + ``targetID`` comes from the environment and ``targetType`` from the + manager the handler picked, and no statement can assert either -- the + pairing is only visible in the request body. There is no live coverage + of a workspace target left either: ``TestJobsFusion`` targets a v2 + cluster, so this is what holds the write-path vocabulary in place. + """ + from singlestoredb.management.job import TargetType + + cluster_id = '11111111-1111-4111-8111-111111111111' + starter_id = '22222222-2222-4222-8222-222222222222' + + body = self._job_request_body( + "RUN JOB USING NOTEBOOK 'nb.ipynb' " + "WITH RUNTIME 'notebooks-cpu-small'", + SINGLESTOREDB_WORKSPACE=cluster_id, + SINGLESTOREDB_DEFAULT_DATABASE='dbtest', + ) + assert body['targetConfig'] == dict( + databaseName='dbtest', + targetID=cluster_id, + targetType=TargetType.CLUSTER.value, + ) + assert body['schedule']['mode'] == 'Once' + assert body['executionConfig']['notebookPath'] == 'nb.ipynb' + assert body['executionConfig']['runtimeName'] == 'notebooks-cpu-small' + + # A starter deployment is named by its own variable, wins over the + # regular one, and takes the matching targetType. + body = self._job_request_body( + "RUN JOB USING NOTEBOOK 'nb.ipynb'", + SINGLESTOREDB_WORKSPACE=cluster_id, + SINGLESTOREDB_VIRTUAL_WORKSPACE=starter_id, + SINGLESTOREDB_DEFAULT_DATABASE='dbtest', + ) + assert body['targetConfig']['targetID'] == starter_id + assert body['targetConfig']['targetType'] == \ + TargetType.VIRTUAL_CLUSTER.value + + # SCHEDULE JOB assembles the same target, and is the only one of the + # two that can carry RESUME TARGET. + body = self._job_request_body( + "SCHEDULE JOB USING NOTEBOOK 'nb.ipynb' WITH MODE 'Recurring' " + 'EXECUTE EVERY 2 HOURS RESUME TARGET', + SINGLESTOREDB_WORKSPACE=cluster_id, + SINGLESTOREDB_DEFAULT_DATABASE='dbtest', + ) + assert body['targetConfig'] == dict( + databaseName='dbtest', + resumeTarget=True, + targetID=cluster_id, + targetType=TargetType.CLUSTER.value, + ) + assert body['schedule'] == dict( + mode='Recurring', executionIntervalInMinutes=120, + ) + + # The whole targetConfig hangs off the database variable: without it + # the job is submitted with no target at all, whatever deployment the + # environment names. + body = self._job_request_body( + "RUN JOB USING NOTEBOOK 'nb.ipynb'", + SINGLESTOREDB_WORKSPACE=cluster_id, + ) + assert 'targetConfig' not in body + + def _fusion_env(self, **values): + """Run with only the deployment variables in ``values`` set.""" + from unittest.mock import patch + + ctx = patch.dict(os.environ) + ctx.start() + self.addCleanup(ctx.stop) + for name in ( + 'SINGLESTOREDB_WORKSPACE', + 'SINGLESTOREDB_VIRTUAL_WORKSPACE', + 'SINGLESTOREDB_WORKSPACE_GROUP', + 'SINGLESTOREDB_PROJECT', + 'SINGLESTOREDB_DEFAULT_DATABASE', + ): + os.environ.pop(name, None) + os.environ.update(values) + + def test_project_resolves_the_clause_and_nothing_else(self): + """ + ``IN PROJECT`` is the only thing ``get_project`` reads. + + Absent the clause it returns ``None``, leaving the choice to + ``ClusterManager._resolve_project_id``, which reads the project off the + current deployment. ``SINGLESTOREDB_PROJECT`` is not consulted: it names + an inference API project, so resolving it here turned every notebook's + ``CREATE CLUSTER`` into a 404. + """ + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + + project_id = '11111111-1111-4111-8111-111111111111' + by_name = MagicMock() + by_name.name = 'My Project' + manager = MagicMock() + manager.projects = [by_name] + + with patch.object(utils, 'get_cluster_manager', return_value=manager): + self._fusion_env() + assert utils.get_project( + dict(in_project=dict(project_id=project_id)), + ) is manager.get_project.return_value + manager.get_project.assert_called_once_with(project_id) + + assert utils.get_project( + dict(in_project=dict(project_name='My Project')), + ) is by_name + + assert utils.get_project({}) is None + + # Still None with the environment variable set. + self._fusion_env(SINGLESTOREDB_PROJECT=project_id) + assert utils.get_project({}) is None + + def test_deployment_refuses_the_group_environment_variable(self): + """ + ``SINGLESTOREDB_WORKSPACE_GROUP`` holds a group ID, not a cluster ID. + + v2 reports the group only as ``Cluster.group`` and has no route to + look it up, so guessing which cluster was meant could target the wrong + deployment. + """ + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + + with patch.object(utils, 'get_cluster_manager', return_value=MagicMock()): + self._fusion_env( + SINGLESTOREDB_WORKSPACE_GROUP='11111111-1111-4111-8111-111111111111', + ) + with self.assertRaises(KeyError) as cm: + utils.get_deployment({}) + + msg = str(cm.exception) + assert 'SINGLESTOREDB_WORKSPACE_GROUP' in msg + assert 'SINGLESTOREDB_WORKSPACE' in msg + + def test_in_group_resolves_a_workspace_group_against_v1(self): + """ + ``IN GROUP`` names a v1 workspace group, by name and by ID. + + Stage is attached to the group itself at v1, so a group names a Stage on + its own. The cluster manager must not be touched at all: a group ID is + not a cluster ID, and looking one up as the other is what made this + spelling miss. + """ + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + from singlestoredb.warnings import DeprecatedFeatureWarning + + group_id = '11111111-1111-4111-8111-111111111111' + group = MagicMock() + group.id = group_id + group.name = 'wsg1' + + v1 = MagicMock() + v1.workspace_groups = [group] + v1.get_workspace_group.return_value = group + clusters = MagicMock() + + def resolve(params): + with patch.object(utils, 'get_workspace_manager', return_value=v1), \ + patch.object( + utils, 'get_cluster_manager', return_value=clusters, + ): + self._fusion_env() + with self.assertWarns(DeprecatedFeatureWarning): + return utils.get_deployment(params) + + for params in ( + dict(group=dict(group_name='wsg1')), + {'in': dict(in_group=dict(group_name='wsg1'))}, + dict(group=dict(group_id=group_id)), + {'in': dict(in_group=dict(group_id=group_id))}, + ): + assert resolve(params) is group, params + + clusters.assert_not_called() + assert not clusters.method_calls, clusters.method_calls + + def test_in_group_falls_back_to_a_starter_workspace(self): + """ + A name or ID that is no group's is tried as a starter workspace. + + A starter workspace owns its Stage the same way a group does and was + reachable through this spelling before, so it stays reachable. + """ + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + from singlestoredb.warnings import DeprecatedFeatureWarning + + starter_id = '22222222-2222-4222-8222-222222222222' + starter = MagicMock() + starter.id = starter_id + starter.name = 'starter1' + + v1 = MagicMock() + v1.workspace_groups = [] + v1.starter_workspaces = [starter] + v1.get_workspace_group.side_effect = s2.ManagementError(errno=404) + v1.get_starter_workspace.return_value = starter + + def resolve(params): + with patch.object(utils, 'get_workspace_manager', return_value=v1): + self._fusion_env() + with self.assertWarns(DeprecatedFeatureWarning): + return utils.get_deployment(params) + + for params in ( + {'in': dict(in_group=dict(group_name='starter1'))}, + {'in': dict(in_group=dict(group_id=starter_id))}, + ): + assert resolve(params) is starter, params + + def test_bare_in_falls_back_to_a_workspace_group(self): + """ + A bare ``IN`` that matches no cluster is tried as a workspace group. + + A bare ``IN`` named a workspace group before the Stage commands moved to + v2 -- a group was the only kind of Stage owner there was -- so a + statement written then keeps working, and silently: ``IN`` is the + spelling to use for either kind of owner, so there is nothing about the + statement to warn about. Only ``IN GROUP`` warns. + """ + import warnings + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + + group_id = '11111111-1111-4111-8111-111111111111' + group = MagicMock() + group.id = group_id + group.name = 'wsg1' + + v1 = MagicMock() + v1.workspace_groups = [group] + v1.starter_workspaces = [] + v1.get_workspace_group.return_value = group + + clusters = MagicMock() + clusters.clusters = [] + clusters.starter_clusters = [] + clusters.get_cluster.side_effect = s2.ManagementError(errno=404) + clusters.get_starter_cluster.side_effect = s2.ManagementError(errno=404) + + def resolve(params): + with patch.object(utils, 'get_workspace_manager', return_value=v1), \ + patch.object( + utils, 'get_cluster_manager', return_value=clusters, + ): + self._fusion_env() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + return utils.get_deployment(params), caught + + for params in ( + dict(deployment_name='wsg1'), + {'in': dict(in_deployment=dict(deployment_name='wsg1'))}, + {'in': dict(in_deployment=dict(deployment_id=group_id))}, + ): + found, caught = resolve(params) + assert found is group, params + assert not caught, (params, [str(x.message) for x in caught]) + + def test_bare_in_prefers_a_cluster_over_a_group_of_the_same_name(self): + """ + The fallback is second, so nothing that resolves today changes meaning. + + A name that is both a cluster's and a workspace group's has to stay the + cluster's, and quietly: the fallback was not reached, so there is + nothing deprecated about the statement. + """ + import warnings + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + + cluster = MagicMock() + cluster.name = 'shared-name' + clusters = MagicMock() + clusters.clusters = [cluster] + + v1 = MagicMock() + + with patch.object(utils, 'get_workspace_manager', return_value=v1), \ + patch.object( + utils, 'get_cluster_manager', return_value=clusters, + ): + self._fusion_env() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + found = utils.get_deployment( + dict(deployment_name='shared-name'), + ) + + assert found is cluster + assert not caught, [str(x.message) for x in caught] + # The v1 manager must not even be built: the fallback is the only thing + # that needs it, and it was not reached. + assert not v1.method_calls, v1.method_calls + + def test_in_group_miss_names_the_workspace_group(self): + """A miss says what was looked for, not what a cluster would be.""" + import warnings + from unittest.mock import MagicMock + from unittest.mock import patch + + from singlestoredb.fusion.handlers import utils + + group_id = '11111111-1111-4111-8111-111111111111' + v1 = MagicMock() + v1.workspace_groups = [] + v1.starter_workspaces = [] + v1.get_workspace_group.side_effect = s2.ManagementError(errno=404) + v1.get_starter_workspace.side_effect = s2.ManagementError(errno=404) + + def message(params): + with patch.object(utils, 'get_workspace_manager', return_value=v1): + self._fusion_env() + with self.assertRaises(KeyError) as cm: + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + utils.get_deployment(params) + return str(cm.exception) + + for params, needle in ( + ({'in': dict(in_group=dict(group_name='wsg1'))}, 'wsg1'), + ({'in': dict(in_group=dict(group_id=group_id))}, group_id), + ): + msg = message(params) + assert needle in msg, msg + assert 'workspace group' in msg, msg + @pytest.mark.management +@pytest.mark.management_v1 class TestWorkspaceFusion(unittest.TestCase): + """ + The WORKSPACE and WORKSPACE GROUP grammar, which is the v1 vocabulary. + + Marked ``management_v1`` so it switches off with the rest of the v1 + coverage; ``TestClusterFusion*`` is the v2 replacement. The grammar + itself, and therefore this suite, goes away with ``management/v1/``. + """ id: str = secrets.token_hex(8) dbname: str = '' @@ -103,9 +963,12 @@ class TestWorkspaceFusion(unittest.TestCase): def setUpClass(cls): sql_file = os.path.join(os.path.dirname(__file__), 'test.sql') cls.dbname, cls.dbexisted = utils.load_sql(sql_file) - mgr = s2.manage_workspaces() + # Pinned: manage_workspaces() follows the management.version + # option, and Fusion is v1-only. + mgr = s2.manage_workspaces(version='v1') + # US-only: no test here asserts anything about these groups' regions, + # and creation in some non-US regions fails with a control-plane 500. us_regions = [x for x in mgr.regions if x.name.startswith('US')] - non_us_regions = [x for x in mgr.regions if not x.name.startswith('US')] wg = mgr.create_workspace_group( f'A Fusion Testing {cls.id}', region=random.choice(us_regions), @@ -120,7 +983,7 @@ def setUpClass(cls): cls.workspace_groups.append(wg) wg = mgr.create_workspace_group( f'C Fusion Testing {cls.id}', - region=random.choice(non_us_regions), + region=random.choice(us_regions), firewall_ranges=[], ) cls.workspace_groups.append(wg) @@ -158,6 +1021,28 @@ def tearDown(self): # traceback.print_exc() pass + def test_stage_in_group_addresses_the_workspace_group(self): + """ + ``SHOW STAGE FILES IN GROUP`` reaches a real v1 workspace group. + + The Stage commands themselves are v2, so this is the one clause of + theirs that belongs in this suite: it names a workspace group, whose + Stage lives at ``stage/{group_id}/fs/``. A freshly created group's Stage + is empty, which is enough to prove the route was reached -- a cluster + lookup would have raised instead. + """ + from singlestoredb.warnings import DeprecatedFeatureWarning + + wg = type(self).workspace_groups[0] + + for clause in [ + f"in group id '{wg.id}'", + f"in group '{wg.name}'", + ]: + with self.assertWarns(DeprecatedFeatureWarning): + self.cur.execute(f'show stage files {clause}') + assert len(list(self.cur)) == 0, clause + def test_show_regions(self): self.cur.execute('show regions') regs = list(self.cur) @@ -255,7 +1140,7 @@ def test_show_workspace_groups(self): assert names == [f'C Fusion Testing {self.id}', f'B Fusion Testing {self.id}'] def test_show_workspaces(self): - mgr = s2.manage_workspaces() + mgr = s2.manage_workspaces(version='v1') wg = mgr.workspace_groups[f'B Fusion Testing {self.id}'] self.cur.execute( @@ -271,20 +1156,26 @@ def test_show_workspaces(self): f'"B Fusion Testing {self.id}" with size S-00', ) - time.sleep(30) - iterations = 20 + # Wait for the three to be listed, not for them to be ACTIVE. Nothing + # below asserts a state value -- 'State' is checked as a column name, + # never for its contents -- so all this test needs is that SHOW + # WORKSPACES can see them. Requiring ACTIVE cost around 450 seconds a + # run for no assertion, and at a 30 second interval most of that was + # overshoot. Polled through timing.sleep so a traced run accounts for + # it; a bare time.sleep here was invisible to the tracer and landed in + # the unlabelled 'other' bucket. + wanted = ('show-ws-1', 'show-ws-2', 'show-ws-3') + deadline = time.time() + 600 while True: - wgs = wg.workspaces - states = [ - x.state for x in wgs - if x.name in ('show-ws-1', 'show-ws-2', 'show-ws-3') - ] - if len(states) == 3 and states.count('ACTIVE') == 3: + listed = [x.name for x in wg.workspaces if x.name in wanted] + if len(listed) == 3: break - iterations -= 1 - if not iterations: - raise RuntimeError('timed out waiting for workspaces to start') - time.sleep(30) + if time.time() >= deadline: + raise RuntimeError( + 'timed out waiting for workspaces to be listed; ' + f'saw {sorted(listed)}', + ) + timing.sleep(5, 'workspace listed') # SHOW self.cur.execute(f'show workspaces in group "B Fusion Testing {self.id}"') @@ -370,7 +1261,7 @@ def test_show_workspaces(self): assert names == ['show-ws-3', 'show-ws-2'] def test_create_drop_workspace(self): - mgr = s2.manage_workspaces() + mgr = s2.manage_workspaces(version='v1') wg = mgr.workspace_groups[f'A Fusion Testing {self.id}'] self.cur.execute( @@ -434,10 +1325,14 @@ def _wait_workspace_group_gone(self, mgr, wg_name, timeout=60, interval=2): time.sleep(interval) def test_create_drop_workspace_group(self): - mgr = s2.manage_workspaces() + mgr = s2.manage_workspaces(version='v1') reg = [x for x in mgr.regions if x.name.startswith('US')][0] - wg_name = f'Create WG Test {id(self)}' + # Random, not id(self): an address repeats across processes, so two + # workers running this test could pick the same name, and it reads + # nothing like a generated name to anyone looking at the organization. + # Whatever this is, it has to keep matching cleanup_deployments. + wg_name = f'Create WG Test {secrets.token_hex(8)}' try: self.cur.execute( @@ -480,41 +1375,575 @@ def test_create_drop_workspace_group(self): self.cur.execute(f'drop workspace group if exists id {wg_id}') finally: + # Only what is still live: the body drops the group itself, and a + # terminated record can still be listed for a while afterwards. + # Failures are reported rather than swallowed -- that is the + # difference between a group that went away and one still billing. + for wg in [ + x for x in mgr.workspace_groups + if x.name == wg_name and x.terminated_at is None + ]: + try: + wg.terminate(force=True) + except Exception as exc: + print( + f'Could not terminate workspace group {wg_name!r}; ' + f'it may still be live: {exc}', + ) + + +class _ClusterFusionMixin: + """ + Plumbing shared by the CLUSTER fusion suites. + + These are the v2 mirror of :class:`TestWorkspaceFusion`, flat rather than + nested. A cluster is created in one statement where a workspace needed + two, so there is no group fixture and no ``IN GROUP`` clause anywhere. + Names are lowercase and hyphenated because ``POST /v2/clusters`` enforces + ``[a-z0-9]([a-z0-9-]*[a-z0-9])?`` at 1-32 characters (audit item 7) -- + the spaced names the v1 suite uses are rejected. + + This was one class deploying three clusters in ``setUpClass``, which every + test then waited out whether or not it touched a cluster: the two + lifecycle tests deploy their own and the region, project and grammar + tests need none at all, yet all of them paid for three. The classes below + declare what they need in :attr:`fixture_prefixes` instead, so the + cluster-less ones start immediately and no class deploys more than it + reads. + + Not a ``TestCase``, and named with a leading underscore: pytest collects + any ``Test``-prefixed ``TestCase`` subclass it can reach, so a base that + was either would run every inherited test a second time under a fixture + of its own. + """ + + #: Prefixes of the shared clusters to deploy before this class's tests, + #: named ``-fusion-cluster-``. Empty means the class needs no + #: deployment, which is true of most of them. + fixture_prefixes: Tuple[str, ...] = () + + #: Set per class in setUpClass rather than once for the module, so the + #: ``LIKE`` patterns in a class can only ever match clusters that class + #: created. The exact-count assertion in ``test_show_clusters_like`` used + #: to rely on the lifecycle tests sorting alphabetically after it and + #: their clusters leaving the list endpoint in time; with a per-class id + #: it holds whatever else is running. + id: str = '' + dbname: str = '' + dbexisted: bool = False + clusters: List[Any] = [] + manager: Any = None + project_id: str = '' + us_regions: List[Any] = [] + + @classmethod + def _project_id(cls, mgr): + """Pick the project to deploy into, or skip. POST requires one.""" + from_env = os.environ.get('SINGLESTOREDB_TEST_PROJECT') + if from_env: + return from_env + standard = [x for x in mgr.projects if x.edition == 'STANDARD'] + if not standard: + raise unittest.SkipTest( + 'No STANDARD project in this organization; set ' + 'SINGLESTOREDB_TEST_PROJECT to the project to deploy into', + ) + return standard[0].id + + @classmethod + def setUpClass(cls): + cls.id = secrets.token_hex(4) + # Rebound per class: a list on the mixin would be one object shared by + # every subclass, so one class's teardown would pop another's clusters. + cls.clusters = [] + + sql_file = os.path.join(os.path.dirname(__file__), 'test.sql') + cls.dbname, cls.dbexisted = utils.load_sql(sql_file) + + # Pinned: the CLUSTER commands are the v2 vocabulary, so the fixture + # must not follow the management.version option out of v2 either. + mgr = s2.manage_clusters(version='v2') + cls.manager = mgr + + cls.us_regions = [ + x for x in mgr.regions + if 'US' in x.name or 'us-' in (x.region_name or '') + ] + if not cls.us_regions: + raise unittest.SkipTest('No US regions reported by the v2 API') + + cls.project_id = cls._project_id(mgr) + + for prefix in cls.fixture_prefixes: + region = random.choice(cls.us_regions) + cls.clusters.append( + mgr.create_cluster( + f'{prefix}-fusion-cluster-{cls.id}', + region=region, + size='S-00', + project=cls.project_id, + wait_on_active=True, + wait_timeout=1200, + ), + ) + + @classmethod + def tearDownClass(cls): + if not cls.dbexisted: + utils.drop_database(cls.dbname) + while cls.clusters: + cluster = cls.clusters.pop() try: - mgr.workspace_groups[wg_name].terminate(force=True) + # No wait_on_terminated: teardown only needs the DELETE to + # land, and waiting each cluster out serially costs minutes + # that assert nothing. Anything the DELETE fails to remove is + # swept by utils.cleanup_tracked. The one place termination has + # to be observed is test_create_drop_cluster, which polls the + # listing itself through _wait_cluster_gone. + cluster.terminate(force=True) except Exception: pass + def setUp(self): + self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED') + os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '1' + self.conn = s2.connect(database=type(self).dbname, local_infile=True) + self.cur = self.conn.cursor() + + def tearDown(self): + if self.enabled: + os.environ['SINGLESTOREDB_FUSION_ENABLED'] = self.enabled + else: + del os.environ['SINGLESTOREDB_FUSION_ENABLED'] + + try: + if self.cur is not None: + self.cur.close() + except Exception: + pass + + try: + if self.conn is not None: + self.conn.close() + except Exception: + pass + + +@pytest.mark.management +class TestClusterFusion(_ClusterFusionMixin, unittest.TestCase): + """ + ``SHOW CLUSTERS`` against three deployed clusters. + + Three of them so the ``LIKE``/``ORDER BY``/``LIMIT`` assertions have + something to sort. Nothing here mutates a cluster, which is what makes the + fixture shareable -- ``SUSPEND``/``RESUME`` cannot share it and deploys its + own in :class:`TestClusterFusionSuspendResume`. + """ + + fixture_prefixes = ('a', 'b', 'c') + + def test_show_clusters(self): + self.cur.execute('show clusters') + names = [x[0] for x in self.cur.fetchall()] + assert self.cur.description[0][0] == 'Name' + for prefix in ('a', 'b', 'c'): + assert f'{prefix}-fusion-cluster-{self.id}' in names, names + + def test_show_clusters_columns(self): + self.cur.execute('show clusters') + cols = [x[0] for x in self.cur.description] + assert cols == ['Name', 'ID', 'Region', 'Size', 'State'], cols + + self.cur.execute('show clusters extended') + cols = [x[0] for x in self.cur.description] + assert cols == [ + 'Name', 'ID', 'Region', 'Size', 'State', 'Provider', 'Endpoint', + 'DeploymentType', 'FirewallRanges', 'ProjectName', 'CreatedAt', + 'TerminatedAt', + ], cols + + rows = {x[0]: x for x in self.cur.fetchall()} + row = rows[f'a-fusion-cluster-{self.id}'] + # Region is the provider slug; Cluster has no region object at v2. + assert row[2], row + assert row[5], row + # ProjectName, not the ID: the column reports the name the project + # listing gives for the ID the cluster was deployed into. + project = type(self).manager.projects[type(self).project_id] + assert row[9] == project.name, row + + def test_show_clusters_like(self): + self.cur.execute(f'show clusters like "a-fusion-cluster-{self.id}"') + names = [x[0] for x in self.cur.fetchall()] + assert names == [f'a-fusion-cluster-{self.id}'], names + + self.cur.execute(f'show clusters like "%-fusion-cluster-{self.id}"') + names = [x[0] for x in self.cur.fetchall()] + assert len(names) == 3, names + + def test_show_clusters_order_by_and_limit(self): + self.cur.execute( + f'show clusters like "%-fusion-cluster-{self.id}" order by name', + ) + names = [x[0] for x in self.cur.fetchall()] + assert names == sorted(names), names + + self.cur.execute( + f'show clusters like "%-fusion-cluster-{self.id}" ' + 'order by name desc', + ) + names = [x[0] for x in self.cur.fetchall()] + assert names == sorted(names, reverse=True), names + + self.cur.execute( + f'show clusters like "%-fusion-cluster-{self.id}" ' + 'order by name limit 2', + ) + names = [x[0] for x in self.cur.fetchall()] + assert len(names) == 2, names + + +@pytest.mark.management +class TestClusterFusionReadOnly(_ClusterFusionMixin, unittest.TestCase): + """ + The handlers that read something the organization already has. + + Projects, regions and starter clusters are all pre-existing, so this class + deploys nothing -- these assertions were waiting on three clusters they + never looked at. Nothing here asserts a row count over a listing, so other + suites deploying at the same time cannot disturb them. + """ + + def test_show_projects(self): + self.cur.execute('show projects') + cols = [x[0] for x in self.cur.description] + assert cols == ['Name', 'ID', 'Edition', 'CreatedAt'], cols + ids = [x[1] for x in self.cur.fetchall()] + assert type(self).project_id in ids, ids + + def test_show_cluster_regions(self): + self.cur.execute('show cluster regions') + cols = [x[0] for x in self.cur.description] + # No ID column: v2 assigns no region IDs. This doubles as the live + # check that the region shape is what the wrappers assume. + assert cols == ['Name', 'Provider', 'RegionName'], cols + + rows = self.cur.fetchall() + assert rows + for name, provider, region_name in rows: + assert name, rows + assert provider, rows + assert region_name, rows + # The display name and the provider slug are different senses on + # this route; if they were equal the wrappers would be reading + # the wrong field. + assert region_name != name or ' ' not in name + + def test_show_cluster_regions_like(self): + self.cur.execute('show cluster regions like "US%" order by name') + names = [x[0] for x in self.cur.fetchall()] + assert names, names + assert all(x.startswith('US') for x in names), names + assert names == sorted(names), names + + def test_show_starter_clusters(self): + self.cur.execute('show starter clusters') + cols = [x[0] for x in self.cur.description] + assert cols == ['Name', 'ID', 'DatabaseName'], cols + + self.cur.execute('show starter clusters extended') + cols = [x[0] for x in self.cur.description] + assert cols == [ + 'Name', 'ID', 'DatabaseName', 'Endpoint', 'ProjectName', + ], cols + + def test_drop_starter_cluster_if_exists(self): + """IF EXISTS must swallow the miss; the bare form must not.""" + with self.assertRaises(KeyError): + self.cur.execute('drop starter cluster "no-such-starter-xyz"') + self.cur.execute('drop starter cluster if exists "no-such-starter-xyz"') + + +@pytest.mark.management +class TestClusterFusionCreateDrop(_ClusterFusionMixin, unittest.TestCase): + """ + ``CREATE CLUSTER`` and ``DROP CLUSTER`` end to end. + + Deploys nothing up front: the test creates, drops and recreates a cluster + of its own, so the three shared fixtures it used to inherit were pure cost. + Alone in its class because it is the longest test in the repo -- most of + twenty minutes, nearly all of it provisioning -- and anything sharing the + class would queue behind it. + """ + + def _wait_cluster_gone(self, name, timeout=180, interval=5): + """ + Poll until the LIST endpoint agrees the cluster is gone. + + The mirror of ``_wait_workspace_group_gone``: ``WAIT ON TERMINATED`` + polls ``GET /v2/clusters/{id}``, and ``GET /v2/clusters`` can lag + behind it, so a create-drop-create sequence sees a stale record. + """ + mgr = type(self).manager + deadline = time.time() + timeout + while True: + found = [x for x in mgr.clusters if x.name == name] + if not found or all(x.terminated_at is not None for x in found): + return + if time.time() >= deadline: + self.fail( + f'cluster {name!r} still active in the list endpoint ' + f'after {timeout}s: {found!r}', + ) + time.sleep(interval) + + def test_create_drop_cluster(self): + mgr = type(self).manager + name = f'd-fusion-cluster-{self.id}' + region = type(self).us_regions[0] + + try: + self.cur.execute( + f'create cluster "{name}" in region "{region.region_name}" ' + f'using provider "{region.provider}" ' + f'in project id "{type(self).project_id}" ' + 'with size "S-00" wait on active', + ) + + # Unlike CREATE WORKSPACE GROUP, this returns a row -- the + # generated password appears in the create response and nowhere + # else, so a caller who cannot see it has no admin access. + row = self.cur.fetchall() + cols = [x[0] for x in self.cur.description] + assert cols == ['Name', 'ID', 'Endpoint', 'AdminPassword'], cols + assert len(row) == 1, row + assert row[0][0] == name, row + assert row[0][1], row + + live = [ + x for x in mgr.clusters + if x.name == name and x.terminated_at is None + ] + assert len(live) == 1, live + cluster_id = live[0].id + + # IF NOT EXISTS on a live cluster is a no-op + self.cur.execute( + f'create cluster if not exists "{name}" ' + f'in region "{region.region_name}" ' + f'in project id "{type(self).project_id}"', + ) + live = [ + x for x in mgr.clusters + if x.name == name and x.terminated_at is None + ] + assert len(live) == 1, live + + # Drop by name + self.cur.execute(f'drop cluster "{name}" wait on terminated') + self._wait_cluster_gone(name) + + # Create again, drop by ID + self.cur.execute( + f'create cluster "{name}" in region "{region.region_name}" ' + f'in project id "{type(self).project_id}" wait on active', + ) + live = [ + x for x in mgr.clusters + if x.name == name and x.terminated_at is None + ] + assert len(live) == 1, live + cluster_id = live[0].id + + self.cur.execute(f'drop cluster id "{cluster_id}" wait on terminated') + self._wait_cluster_gone(name) + + # Drop non-existent by ID + with self.assertRaises(KeyError): + self.cur.execute(f'drop cluster id "{cluster_id}"') + + # ... and with IF EXISTS + self.cur.execute(f'drop cluster if exists id "{cluster_id}"') + + # Drop non-existent by name, both ways + with self.assertRaises(KeyError): + self.cur.execute('drop cluster "no-such-cluster-xyz"') + self.cur.execute('drop cluster if exists "no-such-cluster-xyz"') + + finally: + for cluster in mgr.clusters: + if cluster.name == name and cluster.terminated_at is None: + try: + cluster.terminate() + except Exception: + pass + + +@pytest.mark.management +class TestClusterFusionSuspendResume(_ClusterFusionMixin, unittest.TestCase): + """ + ``SUSPEND CLUSTER`` and ``RESUME CLUSTER``. + + Deploys one cluster rather than sharing :class:`TestClusterFusion`'s three, + for two reasons: it needs exactly one, and it is the only test here that + changes a fixture's state, so sharing would leave the ``SHOW`` assertions + reading a cluster mid-suspend. + """ + + fixture_prefixes = ('a',) + + def test_suspend_resume_cluster(self): + name = f'a-fusion-cluster-{self.id}' + mgr = type(self).manager + + self.cur.execute(f'suspend cluster "{name}" wait on suspended') + state = [x for x in mgr.clusters if x.name == name][0].state + assert state.upper() == 'SUSPENDED', state + + self.cur.execute(f'resume cluster "{name}" wait on resumed') + state = [x for x in mgr.clusters if x.name == name][0].state + assert state.upper() == 'ACTIVE', state + @pytest.mark.management +class TestClusterFusionProject(_ClusterFusionMixin, unittest.TestCase): + """ + How ``IN PROJECT`` resolves, and which spellings must not parse. + + Deploys nothing: the one test here that creates a cluster deliberately does + not wait it out, and the rest assert a rejection. + """ + + def test_create_cluster_without_project(self): + """ + Omitting IN PROJECT is only valid in a single-project organization. + + ``POST /v2/clusters`` requires ``projectID``, so the handler falls + through to ``_resolve_project_id()``, which picks the only project or + raises naming the candidates. Either outcome is correct; silently + choosing one of several would not be. + """ + mgr = type(self).manager + name = f'e-fusion-cluster-{self.id}' + region = type(self).us_regions[0] + + if len(mgr.projects) == 1: + raise unittest.SkipTest( + 'single-project organization; the ambiguous path is what ' + 'this test is for', + ) + + with self.assertRaises(Exception): + self.cur.execute( + f'create cluster "{name}" in region "{region.region_name}"', + ) + + # Nothing should have been created + live = [ + x for x in mgr.clusters + if x.name == name and x.terminated_at is None + ] + assert not live, live + + def test_create_cluster_named_project(self): + """ + ``IN PROJECT ""`` resolves the name to the project's ID. + + Deliberately no ``WAIT ON ACTIVE``. The ``projectID`` is settled by the + time ``POST /v2/clusters`` answers -- the create response carries the + cluster ID, and ``GET /v2/clusters/{id}`` reports the project straight + away -- so provisioning the cluster the rest of the way would add + minutes of waiting and assert nothing this does not already prove. + """ + mgr = type(self).manager + project = [ + x for x in mgr.projects if x.id == type(self).project_id + ][0] + name = f'f-fusion-cluster-{self.id}' + region = type(self).us_regions[0] + + cluster_id = None + try: + self.cur.execute( + f'create cluster "{name}" in region "{region.region_name}" ' + f'in project "{project.name}" with size "S-00"', + ) + row = self.cur.fetchall() + assert len(row) == 1, row + assert row[0][0] == name, row + cluster_id = row[0][1] + assert cluster_id, row + + # Read back through GET /v2/clusters/{id} rather than the listing: + # the create is not waited out, and the LIST endpoint can lag + # behind a cluster it has only just been told about. + assert mgr.get_cluster(cluster_id).project.id == project.id + + finally: + # force=True: the cluster is still PENDING, having never been + # waited out, and a termination request is refused otherwise. + if cluster_id is not None: + try: + mgr.get_cluster(cluster_id).terminate(force=True) + except Exception: + pass + else: + for cluster in mgr.clusters: + if cluster.name == name and cluster.terminated_at is None: + try: + cluster.terminate(force=True) + except Exception: + pass + + def test_region_id_does_not_parse(self): + """v2 has no region IDs, so the v1 spelling must be rejected.""" + with self.assertRaises(Exception): + self.cur.execute( + f'create cluster "g-fusion-cluster-{self.id}" ' + 'in region id "abc"', + ) + + def test_unknown_project_raises(self): + with self.assertRaises(KeyError): + self.cur.execute( + f'create cluster "h-fusion-cluster-{self.id}" ' + 'in region "us-east-1" ' + 'in project "no such project xyz"', + ) + + +@pytest.mark.management +@pytest.mark.xdist_group(utils.SHARED_CLUSTER_JOBS_GROUP) class TestJobsFusion(unittest.TestCase): - id: str = secrets.token_hex(8) notebook_name: str = 'Scheduling Test.ipynb' dbname: str = '' dbexisted: bool = False manager: None - workspace_group: None - workspace: None + cluster: None job_ids = [] @classmethod def setUpClass(cls): sql_file = os.path.join(os.path.dirname(__file__), 'test.sql') cls.dbname, cls.dbexisted = utils.load_sql(sql_file) - cls.manager = s2.manage_workspaces() - us_regions = [x for x in cls.manager.regions if x.name.startswith('US')] - cls.workspace_group = cls.manager.create_workspace_group( - f'Jobs Fusion Testing {cls.id}', - region=random.choice(us_regions), - firewall_ranges=[], - ) - cls.workspace = cls.workspace_group.create_workspace( - f'jobs-test-{cls.id}', - wait_on_active=True, - ) + + # Switched to v2 along with the JOB handlers. A job runs against a + # deployment, and at v2 a deployment is a cluster -- one create call + # rather than a group plus a workspace. This is the only live exercise + # of the Cluster/VirtualCluster targetType vocabulary. + cls.manager = s2.manage_clusters(version='v2') + + # A shared cluster: a job needs a live deployment to target, and every + # listing here is filtered by job id, so nothing this class asserts can + # see another class's jobs. + cls.cluster = utils.shared_clusters(1)[0] + os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = cls.dbname - os.environ['SINGLESTOREDB_WORKSPACE'] = cls.workspace.id + # SINGLESTOREDB_WORKSPACE is the only deployment variable the notebook + # environment publishes -- there is no SINGLESTOREDB_CLUSTER -- and at + # v2 its value is a cluster ID. + os.environ['SINGLESTOREDB_WORKSPACE'] = cls.cluster.id @classmethod def tearDownClass(cls): @@ -523,15 +1952,14 @@ def tearDownClass(cls): cls.manager.organizations.current.jobs.delete(job_id) except Exception: pass - if cls.workspace_group is not None: - cls.workspace_group.terminate(force=True) + # The cluster is the pool's; see TestStageFusion.tearDownClass. cls.manager = None - cls.workspace_group = None - cls.workspace = None - if os.environ.get('SINGLESTOREDB_WORKSPACE', None) is not None: - del os.environ['SINGLESTOREDB_WORKSPACE'] - if os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE', None) is not None: - del os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] + cls.cluster = None + for envvar in ( + 'SINGLESTOREDB_WORKSPACE', + 'SINGLESTOREDB_DEFAULT_DATABASE', + ): + os.environ.pop(envvar, None) def setUp(self): self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED') @@ -671,8 +2099,10 @@ def test_show_jobs_and_executions(self): assert job[1] == 'show-job' assert job[5] == self.notebook_name assert job[6] == self.dbname - assert job[7] == self.workspace.id - assert job[8] == 'Workspace' + assert job[7] == self.cluster.id + # targetType is 'Cluster' at v2 where v1 reported 'Workspace'; + # this is the assertion that proves the manager really moved. + assert job[8] == 'Cluster' # show jobs with name like "show-job" extended self.cur.execute(f'show jobs {job_id} like "show-job" extended') @@ -692,8 +2122,10 @@ def test_show_jobs_and_executions(self): assert job[1] == 'show-job' assert job[5] == self.notebook_name assert job[6] == self.dbname - assert job[7] == self.workspace.id - assert job[8] == 'Workspace' + assert job[7] == self.cluster.id + # targetType is 'Cluster' at v2 where v1 reported 'Workspace'; + # this is the assertion that proves the manager really moved. + assert job[8] == 'Cluster' assert not job[11] assert job[13] == 5 assert job[14] == 'Recurring' @@ -741,51 +2173,56 @@ def test_show_jobs_and_executions(self): @pytest.mark.management +@pytest.mark.xdist_group(utils.SHARED_CLUSTER_STAGE_GROUP) class TestStageFusion(unittest.TestCase): - id: str = secrets.token_hex(8) dbname: str = 'information_schema' manager: None - workspace_group: None - workspace_group_2: None + cluster: None + cluster_2: None @classmethod def setUpClass(cls): - cls.manager = s2.manage_workspaces() - us_regions = [x for x in cls.manager.regions if x.name.startswith('US')] - cls.workspace_group = cls.manager.create_workspace_group( - f'Stage Fusion Testing 1 {cls.id}', - region=random.choice(us_regions), - firewall_ranges=[], - ) - cls.workspace_group_2 = cls.manager.create_workspace_group( - f'Stage Fusion Testing 2 {cls.id}', - region=random.choice(us_regions), - firewall_ranges=[], - ) - # Wait for both workspace groups to start - time.sleep(5) + # Switched to v2: get_deployment() resolves against clusters, so the + # fixtures must be clusters. The v1 stage path is still covered by + # test_management_v1.py -- switched rather than duplicated, because + # duplicating doubles a suite that already runs for tens of minutes. + cls.manager = s2.manage_clusters(version='v2') + + # Two clusters from the shared pool rather than two of this class's + # own. Nothing here mutates a cluster, and the second one exists only + # so a bare IN can name a deployment other than the default. Deploying + # them was 891s of the run; see utils.shared_clusters. + cls.cluster, cls.cluster_2 = utils.shared_clusters(2) + + # The stage paths below are fixed rather than namespaced, and the + # listings are asserted by exact contents, so both stages have to start + # empty: a pool cluster carries whatever the class before it left + # there. tearDown clears them again after every test. + for cluster in (cls.cluster, cls.cluster_2): + utils.clear_stage(cluster) os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = 'information_schema' - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] = cls.workspace_group.id + # SINGLESTOREDB_WORKSPACE_GROUP would raise at v2: its value is a + # group ID, which v2 reports only as Cluster.group and cannot look + # up, so get_deployment() refuses to guess which cluster was meant + # rather than target the wrong one. + os.environ['SINGLESTOREDB_WORKSPACE'] = cls.cluster.id @classmethod def tearDownClass(cls): - if cls.workspace_group is not None: - cls.workspace_group.terminate(force=True) - if cls.workspace_group_2 is not None: - cls.workspace_group_2.terminate(force=True) + # The clusters are the pool's, not this class's: they stay live for the + # classes that follow and are terminated once, at the end of the + # session, by utils.cleanup_tracked. cls.manager = None - cls.workspace_group = None - cls.workspace_group_2 = None - cls.workspace = None - cls.workspace_2 = None - if os.environ.get('SINGLESTOREDB_WORKSPACE', None) is not None: - del os.environ['SINGLESTOREDB_WORKSPACE'] - if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP', None) is not None: - del os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] - if os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE', None) is not None: - del os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] + cls.cluster = None + cls.cluster_2 = None + for envvar in ( + 'SINGLESTOREDB_WORKSPACE', + 'SINGLESTOREDB_WORKSPACE_GROUP', + 'SINGLESTOREDB_DEFAULT_DATABASE', + ): + os.environ.pop(envvar, None) def setUp(self): self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED') @@ -816,10 +2253,10 @@ def tearDown(self): pass def _clear_stage(self): - if self.workspace_group is not None: + if self.cluster is not None: self.cur.execute(f''' show stage files - in group id '{self.workspace_group.id}' recursive + in id '{self.cluster.id}' recursive ''') files = list(self.cur) folders = [] @@ -829,18 +2266,18 @@ def _clear_stage(self): continue self.cur.execute(f''' drop stage file '{file[0]}' - in group id '{self.workspace_group.id}' + in id '{self.cluster.id}' ''') for folder in folders: self.cur.execute(f''' drop stage folder '{folder[0]}' - in group id '{self.workspace_group.id}' + in id '{self.cluster.id}' ''') - if self.workspace_group_2 is not None: + if self.cluster_2 is not None: self.cur.execute(f''' show stage files - in group id '{self.workspace_group_2.id}' recursive + in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) folders = [] @@ -850,12 +2287,12 @@ def _clear_stage(self): continue self.cur.execute(f''' drop stage file '{file[0]}' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' ''') for folder in folders: self.cur.execute(f''' drop stage folder '{folder[0]}' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' ''') def test_show_stage(self): @@ -915,57 +2352,27 @@ def test_show_stage(self): 'subdir2/', ] - # List files in specific workspace group - self.cur.execute(f''' - show stage files in group id '{self.workspace_group.id}' - ''') - files = list(self.cur) - assert len(files) == 3 - assert list(sorted(x[0] for x in files)) == [ - 'new_test_1.sql', - 'subdir1/', - 'subdir2/', - ] - - self.cur.execute(f''' - show stage files in id '{self.workspace_group.id}' - ''') - files = list(self.cur) - assert len(files) == 3 - assert list(sorted(x[0] for x in files)) == [ - 'new_test_1.sql', - 'subdir1/', - 'subdir2/', - ] - - self.cur.execute(f''' - show stage files in group '{self.workspace_group.name}' - ''') - files = list(self.cur) - assert len(files) == 3 - assert list(sorted(x[0] for x in files)) == [ - 'new_test_1.sql', - 'subdir1/', - 'subdir2/', - ] - - self.cur.execute(f''' - show stage files in '{self.workspace_group.name}' - ''') - files = list(self.cur) - assert len(files) == 3 - assert list(sorted(x[0] for x in files)) == [ + # List files in a specific deployment. A bare IN is the only spelling + # that names one: IN GROUP names a v1 workspace group instead, which is + # a different resource, so it is not tested against a cluster here -- + # TestWorkspaceFusion covers it against a real group. + expected = [ 'new_test_1.sql', 'subdir1/', 'subdir2/', ] + for clause in [ + f"in id '{self.cluster.id}'", + f"in '{self.cluster.name}'", + ]: + self.cur.execute(f'show stage files {clause}') + files = list(self.cur) + assert len(files) == 3, (clause, files) + assert list(sorted(x[0] for x in files)) == expected, clause - # Check other workspace group - self.cur.execute(f''' - show stage files in group '{self.workspace_group_2.name}' - ''') - files = list(self.cur) - assert len(files) == 0 + # Check the other cluster + self.cur.execute(f"show stage files in '{self.cluster_2.name}'") + assert len(list(self.cur)) == 0 # Limit results self.cur.execute(''' @@ -1059,13 +2466,13 @@ def test_download_stage(self): # Copy file to stage 2 self.cur.execute(f''' upload file to stage 'dl_test2.sql' - in group '{self.workspace_group_2.name}' + in '{self.cluster_2.name}' from '{test2_sql}' ''') # Make sure only one file in stage 2 self.cur.execute(f''' - show stage files in group '{self.workspace_group_2.name}' + show stage files in '{self.cluster_2.name}' ''') files = list(self.cur) assert len(files) == 1 @@ -1083,7 +2490,7 @@ def test_download_stage(self): with tempfile.TemporaryDirectory() as tmpdir: self.cur.execute(f''' download stage file 'dl_test2.sql' - in group '{self.workspace_group_2.name}' + in '{self.cluster_2.name}' to '{tmpdir}/dl_test2.sql' ''') with open(os.path.join(tmpdir, 'dl_test2.sql'), 'r') as dl_file: @@ -1114,7 +2521,7 @@ def test_stage_multi_wg_operations(self): # Copy file to stage 2 self.cur.execute(f''' upload file to stage 'new_test2.sql' - in group '{self.workspace_group_2.name}' + in '{self.cluster_2.name}' from '{test2_sql}' ''') @@ -1128,7 +2535,7 @@ def test_stage_multi_wg_operations(self): # Make sure only one file in stage 2 self.cur.execute(f''' - show stage files in group '{self.workspace_group_2.name}' recursive + show stage files in '{self.cluster_2.name}' recursive ''') files = list(self.cur) assert len(files) == 1 @@ -1136,7 +2543,7 @@ def test_stage_multi_wg_operations(self): # Make sure only one file in stage 2 (using IN) self.cur.execute(f''' - show stage files in '{self.workspace_group_2.name}' recursive + show stage files in '{self.cluster_2.name}' recursive ''') files = list(self.cur) assert len(files) == 1 @@ -1144,13 +2551,13 @@ def test_stage_multi_wg_operations(self): # Make subdir self.cur.execute(f''' - create stage folder 'data' in group '{self.workspace_group_2.name}' + create stage folder 'data' in '{self.cluster_2.name}' ''') # Upload file using workspace ID self.cur.execute(f''' upload file to stage 'data/new_test2_sub.sql' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' from '{test2_sql}' ''') @@ -1164,7 +2571,7 @@ def test_stage_multi_wg_operations(self): # Make sure two files in stage 2 self.cur.execute(f''' - show stage files in group id '{self.workspace_group_2.id}' recursive + show stage files in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) assert len(files) == 3 @@ -1175,19 +2582,19 @@ def test_stage_multi_wg_operations(self): with self.assertRaises(OSError): self.cur.execute(f''' upload file to stage 'data/new_test2_sub.sql' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' from '{test2_sql}' ''') self.cur.execute(f''' upload file to stage 'data/new_test2_sub.sql' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' from '{test2_sql}' overwrite ''') # Make sure two files in stage 2 self.cur.execute(f''' - show stage files in group id '{self.workspace_group_2.id}' recursive + show stage files in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) assert len(files) == 3 @@ -1197,7 +2604,7 @@ def test_stage_multi_wg_operations(self): # Test LIKE clause self.cur.execute(f''' show stage files - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' like '%_sub%' recursive ''') files = list(self.cur) @@ -1218,7 +2625,7 @@ def test_stage_multi_wg_operations(self): # Make sure two files in stage 2 self.cur.execute(f''' - show stage files in group id '{self.workspace_group_2.id}' recursive + show stage files in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) assert len(files) == 3 @@ -1229,17 +2636,17 @@ def test_stage_multi_wg_operations(self): with self.assertRaises(OSError): self.cur.execute(f''' drop stage folder 'data' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' ''') self.cur.execute(f''' drop stage file 'data/new_test2_sub.sql' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' ''') # Make sure one file and one directory in stage 2 self.cur.execute(f''' - show stage files in group id '{self.workspace_group_2.id}' recursive + show stage files in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) assert len(files) == 2 @@ -1248,12 +2655,12 @@ def test_stage_multi_wg_operations(self): # Drop stage folder from stage 2 self.cur.execute(f''' drop stage folder 'data' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' ''') # Make sure one file in stage 2 self.cur.execute(f''' - show stage files in group id '{self.workspace_group_2.id}' recursive + show stage files in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) assert len(files) == 1 @@ -1262,12 +2669,12 @@ def test_stage_multi_wg_operations(self): # Drop last file self.cur.execute(f''' drop stage file 'new_test2.sql' - in group id '{self.workspace_group_2.id}' + in id '{self.cluster_2.id}' ''') # Make sure no files in stage 2 self.cur.execute(f''' - show stage files in group id '{self.workspace_group_2.id}' recursive + show stage files in id '{self.cluster_2.id}' recursive ''') files = list(self.cur) assert len(files) == 0 @@ -1279,36 +2686,21 @@ class TestFilesFusion(unittest.TestCase): id: str = secrets.token_hex(8) dbname: str = 'information_schema' manager: None - workspace_group: None @classmethod def setUpClass(cls): - cls.manager = s2.manage_workspaces() - us_regions = [x for x in cls.manager.regions if x.name.startswith('US')] - cls.workspace_group = cls.manager.create_workspace_group( - f'Files Fusion Testing {cls.id}', - region=random.choice(us_regions), - firewall_ranges=[], - ) - # Wait for both workspace groups to start - time.sleep(5) - + # Switched to v2 along with get_files_manager(). No deployment + # fixture: the personal, shared and models spaces are org-scoped, and + # none of the tests below ever referenced the workspace group this + # method used to create -- it was a billable resource created for + # nothing. Dropped rather than converted to a cluster. + cls.manager = s2.manage_clusters(version='v2') os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = 'information_schema' - os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] = cls.workspace_group.id @classmethod def tearDownClass(cls): - if cls.workspace_group is not None: - cls.workspace_group.terminate(force=True) cls.manager = None - cls.workspace_group = None - cls.workspace = None - if os.environ.get('SINGLESTOREDB_WORKSPACE', None) is not None: - del os.environ['SINGLESTOREDB_WORKSPACE'] - if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP', None) is not None: - del os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] - if os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE', None) is not None: - del os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] + os.environ.pop('SINGLESTOREDB_DEFAULT_DATABASE', None) def setUp(self): self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED') diff --git a/singlestoredb/tests/test_management_timing.py b/singlestoredb/tests/test_management_timing.py new file mode 100644 index 000000000..5aac42926 --- /dev/null +++ b/singlestoredb/tests/test_management_timing.py @@ -0,0 +1,458 @@ +#!/usr/bin/env python +# type: ignore +""" +Time accounting for the management API. + +These are unit tests: no token, no deployment and no HTTP. The point of the +module under test is to separate time spent in requests from time spent +sleeping in a ``wait_on_*`` loop, so that is what is asserted -- against a +mocked session, with :func:`time.sleep` patched out. +""" +import contextlib +import io +import unittest +from unittest.mock import MagicMock +from unittest.mock import patch + +import singlestoredb as s2 +from singlestoredb.management import timing + + +FAKE_TOKEN = 'test-token-12345' +FAKE_BASE_URL = 'https://api.example.com' +FAKE_ID = '44444444-4444-4444-8444-444444444444' + + +def _response(status_code=200, body=b'{}', request_body=None, retries=None): + """Return a stand-in for a requests.Response.""" + out = MagicMock() + out.status_code = status_code + out.headers = {'Content-Length': str(len(body))} + out.content = body + out.request.body = request_body + if retries is None: + del out.raw.retries + else: + out.raw.retries.history = retries + return out + + +@contextlib.contextmanager +def _tracing(enabled): + """ + Force the ``management.trace`` option for the duration of the block. + + Both directions are needed: the option is read from + ``SINGLESTOREDB_MANAGEMENT_TRACE``, which a traced test run sets, so the + tests that assert nothing is recorded have to turn it off explicitly rather + than assume it. Any ambient trace is detached for the same reason -- the + conftest opens one around every test in a traced run. + """ + token = timing._active.set(()) + s2.config.set_option('management.trace', enabled) + try: + yield + finally: + s2.config.reset_option('management.trace') + timing._active.reset(token) + + +def _manager(response=None): + """Return a Manager whose session answers with ``response``.""" + from singlestoredb.management.manager import Manager + with patch( + 'singlestoredb.management.manager.get_token', + return_value=FAKE_TOKEN, + ): + mgr = Manager(access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL) + mgr._sess = MagicMock() + mgr._sess.get.return_value = response or _response() + mgr._sess.post.return_value = response or _response() + return mgr + + +class TestRouteOf(unittest.TestCase): + """Aggregation keys. One row per route, not one per resource.""" + + def test_ids_are_collapsed(self): + self.assertEqual( + timing.route_of('get', f'clusters/{FAKE_ID}'), + 'GET clusters/{id}', + ) + self.assertEqual( + timing.route_of('get', 'jobs/12345/executions'), + 'GET jobs/{id}/executions', + ) + + def test_route_segments_are_kept(self): + self.assertEqual( + timing.route_of('post', 'clusters'), 'POST clusters', + ) + self.assertEqual( + timing.route_of('get', 'regions/sharedtier'), + 'GET regions/sharedtier', + ) + + def test_query_strings_and_slashes_are_dropped(self): + self.assertEqual( + timing.route_of('get', '/clusters/?force=true'), 'GET clusters', + ) + + +class TestRecording(unittest.TestCase): + """What is and is not collected.""" + + def test_nothing_is_recorded_without_a_trace(self): + with _tracing(False): + self.assertFalse(timing.recording()) + # Nothing to assert but that this does not raise or cost + # anything: the event is dropped before an Event is even built. + timing.record_request('get', 'clusters', 1.0, 0.0) + + def test_a_trace_collects_requests(self): + mgr = _manager(_response(body=b'{"a": 1}', request_body=b'{"b": 2}')) + with timing.trace() as trace: + mgr._get('clusters') + mgr._get(f'clusters/{FAKE_ID}') + + self.assertEqual(len(trace.events), 2) + self.assertEqual( + [x.label for x in trace.events], + ['GET clusters', 'GET clusters/{id}'], + ) + for event in trace.events: + self.assertEqual(event.kind, timing.REQUEST) + self.assertEqual(event.status, 200) + self.assertEqual(event.response_bytes, 8) + self.assertEqual(event.request_bytes, 8) + self.assertEqual(event.retries, 0) + self.assertIsNone(event.error) + self.assertGreaterEqual(event.duration, 0.0) + + def test_a_trace_stops_collecting_at_the_end_of_the_block(self): + mgr = _manager() + with timing.trace() as trace: + mgr._get('clusters') + mgr._get('clusters') + self.assertEqual(len(trace.events), 1) + + def test_retries_are_reported(self): + """ + A retried request has to be distinguishable from a slow one. + + Retries happen under ``requests``, so without this the backoff shows up + as one long response and nothing says why. + """ + mgr = _manager(_response(retries=[object(), object()])) + with timing.trace() as trace: + mgr._get('clusters') + self.assertEqual(trace.events[0].retries, 2) + self.assertEqual(trace.stats()[0].retries, 2) + + def test_a_failed_request_is_recorded_with_its_error(self): + import requests + mgr = _manager() + mgr._sess.get.side_effect = requests.exceptions.ConnectionError('boom') + with timing.trace() as trace: + with self.assertRaises(s2.ManagementError): + mgr._get('clusters') + self.assertEqual(len(trace.events), 1) + self.assertEqual(trace.events[0].error, 'ConnectionError') + self.assertIsNone(trace.events[0].status) + self.assertEqual(trace.stats()[0].errors, 1) + + def test_an_http_error_is_still_a_recorded_request(self): + mgr = _manager(_response(status_code=404, body=b'nope')) + with timing.trace() as trace: + with self.assertRaises(s2.ManagementError): + mgr._get('clusters') + self.assertEqual(trace.events[0].status, 404) + + def test_nested_traces_both_collect(self): + mgr = _manager() + with timing.trace() as outer: + mgr._get('clusters') + with timing.trace() as inner: + mgr._get('regions') + mgr._get('clusters') + + self.assertEqual([x.label for x in inner.events], ['GET regions']) + self.assertEqual( + [x.label for x in outer.events], + ['GET clusters', 'GET regions', 'GET clusters'], + ) + + +class TestWaiting(unittest.TestCase): + """Polling sleeps, which is where the wall clock usually goes.""" + + def test_sleep_is_recorded_as_a_wait(self): + with patch('singlestoredb.management.timing.time.sleep') as slept: + with timing.trace() as trace: + timing.sleep(20, 'cluster state -> active') + slept.assert_called_once_with(20) + self.assertEqual(len(trace.events), 1) + self.assertEqual(trace.events[0].kind, timing.WAIT) + self.assertEqual(trace.events[0].label, 'cluster state -> active') + + def test_sleep_still_sleeps_when_nothing_is_recording(self): + with patch('singlestoredb.management.timing.time.sleep') as slept: + timing.sleep(20, 'cluster state -> active') + slept.assert_called_once_with(20) + + def test_timed_labels_blocking_work_that_is_neither(self): + with timing.trace() as trace: + with timing.timed('cluster endpoint connect'): + pass + self.assertEqual(trace.events[0].kind, timing.WAIT) + self.assertEqual(trace.events[0].label, 'cluster endpoint connect') + + def test_timed_records_even_when_the_block_raises(self): + with timing.trace() as trace: + with self.assertRaises(ValueError): + with timing.timed('cluster endpoint connect'): + raise ValueError('boom') + self.assertEqual(len(trace.events), 1) + + def test_wait_on_state_records_one_wait_per_poll(self): + """ + The ``wait_on_*`` loops are the reason this module exists. + + Three polls of a cluster that is still PENDING have to show up as three + waits and two requests, not as one opaque 40 seconds. + """ + mgr = _manager() + mgr.obj_type = 'cluster' + pending, active = MagicMock(), MagicMock() + pending.state = 'PENDING' + pending.id = FAKE_ID + active.state = 'ACTIVE' + active.id = FAKE_ID + mgr.get_cluster = MagicMock(side_effect=[pending, active]) + + with patch('singlestoredb.management.timing.time.sleep'): + with timing.trace() as trace: + out = mgr._wait_on_state(pending, 'ACTIVE', interval=20) + + self.assertIs(out, active) + waits = [x for x in trace.events if x.kind == timing.WAIT] + self.assertEqual(len(waits), 2) + self.assertEqual({x.label for x in waits}, {'cluster state -> active'}) + + +class TestPollCost(unittest.TestCase): + """ + ``wait_timeout`` has to be a duration, not a poll count. + + The loops used to charge every iteration a flat ``interval`` and never + counted the refetch between sleeps. Since the session gained retries and a + 180 second read timeout a single poll can cost minutes, so a caller asking + to wait 600 seconds could wait for an hour without a timeout being raised. + """ + + @contextlib.contextmanager + def _clock(self, per_call=0.0): + """ + Run with a fake monotonic clock and no real sleeping. + + ``per_call`` is the number of seconds each *refetch* is made to appear + to take, which is what the old accounting ignored. + """ + reading = [0.0] + + def advance(): + reading[0] += per_call + return reading[0] + + with patch('singlestoredb.management.timing.time.sleep'): + with patch( + 'singlestoredb.management.timing.now', + side_effect=lambda: reading[0], + ): + yield advance + + def test_poll_cost_floors_at_the_interval(self): + with self._clock(): + self.assertEqual(timing.poll_cost(timing.now(), 10), 10) + + def test_poll_cost_charges_measured_time_when_it_exceeds_the_interval(self): + with self._clock(per_call=100.0) as advance: + started_at = timing.now() + advance() + self.assertEqual(timing.poll_cost(started_at, 10), 100.0) + + def test_a_slow_refetch_counts_against_the_timeout(self): + """ + Six polls of a refetch that costs 100s, not sixty of a nominal 10s. + + The whole point of the fix: ``timeout=600`` means ten minutes of wall + clock, so a poll that really takes 100 seconds exhausts it in six + iterations rather than sixty. + """ + mgr = _manager() + mgr.obj_type = 'cluster' + pending = MagicMock() + pending.state = 'PENDING' + pending.id = FAKE_ID + + with self._clock(per_call=100.0) as advance: + def refetch(id): + advance() + return pending + mgr.get_cluster = MagicMock(side_effect=refetch) + + with self.assertRaises(s2.ManagementError): + mgr._wait_on_state( + pending, 'ACTIVE', interval=10, timeout=600, + ) + + self.assertEqual(mgr.get_cluster.call_count, 6) + + def test_the_interval_floor_still_bounds_a_patched_out_sleep(self): + """ + With ``time.sleep`` patched out the measured time is ~0, so without the + floor in :func:`timing.poll_cost` nothing would ever charge the timeout + and the loop would spin forever. Every offline test that polls relies + on this. + """ + mgr = _manager() + mgr.obj_type = 'cluster' + pending = MagicMock() + pending.state = 'PENDING' + pending.id = FAKE_ID + mgr.get_cluster = MagicMock(return_value=pending) + + with self._clock(): + with self.assertRaises(s2.ManagementError): + mgr._wait_on_state( + pending, 'ACTIVE', interval=20, timeout=60, + ) + + self.assertEqual(mgr.get_cluster.call_count, 3) + + +class TestReporting(unittest.TestCase): + """Totals, aggregates and the summary text.""" + + def _trace(self): + """Return a stopped trace holding known durations.""" + trace = timing.Trace().start() + trace.add(timing.Event(timing.REQUEST, 'GET clusters', 1.0, 0.0, status=200)) + trace.add( + timing.Event(timing.REQUEST, 'GET clusters/{id}', 2.0, 1.0, status=200), + ) + trace.add( + timing.Event(timing.REQUEST, 'GET clusters/{id}', 4.0, 3.0, status=200), + ) + trace.add(timing.Event(timing.WAIT, 'cluster state -> active', 40.0, 7.0)) + return trace.stop() + + def test_totals_split_requests_from_waiting(self): + trace = self._trace() + self.assertEqual(trace.total(timing.REQUEST), 7.0) + self.assertEqual(trace.total(timing.WAIT), 40.0) + self.assertEqual(trace.total(), 47.0) + + def test_stats_aggregate_by_label_slowest_first(self): + stats = self._trace().stats(timing.REQUEST) + self.assertEqual([x.label for x in stats], ['GET clusters/{id}', 'GET clusters']) + first = stats[0] + self.assertEqual(first.calls, 2) + self.assertEqual(first.total, 6.0) + self.assertEqual(first.mean, 3.0) + self.assertEqual(first.min, 2.0) + self.assertEqual(first.max, 4.0) + + def test_unaccounted_time_is_never_negative(self): + # The events claim 47s; a trace that was open for less than that (these + # durations are fabricated) must report 0 rather than a negative. + self.assertEqual(self._trace().unaccounted, 0.0) + + def test_summary_names_the_split_and_the_routes(self): + out = self._trace().summary() + self.assertIn('requests', out) + self.assertIn('waiting', out) + self.assertIn('GET clusters/{id}', out) + self.assertIn('cluster state -> active', out) + + def test_of_holds_the_given_events_and_reports_the_given_elapsed(self): + events = self._trace().events + out = timing.Trace.of(events, 100.0) + self.assertEqual(len(out.events), 4) + self.assertEqual(out.elapsed, 100.0) + self.assertEqual(out.total(timing.WAIT), 40.0) + # Given out of order, reported in completion order. + shuffled = timing.Trace.of(list(reversed(events)), 100.0) + self.assertEqual( + [x.started_at for x in shuffled.events], [0.0, 1.0, 3.0, 7.0], + ) + + def test_of_never_reports_a_negative_elapsed(self): + # The conftest subtracts nested traces' elapsed from their parent's, + # and rounding or an overlapping trace could take that below zero. + self.assertEqual(timing.Trace.of([], -5.0).elapsed, 0.0) + + def test_a_nested_trace_shares_event_objects_with_its_parent(self): + """ + The conftest separates class-fixture time from test time by identity. + + A class-scoped trace spans its tests as well as its fixtures, so the + fixture share is the parent's events minus the children's. That is only + exact because :func:`timing._emit` hands the *same* Event to every + active trace rather than a copy per trace. + """ + mgr = _manager() + with timing.trace() as outer: + mgr._get('regions') + with timing.trace() as inner: + mgr._get('clusters') + + fixture_only = [ + x for x in outer.events if id(x) not in { + id(y) for y in inner.events + } + ] + self.assertEqual([x.label for x in fixture_only], ['GET regions']) + + def test_combine_folds_traces_and_sums_their_elapsed(self): + one, two = self._trace(), self._trace() + combined = timing.Trace.combine([one, two]) + self.assertEqual(len(combined.events), 8) + self.assertEqual(combined.total(timing.WAIT), 80.0) + self.assertAlmostEqual(combined.elapsed, one.elapsed + two.elapsed) + + +class TestStderrLogging(unittest.TestCase): + """The zero-code-change path: SINGLESTOREDB_MANAGEMENT_TRACE.""" + + def test_events_are_logged_when_the_option_is_on(self): + mgr = _manager() + err = io.StringIO() + with _tracing(True): + self.assertTrue(timing.recording()) + with contextlib.redirect_stderr(err): + mgr._get(f'clusters/{FAKE_ID}') + out = err.getvalue() + self.assertIn('GET clusters/{id}', out) + self.assertIn('-> 200', out) + + def test_waits_are_logged_too(self): + err = io.StringIO() + with _tracing(True): + with patch('singlestoredb.management.timing.time.sleep'): + with contextlib.redirect_stderr(err): + timing.sleep(20, 'cluster state -> active') + self.assertIn('cluster state -> active', err.getvalue()) + + def test_nothing_is_logged_when_the_option_is_off(self): + mgr = _manager() + err = io.StringIO() + with _tracing(False): + self.assertFalse(timing.recording()) + with contextlib.redirect_stderr(err): + mgr._get('clusters') + self.assertEqual(err.getvalue(), '') + + +if __name__ == '__main__': + unittest.main() diff --git a/singlestoredb/tests/test_management_utils.py b/singlestoredb/tests/test_management_utils.py new file mode 100644 index 000000000..b7b25f1e5 --- /dev/null +++ b/singlestoredb/tests/test_management_utils.py @@ -0,0 +1,1819 @@ +#!/usr/bin/env python +# type: ignore +""" +Version-neutral unit tests for the management API helpers. + +Nothing here touches a version-specific module or needs a management token or +a container. These were originally written alongside the versioned wrappers +only because that is where the bugs were found. +""" +import datetime +import os +import pathlib +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock +from unittest.mock import patch + +from singlestoredb.exceptions import ManagementError +from singlestoredb.management.utils import normalize_remote_path +from singlestoredb.tests.utils import counting_file_space +from singlestoredb.tests.utils import counting_stage + + +TEST_DIR = pathlib.Path(os.path.dirname(__file__)) + + +class TestFolderTransferPaths(unittest.TestCase): + """Folder helpers must address remote objects with the full remote path + and resolve ``ignore`` globs relative to the local folder.""" + + def _make_stage(self): + from singlestoredb.management.stage import Stage + stage = Stage.__new__(Stage) + stage._manager = MagicMock() + return stage + + def _make_file_space(self): + from singlestoredb.management.files import FileSpace + space = FileSpace.__new__(FileSpace) + space._manager = MagicMock() + return space + + def _make_files_object(self, path, type_='file'): + from singlestoredb.management.files import FilesObject + return FilesObject( + name=path.rsplit('/', 1)[-1], + path=path, + size=0, + type=type_, + format='', + mimetype='', + created=None, + last_modified=None, + writable=True, + ) + + def _make_local_tree(self, tmp): + """Create ``/src/keep.py`` and ``/src/sub/skip.pyc``.""" + import os + root = os.path.join(tmp, 'src') + os.makedirs(os.path.join(root, 'sub')) + keep = os.path.join(root, 'keep.py') + skip = os.path.join(root, 'sub', 'skip.pyc') + for path in (keep, skip): + with open(path, 'w') as f: + f.write('x') + return root, keep, skip + + def test_stage_download_folder_prefixes_remote_paths(self): + import tempfile + stage = self._make_stage() + # listdir strips the stage_path prefix from its results + stage.listdir = MagicMock( + return_value=[ + self._make_files_object('a.txt'), + self._make_files_object('sub/b.txt'), + ], + ) + stage.is_dir = MagicMock(side_effect=lambda p: p == 'remote/folder') + stage._download_file = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + stage.download_folder('remote/folder', tmp, overwrite=True) + requested = [call.args[0] for call in stage._download_file.call_args_list] + self.assertEqual( + requested, ['remote/folder/a.txt', 'remote/folder/sub/b.txt'], + ) + + def test_stage_download_folder_normalizes_prefix(self): + import tempfile + stage = self._make_stage() + stage.listdir = MagicMock( + return_value=[self._make_files_object('a.txt')], + ) + # download_folder normalizes './remote/folder/' before probing. + stage.is_dir = MagicMock(side_effect=lambda p: p == 'remote/folder') + stage._download_file = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + stage.download_folder('./remote/folder/', tmp, overwrite=True) + self.assertEqual( + stage._download_file.call_args_list[0].args[0], + 'remote/folder/a.txt', + ) + + def test_stage_download_folder_uses_listing_type_not_is_dir(self): + """The entry type comes from the listing, so no per-entry is_dir + call is made, and empty remote folders are still created locally.""" + import os + import tempfile + stage = self._make_stage() + stage.listdir = MagicMock( + return_value=[ + self._make_files_object('empty', type_='directory'), + self._make_files_object('a.txt'), + ], + ) + is_dir_calls = [] + + def is_dir(p): + is_dir_calls.append(p) + return p == 'remote' + + stage.is_dir = is_dir + stage._download_file = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + dest = os.path.join(tmp, 'dest') + stage.download_folder('remote', dest, overwrite=True) + # Only the top-level folder check, nothing per entry + self.assertEqual(is_dir_calls, ['remote']) + self.assertTrue(os.path.isdir(os.path.join(dest, 'empty'))) + requested = [call.args[0] for call in stage._download_file.call_args_list] + self.assertEqual(requested, ['remote/a.txt']) + + def test_stage_upload_folder_ignores_folder_patterns(self): + import os + import tempfile + stage = self._make_stage() + stage.exists = MagicMock(return_value=False) + stage.upload_file = MagicMock() + stage.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root = os.path.join(tmp, 'src') + os.makedirs(os.path.join(root, '__pycache__')) + keep = os.path.join(root, 'keep.py') + for path in (keep, os.path.join(root, '__pycache__', 'a.pyc')): + with open(path, 'w') as f: + f.write('x') + stage.upload_folder(root, 'dest', ignore='**/__pycache__') + uploaded = [ + call.args[0] for call in stage.upload_file.call_args_list + ] + self.assertEqual(uploaded, [keep]) + + def test_file_space_upload_folder_ignores_folder_patterns(self): + import os + import tempfile + space = self._make_file_space() + space.upload_file = MagicMock() + space.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root = os.path.join(tmp, 'src') + os.makedirs(os.path.join(root, '__pycache__')) + keep = os.path.join(root, 'keep.py') + for path in (keep, os.path.join(root, '__pycache__', 'a.pyc')): + with open(path, 'w') as f: + f.write('x') + space.upload_folder(root, 'dest', ignore='**/__pycache__') + uploaded = [ + call.kwargs['local_path'] + for call in space.upload_file.call_args_list + ] + self.assertEqual(uploaded, [keep]) + + def test_download_folder_defaults_to_remote_folder_name(self): + """With no local_path, the destination is the remote folder's name + in the current directory.""" + import os + import tempfile + cwd = os.getcwd() + for name, obj, attr in ( + ('Stage', self._make_stage(), '_download_file'), + ('FileSpace', self._make_file_space(), '_download_file'), + ): + obj.listdir = MagicMock( + return_value=[self._make_files_object('a.txt')], + ) + obj.is_dir = MagicMock(return_value=True) + setattr(obj, attr, MagicMock()) + with tempfile.TemporaryDirectory() as tmp: + try: + os.chdir(tmp) + obj.download_folder('remote/folder') + finally: + os.chdir(cwd) + target = getattr(obj, attr).call_args_list[0].args[1] + self.assertEqual( + os.path.normpath(target), + os.path.join('folder', 'a.txt'), + f'{name} wrote to {target}', + ) + + def test_download_folder_root_without_local_path_raises(self): + for obj in (self._make_stage(), self._make_file_space()): + obj.listdir = MagicMock(return_value=[]) + obj.is_dir = MagicMock(return_value=True) + with self.assertRaises(ValueError) as ctx: + obj.download_folder('/') + self.assertIn('local_path must be specified', str(ctx.exception)) + + def test_download_folder_explicit_local_path_unchanged(self): + """Explicit local_path keeps writing directly into that directory.""" + import os + import tempfile + for obj in (self._make_stage(), self._make_file_space()): + obj.listdir = MagicMock( + return_value=[self._make_files_object('a.txt')], + ) + obj.is_dir = MagicMock(return_value=True) + obj._download_file = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + dest = os.path.join(tmp, 'dest') + obj.download_folder('remote/folder', dest, overwrite=True) + self.assertEqual( + obj._download_file.call_args_list[0].args[1], + os.path.join(dest, 'a.txt'), + ) + + def test_upload_folder_builds_slash_separated_remote_paths(self): + """Remote paths must use '/' even when the local platform uses '\\'.""" + import tempfile + stage = self._make_stage() + stage.exists = MagicMock(return_value=False) + stage.upload_file = MagicMock() + stage.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root, _, _ = self._make_local_tree(tmp) + stage.upload_folder(root, 'dest/') + targets = sorted( + call.args[1] for call in stage.upload_file.call_args_list + ) + self.assertEqual(targets, ['dest/keep.py', 'dest/sub/skip.pyc']) + for target in targets: + self.assertNotIn('\\', target) + + def test_stage_upload_folder_strips_leading_prefix_segments(self): + """A '/foo' or './foo' prefix must not survive into the remote path. + + ``listdir`` and ``download_folder`` already pass + ``strip_leading=True``, and the stage routes interpolate the path into + a URL, so a leading '/' would produce a doubled slash. + """ + import tempfile + for prefix in ('/dest', './dest'): + stage = self._make_stage() + stage.exists = MagicMock(return_value=False) + stage.upload_file = MagicMock() + stage.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root, _, _ = self._make_local_tree(tmp) + stage.upload_folder(root, prefix) + targets = sorted( + call.args[1] for call in stage.upload_file.call_args_list + ) + self.assertEqual( + targets, ['dest/keep.py', 'dest/sub/skip.pyc'], + f'prefix {prefix!r} produced {targets}', + ) + + def test_file_space_upload_folder_strips_leading_prefix_segments(self): + """``FileSpace._upload`` builds ``files/fs/{location}/{path}``, so a + leading '/' would request ``files/fs///dest/...``.""" + import tempfile + for prefix in ('/dest', './dest'): + space = self._make_file_space() + space.upload_file = MagicMock() + space.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root, _, _ = self._make_local_tree(tmp) + space.upload_folder(root, prefix) + targets = sorted( + call.kwargs['path'] + for call in space.upload_file.call_args_list + ) + self.assertEqual( + targets, ['dest/keep.py', 'dest/sub/skip.pyc'], + f'prefix {prefix!r} produced {targets}', + ) + + def test_stage_upload_folder_applies_ignore_globs(self): + import tempfile + stage = self._make_stage() + stage.exists = MagicMock(return_value=False) + stage.upload_file = MagicMock() + stage.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root, keep, _ = self._make_local_tree(tmp) + stage.upload_folder(root, 'dest', ignore='**/*.pyc') + uploaded = [ + call.args[0] for call in stage.upload_file.call_args_list + ] + self.assertEqual(uploaded, [keep]) + + def test_stage_upload_folder_applies_ignore_globs_to_cwd(self): + import os + import tempfile + stage = self._make_stage() + stage.exists = MagicMock(return_value=False) + stage.upload_file = MagicMock() + stage.info = MagicMock() + cwd = os.getcwd() + with tempfile.TemporaryDirectory() as tmp: + root, _, _ = self._make_local_tree(tmp) + try: + os.chdir(root) + stage.upload_folder('.', 'dest', ignore='**/*.pyc') + finally: + os.chdir(cwd) + uploaded = [ + call.args[0] for call in stage.upload_file.call_args_list + ] + self.assertEqual(uploaded, ['keep.py']) + + def test_file_space_upload_folder_applies_ignore_globs(self): + import tempfile + space = self._make_file_space() + space.upload_file = MagicMock() + space.info = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + root, keep, _ = self._make_local_tree(tmp) + space.upload_folder(root, 'dest', ignore='**/*.pyc') + uploaded = [ + call.kwargs['local_path'] + for call in space.upload_file.call_args_list + ] + self.assertEqual(uploaded, [keep]) + + def test_file_space_upload_folder_applies_ignore_globs_to_cwd(self): + import os + import tempfile + space = self._make_file_space() + space.upload_file = MagicMock() + space.info = MagicMock() + cwd = os.getcwd() + with tempfile.TemporaryDirectory() as tmp: + root, _, _ = self._make_local_tree(tmp) + try: + os.chdir(root) + space.upload_folder('.', 'dest', ignore='**/*.pyc') + finally: + os.chdir(cwd) + uploaded = [ + call.kwargs['local_path'] + for call in space.upload_file.call_args_list + ] + self.assertEqual(uploaded, ['keep.py']) + + +class TestCustomModelUploadPaths(unittest.TestCase): + """``UPLOAD CUSTOM MODEL`` must not replay the local directory tree into + the models space. The handler is hidden (``_enabled = False``) and so has + no live coverage.""" + + def _run(self, local_path): + from singlestoredb.fusion.handlers.models import UploadCustomModelHandler + handler = UploadCustomModelHandler.__new__(UploadCustomModelHandler) + space = MagicMock() + with patch( + 'singlestoredb.fusion.handlers.models.get_file_space', + return_value=space, + ): + handler.run( + dict( + model_name='mymodel', + local_path=local_path, + overwrite=False, + ), + ) + return space + + def test_single_file_uploads_under_the_model_name(self): + import tempfile + with tempfile.TemporaryDirectory() as tmp: + local = os.path.join(tmp, 'nested', 'weights.bin') + os.makedirs(os.path.dirname(local)) + with open(local, 'w') as f: + f.write('x') + space = self._run(local) + space.upload_folder.assert_not_called() + self.assertEqual( + space._upload_local_file.call_args.kwargs['path'], + 'mymodel/weights.bin', + ) + + def test_a_directory_still_goes_through_upload_folder(self): + import tempfile + with tempfile.TemporaryDirectory() as tmp: + space = self._run(tmp) + space._upload_local_file.assert_not_called() + self.assertEqual( + space.upload_folder.call_args.kwargs['path'], 'mymodel', + ) + + +class TestRecursiveDownloadPathTraversal(unittest.TestCase): + """Recursive download helpers must refuse to write outside ``local_path`` + when the remote listing contains traversal segments (``..``).""" + + def _make_file_location(self): + # FileSpace is a concrete FileLocation subclass; instantiate via + # __new__ to skip its constructor (which expects a real FilesManager). + from singlestoredb.management.files import FileSpace + loc = FileSpace.__new__(FileSpace) + loc._manager = MagicMock() + return loc + + def _make_files_object(self, path, type_='file'): + from singlestoredb.management.files import FilesObject + return FilesObject( + name=path.rsplit('/', 1)[-1], + path=path, + size=0, + type=type_, + format='', + mimetype='', + created=None, + last_modified=None, + writable=True, + ) + + def test_files_download_folder_rejects_traversal(self): + import tempfile + loc = self._make_file_location() + # Listing returns an entry whose path escapes via '..' + loc.listdir = MagicMock( + return_value=[self._make_files_object('../escape.txt')], + ) + loc._download_file = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + target = f'{tmp}/dest' + import os + os.makedirs(target) + with self.assertRaises(ManagementError) as ctx: + loc.download_folder('remote', target, overwrite=True) + self.assertIn('outside destination', str(ctx.exception)) + loc._download_file.assert_not_called() + + def test_files_download_folder_rejects_traversal_directory(self): + import tempfile + loc = self._make_file_location() + # Directory entry that escapes + loc.listdir = MagicMock( + return_value=[self._make_files_object('../evil', type_='directory')], + ) + with tempfile.TemporaryDirectory() as tmp: + target = f'{tmp}/dest' + import os + os.makedirs(target) + with self.assertRaises(ManagementError) as ctx: + loc.download_folder('remote', target, overwrite=True) + self.assertIn('outside destination', str(ctx.exception)) + + def test_stage_download_folder_rejects_traversal(self): + import tempfile + from singlestoredb.management.stage import Stage + stage = Stage.__new__(Stage) + stage.listdir = MagicMock( + return_value=[self._make_files_object('../escape.txt')], + ) + # is_dir(stage_path) must return True (it's a directory); the entry + # type in the listing marks each entry as a file. + stage.is_dir = MagicMock(side_effect=lambda p: p == 'remote') + stage._download_file = MagicMock() + with tempfile.TemporaryDirectory() as tmp: + target = f'{tmp}/dest' + import os + os.makedirs(target) + with self.assertRaises(ManagementError) as ctx: + stage.download_folder('remote', target, overwrite=True) + self.assertIn('outside destination', str(ctx.exception)) + stage._download_file.assert_not_called() + + +class TestUploadRoundTrips(unittest.TestCase): + """An upload must not repeat work it has already done. + + The counts pinned here are the Stage / file space half of what + ``UPLOAD FILE TO STAGE`` costs; the two that resolve ``IN ''`` are + made before a ``Stage`` exists and so cannot be seen from here. For the + whole-statement count, add two. + """ + + def _local_file(self, tmp, content='contents'): + local = os.path.join(tmp, 'local.csv') + with open(local, 'w') as f: + f.write(content) + return local + + def test_a_fresh_upload_costs_one_check_and_one_write(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + stage, manager = counting_stage() + obj = stage.upload_file(local, 'remote.csv') + # Was four: upload_file and _upload each checked exists() + self.assertEqual( + manager.calls, [ + ('GET', 'remote.csv'), # exists() + ('PUT', 'remote.csv'), # the upload + ('GET', 'remote.csv'), # info() for the return value + ], + ) + # The public contract still hands back a populated object + self.assertEqual(obj.name, 'remote.csv') + self.assertEqual(obj.path, 'remote.csv') + self.assertEqual(obj.type, 'file') + self.assertEqual(obj.size, 8) + self.assertTrue(obj.writable) + + def test_an_overwrite_costs_one_check_and_one_delete(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + stage, manager = counting_stage(existing=['remote.csv']) + stage.upload_file(local, 'remote.csv', overwrite=True) + # Was six: the duplicated exists() dragged a second remove() check in, + # and then the remaining exists()/is_dir() pair was the same GET twice + self.assertEqual( + manager.calls, [ + ('GET', 'remote.csv'), # the one metadata fetch + ('DELETE', 'remote.csv'), + ('PUT', 'remote.csv'), + ('GET', 'remote.csv'), # info() for the return value + ], + ) + + def test_an_overwrite_of_a_folder_raises_on_the_one_check(self): + # The IsADirectoryError remove() used to raise through _upload is + # raised by _upload itself now, with the same message. + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + stage, manager = counting_stage(existing=['remote.csv/']) + with self.assertRaises(IsADirectoryError) as ctx: + stage.upload_file(local, 'remote.csv', overwrite=True) + self.assertIn('use rmdir or removedirs', str(ctx.exception)) + self.assertEqual(manager.calls, [('GET', 'remote.csv')]) + + def test_a_conflict_still_raises_and_closes_the_local_file(self): + opened = [] + real_open = open + + def recording_open(*args, **kwargs): + handle = real_open(*args, **kwargs) + opened.append(handle) + return handle + + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + stage, manager = counting_stage(existing=['remote.csv']) + with patch('builtins.open', recording_open): + with self.assertRaises(OSError) as ctx: + stage.upload_file(local, 'remote.csv') + self.assertIn('stage path already exists', str(ctx.exception)) + self.assertEqual(manager.calls, [('GET', 'remote.csv')]) + # The conflict is now detected inside _upload, which is after the + # local file has been opened, so that handle has to close on the way + # out rather than wait for the collector + self.assertTrue(opened) + self.assertTrue(all(handle.closed for handle in opened)) + + def test_a_local_directory_is_rejected_before_any_request(self): + with tempfile.TemporaryDirectory() as tmp: + stage, manager = counting_stage() + with self.assertRaises(IsADirectoryError): + stage.upload_file(tmp, 'remote.csv') + self.assertEqual(manager.calls, []) + + def test_the_fusion_path_skips_the_metadata_request(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + stage, manager = counting_stage() + out = stage._upload_local_file(local, 'remote.csv', fetch_info=False) + self.assertIsNone(out) + self.assertEqual( + manager.calls, [('GET', 'remote.csv'), ('PUT', 'remote.csv')], + ) + + def test_the_fusion_handler_takes_that_path(self): + from singlestoredb.fusion.handlers.stage import UploadStageFileHandler + handler = UploadStageFileHandler.__new__(UploadStageFileHandler) + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + stage, manager = counting_stage() + with patch( + 'singlestoredb.fusion.handlers.stage.get_deployment', + return_value=SimpleNamespace(stage=stage), + ): + handler.run( + dict( + local_path=local, + stage_path='remote.csv', + overwrite=False, + ), + ) + self.assertEqual( + manager.calls, [('GET', 'remote.csv'), ('PUT', 'remote.csv')], + ) + + def test_a_file_space_upload_costs_the_same(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + space, manager = counting_file_space() + obj = space.upload_file(local, 'remote.csv') + fresh = list(manager.calls) + + manager.calls.clear() + out = space._upload_local_file( + local, 'other.csv', fetch_info=False, + ) + self.assertEqual( + fresh, [ + ('GET', 'remote.csv'), + ('PUT', 'remote.csv'), + ('GET', 'remote.csv'), + ], + ) + self.assertEqual(obj.type, 'file') + self.assertIsNone(out) + self.assertEqual( + manager.calls, [('GET', 'other.csv'), ('PUT', 'other.csv')], + ) + + def test_a_file_space_conflict_names_the_file_space(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + space, _ = counting_file_space(existing=['remote.csv']) + with self.assertRaises(OSError) as ctx: + space.upload_file(local, 'remote.csv') + self.assertIn('file path already exists', str(ctx.exception)) + + def test_a_file_space_overwrite_also_checks_once(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + space, manager = counting_file_space(existing=['remote.csv']) + space.upload_file(local, 'remote.csv', overwrite=True) + self.assertEqual( + manager.calls, [ + ('GET', 'remote.csv'), + ('DELETE', 'remote.csv'), + ('PUT', 'remote.csv'), + ('GET', 'remote.csv'), + ], + ) + + def test_a_file_space_overwrite_of_a_folder_raises(self): + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + space, manager = counting_file_space(existing=['remote.csv/']) + with self.assertRaises(IsADirectoryError) as ctx: + space.upload_file(local, 'remote.csv', overwrite=True) + self.assertIn('file path is a directory', str(ctx.exception)) + self.assertEqual(manager.calls, [('GET', 'remote.csv')]) + + def test_a_folder_upload_pays_the_saving_per_file(self): + with tempfile.TemporaryDirectory() as tmp: + root = os.path.join(tmp, 'src') + os.makedirs(root) + for name in ('a.csv', 'b.csv'): + with open(os.path.join(root, name), 'w') as f: + f.write('x') + stage, manager = counting_stage(existing=['dest/']) + stage.upload_folder(root, 'dest') + # Two files: one exists() + one PUT + one info() each, plus the + # exists() / is_dir() on the destination and the closing info(). + # Was eleven -- one duplicated exists() per file. + self.assertEqual(len(manager.calls), 9) + self.assertEqual(manager.counts()['PUT'], 2) + + +class TestRemotePathUtils(unittest.TestCase): + """Test cases for remote path normalization (no server required).""" + + def test_local_separators_converted(self): + # A prefix built with os.path.join on Windows keeps a trailing '\' + assert normalize_remote_path('llama3\\') == 'llama3' + assert normalize_remote_path('a\\b\\c.txt') == 'a/b/c.txt' + assert normalize_remote_path(pathlib.PurePosixPath('a/b')) == 'a/b' + + def test_duplicate_and_trailing_separators_collapsed(self): + assert normalize_remote_path('a//b/') == 'a/b' + assert normalize_remote_path('a/b///') == 'a/b' + assert normalize_remote_path('a\\\\b\\') == 'a/b' + + def test_strip_leading(self): + assert normalize_remote_path('./a/b', strip_leading=True) == 'a/b' + assert normalize_remote_path('/a/b', strip_leading=True) == 'a/b' + assert normalize_remote_path('.\\a\\b', strip_leading=True) == 'a/b' + assert normalize_remote_path('/', strip_leading=True) == '' + assert normalize_remote_path('', strip_leading=True) == '' + + def test_strip_leading_off_by_default(self): + assert normalize_remote_path('/a/b') == '/a/b' + + def test_joining_produces_valid_remote_path(self): + # Regression: 'llama3\/file' was produced before normalization + prefix = normalize_remote_path('llama3\\') + assert f'{prefix}/file' == 'llama3/file' + + def test_listdir_style_suffix(self): + # The listdir call sites append '/' after normalizing + assert normalize_remote_path('llama3\\', strip_leading=True) + '/' \ + == 'llama3/' + assert normalize_remote_path('/', strip_leading=True) + '/' == '/' + + +class TestSecretFromDictTimestamps(unittest.TestCase): + """ + Coverage for ``Secret.from_dict`` running its timestamp fields + through ``to_datetime``. + """ + + def test_timestamps_parsed_to_datetime(self): + from singlestoredb.management.organization import Secret + + obj = { + 'secretID': 'sec-1', + 'name': 'my-secret', + 'createdBy': 'user-a', + 'createdAt': '2024-01-01T00:00:00Z', + 'lastUpdatedBy': 'user-b', + 'lastUpdatedAt': '2024-02-15T12:34:56Z', + 'value': 'shh', + 'deletedBy': None, + 'deletedAt': None, + } + sec = Secret.from_dict(obj) + self.assertIsInstance(sec.created_at, datetime.datetime) + self.assertEqual(sec.created_at.year, 2024) + self.assertIsInstance(sec.last_updated_at, datetime.datetime) + self.assertEqual(sec.last_updated_at.minute, 34) + self.assertIsNone(sec.deleted_at) + + def test_missing_timestamps_become_none(self): + from singlestoredb.management.organization import Secret + + obj = { + 'secretID': 'sec-1', + 'name': 'my-secret', + 'createdBy': 'user-a', + 'lastUpdatedBy': 'user-b', + } + sec = Secret.from_dict(obj) + self.assertIsNone(sec.created_at) + self.assertIsNone(sec.last_updated_at) + self.assertIsNone(sec.deleted_at) + + +class TestTTLProperty(unittest.TestCase): + """A ttl_property caches per instance, not per class.""" + + @staticmethod + def _counter_class(): + from singlestoredb.management.utils import ttl_property + + class Counter: + def __init__(self): + self.calls = 0 + + @ttl_property(datetime.timedelta(hours=1)) + def value(self): + self.calls += 1 + return self.calls + + return Counter + + def test_repeated_reads_are_served_from_the_cache(self): + obj = self._counter_class()() + self.assertEqual(obj.value, 1) + self.assertEqual(obj.value, 1) + self.assertEqual(obj.calls, 1) + + def test_each_instance_caches_its_own_value(self): + # Two managers may hold different tokens, so one must never be served + # the other's copy. + cls = self._counter_class() + first, second = cls(), cls() + self.assertEqual(first.value, 1) + self.assertEqual(second.value, 1) + self.assertEqual(first.calls, 1) + self.assertEqual(second.calls, 1) + + def test_an_expired_value_is_refetched(self): + cls = self._counter_class() + obj = cls() + self.assertEqual(obj.value, 1) + type(obj).__dict__['value'].ttl = datetime.timedelta(0) + self.assertEqual(obj.value, 2) + + def test_reset_discards_the_cached_value(self): + cls = self._counter_class() + obj = cls() + self.assertEqual(obj.value, 1) + type(obj).__dict__['value'].reset(obj) + self.assertEqual(obj.value, 2) + + +class TestManagerTransport(unittest.TestCase): + """ + Retries and timeouts on the session every manager shares. + + The long ``wait_on_active`` loops poll for twenty minutes, and a + keep-alive connection the far end closed while the client slept surfaces + as ``RemoteDisconnected`` on the next poll -- which used to fail the whole + operation, leaving a live cluster behind. + """ + + def _manager(self): + from singlestoredb.management.manager import Manager + return Manager(access_token='fake-token', base_url='https://example.com') + + def test_retries_are_mounted_for_both_schemes(self): + mgr = self._manager() + for prefix in ('http://', 'https://'): + retries = mgr._sess.get_adapter(prefix + 'x').max_retries + self.assertGreater(retries.total, 0) + + def test_post_is_not_replayed(self): + # A dropped connection does not say whether the server acted on the + # request, and a replayed POST /clusters deploys twice. + retries = self._manager()._sess.get_adapter('https://x').max_retries + self.assertNotIn('POST', retries.allowed_methods) + self.assertIn('GET', retries.allowed_methods) + self.assertIn('DELETE', retries.allowed_methods) + + def test_transient_statuses_are_retried(self): + retries = self._manager()._sess.get_adapter('https://x').max_retries + for status in (429, 502, 503, 504): + self.assertIn(status, retries.status_forcelist) + self.assertNotIn(404, retries.status_forcelist) + # _check has to be the one to raise, so it can quote the body. + self.assertFalse(retries.raise_on_status) + + def test_a_default_timeout_is_applied(self): + mgr = self._manager() + mgr._sess.get = MagicMock() + mgr._doit('get', 'clusters') + self.assertEqual( + mgr._sess.get.call_args[1]['timeout'], + (10.0, 180.0), + ) + + def test_an_explicit_timeout_wins(self): + mgr = self._manager() + mgr._sess.get = MagicMock() + mgr._doit('get', 'clusters', timeout=1) + self.assertEqual(mgr._sess.get.call_args[1]['timeout'], 1) + + def test_a_transport_failure_names_the_route(self): + import requests + + mgr = self._manager() + mgr._sess.get = MagicMock( + side_effect=requests.exceptions.ConnectionError( + 'Connection aborted.', + ), + ) + with self.assertRaises(ManagementError) as cm: + mgr._doit('get', 'clusters/abc') + + msg = str(cm.exception) + self.assertIn('ConnectionError', msg) + self.assertIn('GET', msg) + self.assertIn('clusters/abc', msg) + + +class TestWaitOnEndpoint(unittest.TestCase): + """ + ``Manager._wait_on_endpoint`` polls a new deployment by connecting to it. + + It only runs inside the notebook environment, which is why the loop having + no exit on success went unnoticed: a successful connect fell out of the + ``try`` and straight back into ``while True``, so the only ways out were an + access-denied error or the timeout. + """ + + def _manager(self): + from singlestoredb.management.manager import Manager + mgr = Manager(access_token='fake-token', base_url='https://example.com') + mgr.obj_type = 'cluster' + return mgr + + def test_a_successful_connect_ends_the_wait(self): + mgr = self._manager() + out = MagicMock() + + with patch.dict( + os.environ, {'SINGLESTOREDB_WORKLOAD_TYPE': 'notebook'}, + ): + with patch('singlestoredb.management.timing.time.sleep'): + result = mgr._wait_on_endpoint(out, interval=1, timeout=10) + + self.assertIs(result, out) + # Once, not until the timeout ran out. + self.assertEqual(out.connect.call_count, 1) + + def test_nothing_is_waited_on_outside_the_notebook_environment(self): + mgr = self._manager() + out = MagicMock() + + with patch.dict(os.environ, {'SINGLESTOREDB_WORKLOAD_TYPE': ''}): + result = mgr._wait_on_endpoint(out, interval=1, timeout=10) + + self.assertIs(result, out) + out.connect.assert_not_called() + + def test_a_refused_connection_is_retried_until_the_timeout(self): + mgr = self._manager() + out = MagicMock() + out.connect = MagicMock(side_effect=OSError('connection refused')) + + with patch.dict( + os.environ, {'SINGLESTOREDB_WORKLOAD_TYPE': 'notebook'}, + ): + with patch('singlestoredb.management.timing.time.sleep'): + with self.assertRaises(ManagementError) as cm: + mgr._wait_on_endpoint(out, interval=10, timeout=30) + + self.assertIn('endpoint', str(cm.exception)) + self.assertEqual(out.connect.call_count, 4) + + +class TestDeploymentTracking(unittest.TestCase): + """ + The sweeper in ``tests/utils.py`` that keeps test runs from leaking + billable deployments. + """ + + def setUp(self): + from singlestoredb.tests import utils + self.utils = utils + self.saved = list(utils._tracked) + utils._tracked.clear() + self.saved_in_flight = list(utils._in_flight) + utils._in_flight.clear() + self.addCleanup(self._restore) + self.owner = utils.get_owner() + self.addCleanup(lambda: utils.set_owner(self.owner)) + + def _restore(self): + self.utils._tracked.clear() + self.utils._tracked.extend(self.saved) + self.utils._in_flight.clear() + self.utils._in_flight.extend(self.saved_in_flight) + + def _deployment(self, name, terminated_at=None, state='ACTIVE'): + """A stand-in that is not a Mock, so tracking does not skip it.""" + class Deployment: + def __init__(self): + self.name = name + self.id = name + self.terminated_at = terminated_at + self.state = state + self.terminated_with = None + self._manager = object() + + def refresh(self): + return self + + def terminate(self, force=False): + self.terminated_with = force + + return Deployment() + + def test_mocked_deployments_are_not_tracked(self): + # The unit tests create objects from patched _post calls; sweeping + # those would be a round trip and a warning per fake object. + self.utils.track(MagicMock()) + self.assertEqual(self.utils._tracked, []) + + def test_a_tracked_deployment_is_terminated_with_force(self): + obj = self._deployment('wg-1') + self.utils.track(obj) + self.assertEqual(len(self.utils.cleanup_tracked()), 1) + self.assertTrue(obj.terminated_with) + self.assertEqual(self.utils._tracked, []) + + def test_an_already_terminated_deployment_is_left_alone(self): + obj = self._deployment('wg-1', terminated_at='2026-01-01T00:00:00Z') + self.utils.track(obj) + self.assertEqual(self.utils.cleanup_tracked(), []) + self.assertIsNone(obj.terminated_with) + + def test_a_deployment_that_no_longer_exists_is_left_alone(self): + obj = self._deployment('wg-1') + obj.refresh = MagicMock( + side_effect=ManagementError(errno=404, msg='not found'), + ) + self.utils.track(obj) + self.assertEqual(self.utils.cleanup_tracked(), []) + self.assertIsNone(obj.terminated_with) + self.assertEqual(self.utils._tracked, []) + + def test_a_refresh_that_fails_transiently_is_still_terminated(self): + """Only a 404 means gone. Guessing "gone" on a 503 would skip the + termination and leave the deployment running and billing.""" + for exc in ( + ManagementError(errno=503, msg='service unavailable'), + KeyError('connection dropped'), + ): + obj = self._deployment('wg-1') + obj.refresh = MagicMock(side_effect=exc) + self.utils.track(obj) + self.assertEqual( + self.utils.cleanup_tracked(), ["Deployment 'wg-1'"], + f'{exc!r} was taken as already gone', + ) + self.assertTrue(obj.terminated_with) + + def test_children_are_terminated_before_their_parents(self): + group = self._deployment('wg-1') + space = self._deployment('ws-1') + order = [] + for obj in (group, space): + obj.terminate = lambda force=False, obj=obj: order.append(obj.name) + self.utils.track(group) + self.utils.track(space) + self.utils.cleanup_tracked() + self.assertEqual(order, ['ws-1', 'wg-1']) + + def test_a_sweep_is_limited_to_one_owner(self): + self.utils.set_owner('mod.ClassA') + first = self.utils.track(self._deployment('a')) + self.utils.set_owner('mod.ClassB') + second = self.utils.track(self._deployment('b')) + + self.assertEqual(self.utils.cleanup_tracked('mod.ClassA'), ["Deployment 'a'"]) + self.assertTrue(first.terminated_with) + self.assertIsNone(second.terminated_with) + + # ... and the rest still goes at the end of the session. + self.assertEqual(len(self.utils.cleanup_tracked()), 1) + self.assertTrue(second.terminated_with) + + def test_a_failed_termination_does_not_stop_the_sweep(self): + first = self._deployment('a') + first.terminate = MagicMock(side_effect=RuntimeError('boom')) + second = self._deployment('b') + self.utils.track(first) + self.utils.track(second) + + # Nothing raises: this runs outside any test, where an exception is + # reported against whatever happens to run next. + self.assertEqual(self.utils.cleanup_tracked(), ["Deployment 'b'"]) + self.assertTrue(second.terminated_with) + + # 'a' stays tracked so the end-of-session sweep retries it. Dropping it + # here is how one transient error used to leak a cluster for good. + self.assertEqual(self.utils.tracked_labels(), ["Deployment 'a'"]) + first.terminate = MagicMock() + self.assertEqual(self.utils.cleanup_tracked(), ["Deployment 'a'"]) + self.assertEqual(self.utils.tracked_labels(), []) + + def test_a_create_that_fails_while_waiting_still_tracks_the_orphan(self): + """The headline leak: a creator makes the deployment and only then + waits for it, so a wait that times out raises after the server has a + live cluster. Tracking wraps the return value, so without recovery + nothing registers it -- silently, with no summary line.""" + orphan = self._deployment('cl-test-shared-0-abc') + receiver = SimpleNamespace(clusters=[orphan]) + + def create_then_fail_waiting(recv, name, **kwargs): + # What create_cluster does: the cluster exists by now, and the + # wait is what raises. + raise ManagementError(msg=f'Exceeded waiting time for {name}') + + wrapped = self.utils._tracking_wrapper( + create_then_fail_waiting, lambda recv: recv.clusters, + ) + with self.assertRaises(ManagementError): + wrapped(receiver, 'cl-test-shared-0-abc', wait_on_active=True) + + self.assertEqual( + self.utils.tracked_labels(), + [ + "Deployment 'cl-test-shared-0-abc' (left behind by a failed " + 'create)', + ], + ) + self.assertEqual(len(self.utils.cleanup_tracked()), 1) + self.assertTrue(orphan.terminated_with) + + def test_an_interrupt_during_the_wait_also_recovers_the_orphan(self): + """Ctrl-C during wait_on_active leaves the same live cluster a timeout + does, so the wrapper catches BaseException rather than Exception.""" + orphan = self._deployment('cl-1') + receiver = SimpleNamespace(clusters=[orphan]) + + def interrupted(recv, name, **kwargs): + raise KeyboardInterrupt + + wrapped = self.utils._tracking_wrapper( + interrupted, lambda recv: recv.clusters, + ) + with self.assertRaises(KeyboardInterrupt): + wrapped(receiver, 'cl-1') + self.assertEqual(len(self.utils._tracked), 1) + + def test_a_mocked_receiver_does_not_track_what_it_returns(self): + """A unit test's stubbed ``get_cluster`` hands back a real Cluster + whose ``_manager`` is None, or a bare sentinel. ``track()`` calls + anything it cannot place real -- rightly, since guessing "fake" leaks + a billable cluster -- so it would register both, and the end-of-session + summary would report phantom live deployments. The receiver's verdict + is what decides.""" + mgr = SimpleNamespace( + _get=MagicMock(), _post=MagicMock(), _delete=MagicMock(), + ) + returned = self._deployment('my-cluster') + returned._manager = None + + for value in (returned, 'sentinel'): + wrapped = self.utils._tracking_wrapper( + lambda recv, name, value=value, **kwargs: value, + lambda recv: [], + ) + self.assertIs(wrapped(mgr, 'my-cluster'), value) + + self.assertEqual(self.utils.tracked_labels(), []) + + def test_a_real_receiver_still_tracks_what_it_returns(self): + """The other side of the check above: an unrecognisable return value + from a real manager is still swept, because a cluster left running + costs money and a redundant terminate costs one round trip.""" + mgr = SimpleNamespace(_get=object(), _post=object(), _delete=object()) + returned = self._deployment('cl-1') + returned._manager = None + + wrapped = self.utils._tracking_wrapper( + lambda recv, name, **kwargs: returned, lambda recv: [], + ) + wrapped(mgr, 'cl-1') + self.assertEqual(self.utils.tracked_labels(), ["Deployment 'cl-1'"]) + + def test_a_mocked_receiver_is_not_searched_for_orphans(self): + """The unit tests drive these creators with patched transports; a + failure there names nothing real to recover.""" + def boom(recv, name, **kwargs): + raise ManagementError(msg='boom') + + wrapped = self.utils._tracking_wrapper(boom, lambda recv: recv.clusters) + with self.assertRaises(ManagementError): + wrapped(MagicMock(), 'cl-1') + self.assertEqual(self.utils._tracked, []) + + def test_a_real_manager_with_a_patched_post_is_not_searched_either(self): + """The receiver of these creators is the manager itself, which has no + ``_manager``. Checking for one made a real manager with a patched + ``_post`` -- what the unit tests drive -- read as live, so the recovery + fired an actual management API GET from a unit test.""" + receiver = SimpleNamespace(_get=object(), _post=MagicMock()) + + def boom(recv, name, **kwargs): + raise ManagementError(msg='boom') + + def finder(recv): + raise AssertionError('recovery called the live API') + + wrapped = self.utils._tracking_wrapper(boom, finder) + with self.assertRaises(ManagementError): + wrapped(receiver, 'cl-1') + self.assertEqual(self.utils._tracked, []) + self.assertEqual(self.utils._in_flight, []) + + def test_a_create_killed_mid_wait_is_recovered_by_the_sweep(self): + """The second half of the same leak: a create that has POSTed and is + blocked in ``wait_on_active`` is not tracked yet, and the wrapper's + ``except`` never runs if the process is killed. So the shutdown sweep + recovers whatever is in flight before it walks ``_tracked``.""" + orphan = self._deployment('cl-1') + receiver = SimpleNamespace(clusters=[orphan]) + + def create_then_wait(recv, name, **kwargs): + # Stands in for the sweep firing from SIGTERM/atexit while the + # wait is still blocked. + self.assertEqual(len(self.utils._in_flight), 1) + self.utils.recover_in_flight() + raise AssertionError('the process would have been killed here') + + wrapped = self.utils._tracking_wrapper( + create_then_wait, lambda recv: recv.clusters, + ) + with self.assertRaises(AssertionError): + wrapped(receiver, 'cl-1', wait_on_active=True) + + # Recovered once, not twice: the entry is popped as it is drained, so + # the wrapper's own except finds nothing left to recover. + self.assertEqual( + self.utils.tracked_labels(), + ["Deployment 'cl-1' (left behind by a failed create)"], + ) + self.assertEqual(self.utils._in_flight, []) + + def test_a_finished_create_leaves_nothing_in_flight(self): + receiver = SimpleNamespace(clusters=[]) + + def finder(recv): + return recv.clusters + + made = self._deployment('cl-1') + wrapped = self.utils._tracking_wrapper( + lambda recv, name, **kwargs: made, finder, + ) + self.assertIs(wrapped(receiver, 'cl-1'), made) + self.assertEqual(self.utils._in_flight, []) + self.assertEqual(len(self.utils._tracked), 1) + + def boom(recv, name, **kwargs): + raise ManagementError(msg='boom') + + receiver.clusters = [self._deployment('cl-2')] + with self.assertRaises(ManagementError): + self.utils._tracking_wrapper(boom, finder)(receiver, 'cl-2') + self.assertEqual(self.utils._in_flight, []) + # The orphan was recovered once, not once per code path. + self.assertEqual(len(self.utils._tracked), 2) + + def test_a_mocked_receiver_never_enters_the_in_flight_list(self): + def create(recv, name, **kwargs): + raise AssertionError(str(self.utils._in_flight)) + + wrapped = self.utils._tracking_wrapper( + create, lambda recv: recv.clusters, + ) + with self.assertRaises(AssertionError) as raised: + wrapped(MagicMock(), 'cl-1') + self.assertEqual(str(raised.exception), '[]') + self.assertEqual(self.utils._in_flight, []) + + def test_recovering_nothing_in_flight_is_a_no_op(self): + self.utils.recover_in_flight() + self.assertEqual(self.utils._tracked, []) + + def test_orphan_recovery_matches_on_the_name_keyword_too(self): + orphan = self._deployment('cl-1') + receiver = SimpleNamespace(clusters=[self._deployment('other'), orphan]) + self.utils._recover_orphan( + receiver, lambda recv: recv.clusters, (), {'name': 'cl-1'}, + ) + self.assertEqual(len(self.utils._tracked), 1) + self.assertIs(self.utils._tracked[0][2], orphan) + + def test_orphan_recovery_never_raises(self): + """It runs while the caller's exception is propagating, so a failure + here must not replace the real error.""" + receiver = SimpleNamespace() + self.utils._recover_orphan( + receiver, + lambda recv: recv.clusters, # AttributeError + ('cl-1',), + {}, + ) + self.assertEqual(self.utils._tracked, []) + + def test_untrack_drops_a_deployment(self): + obj = self.utils.track(self._deployment('a')) + self.utils.untrack(obj) + self.assertEqual(self.utils.cleanup_tracked(), []) + self.assertIsNone(obj.terminated_with) + + def test_a_mocked_creation_is_recognised_by_its_manager(self): + mgr = MagicMock() + self.assertTrue(self.utils._creator_is_mocked(mgr)) + + real = SimpleNamespace(_get=object(), _post=object(), _delete=object()) + self.assertFalse(self.utils._creator_is_mocked(real)) + + # A created object reaches its manager through _manager. + self.assertTrue( + self.utils._creator_is_mocked(SimpleNamespace(_manager=mgr)), + ) + self.assertFalse( + self.utils._creator_is_mocked(SimpleNamespace(_manager=real)), + ) + + def test_a_deployment_without_a_manager_is_still_tracked(self): + # The fail-safe bias: an unrecognisable object counts as real. A fake + # deployment swept is a round trip and a warning, whereas a real one + # skipped is a cluster left running and billing. + obj = self._deployment('a') + obj._manager = None + self.utils.track(obj) + self.assertEqual(len(self.utils._tracked), 1) + + self.assertFalse(self.utils._is_mocked(obj)) + self.assertFalse(self.utils._is_mocked(SimpleNamespace(name='b'))) + + def test_every_creation_method_is_wrapped(self): + # A rename that silently stops tracking is how a cluster leaks. + import importlib + + self.utils.install_deployment_tracking() + for module_name, class_name, method_name, _ in self.utils._CREATORS: + klass = getattr(importlib.import_module(module_name), class_name) + method = getattr(klass, method_name, None) + self.assertIsNotNone( + method, f'{class_name}.{method_name} no longer exists', + ) + self.assertTrue( + hasattr(method, '__wrapped__'), + f'{class_name}.{method_name} is not tracked', + ) + + def test_every_creator_takes_name_first_and_has_a_finder(self): + """``_recover_orphan`` reads the name from the first argument and + searches the collection the finder returns, so both have to hold.""" + import importlib + import inspect + + for module_name, class_name, method_name, finder in \ + self.utils._CREATORS: + klass = getattr(importlib.import_module(module_name), class_name) + method = getattr(klass, method_name) + params = list( + inspect.signature( + getattr(method, '__wrapped__', method), + ).parameters, + ) + self.assertEqual( + params[:2], ['self', 'name'], + f'{class_name}.{method_name} no longer takes name first, so ' + 'a failed create would not be recoverable', + ) + self.assertTrue(callable(finder)) + + +class TestSharedClusterPool(unittest.TestCase): + """ + The pool in ``tests/utils.py`` that keeps the Stage and Job suites from + deploying a cluster apiece. + """ + + def setUp(self): + from singlestoredb.tests import utils + self.utils = utils + + self.saved_pool = list(utils._pool) + self.saved_skip = utils._pool_skip + self.saved_tracked = list(utils._tracked) + self.saved_owner = utils.get_owner() + utils._pool.clear() + utils._pool_skip = None + utils._tracked.clear() + self.addCleanup(self._restore) + + self.created = [] + + def _restore(self): + self.utils._pool[:] = self.saved_pool + self.utils._pool_skip = self.saved_skip + self.utils._tracked[:] = self.saved_tracked + self.utils.set_owner(self.saved_owner) + + def _manager(self, regions=('US East 1',), projects=('STANDARD',)): + """ + A stand-in cluster manager. + + Not a ``Mock``: ``utils.track`` skips anything that came out of a + mocked manager, and the owner a pool cluster is tracked under is the + whole point of the pool. ``create_cluster`` calls ``track`` itself + because that is what ``install_deployment_tracking`` does to the real + method. + """ + created = self.created + utils = self.utils + + class Region: + def __init__(self, name): + self.name = name + self.region_name = name + + class Project: + def __init__(self, edition): + self.edition = edition + self.id = f'project-{edition}' + + class Cluster: + def __init__(self, name, kwargs): + self.name = name + self.id = f'id-of-{name}' + self.terminated_at = None + self.state = 'ACTIVE' + self.kwargs = kwargs + self._manager = object() + + def refresh(self): + return self + + def terminate(self, force=False): + pass + + # Bound outside the class body: a comprehension there cannot see the + # enclosing function's names. + region_list = [Region(x) for x in regions] + project_list = [Project(x) for x in projects] + + class Manager: + regions = region_list + projects = project_list + + def create_cluster(self, name, **kwargs): + # The owner in force at creation time is what decides whether + # the per-class sweep eats the pool. + created.append((name, utils.get_owner(), kwargs)) + return utils.track(Cluster(name, kwargs)) + + return Manager() + + def _patched(self, **kwargs): + import singlestoredb as s2 + return patch.object(s2, 'manage_clusters', return_value=self._manager(**kwargs)) + + def test_the_pool_is_built_once(self): + with self._patched(): + first = self.utils.shared_clusters(2) + second = self.utils.shared_clusters(2) + + self.assertEqual([x.id for x in first], [x.id for x in second]) + self.assertEqual(len(self.created), 2) + + def test_the_pool_grows_to_the_largest_request(self): + with self._patched(): + one = self.utils.shared_clusters(1) + two = self.utils.shared_clusters(2) + + # The second call adds a cluster rather than replacing the first. + self.assertEqual(len(self.created), 2) + self.assertEqual(two[0].id, one[0].id) + self.assertEqual(len(two), 2) + + def test_pool_clusters_are_tracked_under_the_empty_owner(self): + # A pool cluster tracked under the class that asked for it first would + # be terminated by conftest's per-class sweep the moment the run moved + # on -- so the pool would die after one consumer. + self.utils.set_owner('mod.ClassA') + with self._patched(): + self.utils.shared_clusters(2) + + self.assertEqual([x[1] for x in self.created], ['', '']) + self.assertEqual([x[0] for x in self.utils._tracked], ['', '']) + + # The owner the caller was running under is put back... + self.assertEqual(self.utils.get_owner(), 'mod.ClassA') + # ... and a sweep of that class leaves the pool alone. + self.assertEqual(self.utils.cleanup_tracked('mod.ClassA'), []) + self.assertEqual(len(self.utils._tracked), 2) + # Only the end-of-session sweep, which matches every owner, takes it. + self.assertEqual(len(self.utils.cleanup_tracked()), 2) + + def test_pool_names_are_swept_by_the_maintenance_script(self): + from singlestoredb.tests import cleanup_deployments + + with self._patched(): + self.utils.shared_clusters(1) + + self.assertTrue( + cleanup_deployments.is_test_deployment(self.created[0][0]), + self.created[0][0], + ) + # POST /v2/clusters caps a name at 32 characters. + self.assertLessEqual(len(self.created[0][0]), 32) + + def test_a_pool_cluster_is_deployed_where_its_consumers_deployed_theirs(self): + with self._patched(): + self.utils.shared_clusters(1) + + _, _, kwargs = self.created[0] + self.assertEqual(kwargs['size'], 'S-00') + self.assertEqual(kwargs['project'], 'project-STANDARD') + self.assertEqual(kwargs['firewall_ranges'], ['0.0.0.0/0']) + self.assertTrue(kwargs['wait_on_active']) + + def test_no_us_region_skips_rather_than_failing(self): + with self._patched(regions=('EU West 1',)): + with self.assertRaises(unittest.SkipTest): + self.utils.shared_clusters(1) + + # Cached: the next class to ask skips without repeating the + # lookups, and nothing was deployed. + with self.assertRaises(unittest.SkipTest): + self.utils.shared_clusters(1) + + self.assertEqual(self.created, []) + + def test_no_standard_project_skips_rather_than_failing(self): + with self._patched(projects=('SHARED',)): + with self.assertRaises(unittest.SkipTest) as cm: + self.utils.shared_clusters(1) + self.assertIn('SINGLESTOREDB_TEST_PROJECT', str(cm.exception)) + self.assertEqual(self.created, []) + + def test_an_explicit_project_does_not_need_a_standard_one(self): + with patch.dict( + os.environ, {'SINGLESTOREDB_TEST_PROJECT': 'chosen-project'}, + ): + with self._patched(projects=('SHARED',)): + self.utils.shared_clusters(1) + + self.assertEqual(self.created[0][2]['project'], 'chosen-project') + + +class TestClearStage(unittest.TestCase): + """ + Emptying a pooled deployment's stage, which is what lets a class that + asserts exact stage listings borrow a cluster another class has used. + """ + + def setUp(self): + from singlestoredb.tests import utils + self.utils = utils + + def _deployment(self, entries, failing=()): + removed = [] + + class Obj: + def __init__(self, path, type): + self.path = path + self.type = type + + class Stage: + def listdir(self, path='/', *, recursive=False, return_objects=False): + assert return_objects + return [Obj(p, t) for p, t in entries] + + def remove(self, path): + if path in failing: + raise OSError('nope') + removed.append(('remove', path)) + + def removedirs(self, path): + if path in failing: + raise OSError('nope') + removed.append(('removedirs', path)) + + class Deployment: + stage = Stage() + + return Deployment(), removed + + def test_files_are_removed_and_folders_go_recursively(self): + deployment, removed = self._deployment( + [('test.sql', 'file'), ('data/', 'directory')], + ) + self.utils.clear_stage(deployment) + self.assertEqual( + removed, [('remove', 'test.sql'), ('removedirs', 'data/')], + ) + + def test_a_path_that_will_not_go_does_not_stop_the_rest(self): + deployment, removed = self._deployment( + [('stuck.sql', 'file'), ('test.sql', 'file')], + failing=('stuck.sql',), + ) + self.utils.clear_stage(deployment) + self.assertEqual(removed, [('remove', 'test.sql')]) + + +class TestLeftoverDeploymentPatterns(unittest.TestCase): + """ + The maintenance sweep runs against a real organization, so it must match + the names the suite generates and nothing else. + """ + + def setUp(self): + from singlestoredb.tests import cleanup_deployments + self.mod = cleanup_deployments + + def test_generated_names_match(self): + for name in ( + 'wg-test-abcDEF_12', + 'ws-test-abcDEF-x', + 'cl-test-abcDEF', + 'cl-test-shared-0-deadbeef', + 'starter-ws-test-abcDEF', + 'starter-cl-test-abcDEF', + 'A Fusion Testing deadbeefdeadbeef', + 'C Fusion Testing deadbeef', + 'd-fusion-cluster-deadbeef', + 'jobs-fusion-deadbeef', + 'stage-fusion-2-deadbeef', + 'Create WG Test deadbeefdeadbeef', + # The decimal id(self) that test named it with before, so groups + # stranded by older runs are still reachable + 'Create WG Test 140234981234', + ): + self.assertTrue(self.mod.is_test_deployment(name), name) + + def test_retired_names_still_match(self): + # main still creates these, and it carries no sweep at all, so they + # keep arriving. Stranded deployments are billed whichever revision + # made them. + for name in ( + 'Stage Fusion Testing 1 f00e4647f2c664fb', + 'Stage Fusion Testing 2 f00e4647f2c664fb', + 'Files Fusion Testing 1beb5e18ba06e135', + # Unattributed -- no revision here generates it -- but present in + # the organization and swept on the owner's say-so + 'Group 3fed3756', + 'Group 3fed37563fed3756', + ): + self.assertTrue(self.mod.is_test_deployment(name), name) + + def test_names_a_person_chose_do_not_match(self): + for name in ( + None, + '', + 'my-production-cluster', + 'wg-test', + 'prod wg-test-x', + 'analytics-fusion-cluster', + 'Fusion Testing', + 'a-fusion-cluster-deadbeef-prod', + # The 'Group ' pattern must not reach a name a person or the + # portal produced -- that is someone's live workspace group + 'Group 1', + 'Group 2', + 'Group deadbeef prod', + ): + self.assertFalse(self.mod.is_test_deployment(name), name) + + def _cluster(self, name, hours=None, naive=False): + # A naive created_at is what the API sends when it omits the zone: the + # instant is still UTC, the tzinfo is just missing. + now = datetime.datetime.now(tz=datetime.timezone.utc) + if naive: + now = now.replace(tzinfo=None) + + class Cluster: + def __init__(self): + self.name = name + self.id = name + self.terminated_at = None + self.created_at = ( + None if hours is None + else now - datetime.timedelta(hours=hours) + ) + + return Cluster() + + def _find(self, clusters, **kwargs): + """Run find_leftovers against a fixed cluster list.""" + import singlestoredb as s2 + + manager = MagicMock() + manager.clusters = clusters + manager.starter_clusters = [] + with patch.object( + s2, 'manage_clusters', return_value=manager, + ), patch.object( + s2, 'manage_workspaces', side_effect=RuntimeError('no v1'), + ): + found, spared, self.unmatched = self.mod.find_leftovers(**kwargs) + return [x[1].name for x in found], spared + + def test_the_age_filter_spares_a_deployment_a_live_run_may_own(self): + # A parallel run's fixtures are named exactly like stranded ones, so + # age is the only thing keeping this from killing them mid-test. + old = self._cluster('cl-test-old', hours=5) + new = self._cluster('cl-test-new', hours=0.5) + terminated = self._cluster('cl-test-gone', hours=5) + terminated.terminated_at = 'yes' + + names, spared = self._find([old, new, terminated], older_than=2) + + self.assertEqual(names, ['cl-test-old']) + self.assertEqual(len(spared), 1) + self.assertIn('cl-test-new', spared[0]) + + def test_the_default_spares_anything_a_run_could_still_own(self): + # Not zero: a default that swept every match would make running this + # during a test run destructive. + self.assertGreaterEqual(self.mod.DEFAULT_MIN_AGE_HOURS, 1) + names, spared = self._find([ + self._cluster('cl-test-mid-run', hours=1), + ]) + self.assertEqual(names, []) + self.assertEqual(len(spared), 1) + + def test_an_unreported_creation_time_is_spared_by_default(self): + names, spared = self._find([self._cluster('cl-test-ageless')]) + self.assertEqual(names, []) + self.assertIn('cl-test-ageless', spared[0]) + + names, _ = self._find( + [self._cluster('cl-test-ageless')], include_unknown_age=True, + ) + self.assertEqual(names, ['cl-test-ageless']) + + def test_a_naive_timestamp_is_read_as_utc(self): + # Reading it as local time would overstate the age east of UTC and + # sweep a deployment a live run owns. + obj = self._cluster('cl-test-naive', hours=1, naive=True) + self.assertAlmostEqual(self.mod._age_hours(obj), 1, delta=0.1) + + def test_an_unrecognized_name_is_reported_not_swept(self): + # The failure this guards against is silent accumulation: a test that + # names a deployment outside PATTERNS leaves strays the sweep reports + # as 'none found'. + names, _ = self._find([ + self._cluster('cl-test-known', hours=10), + self._cluster('some-persons-cluster', hours=10), + ]) + self.assertEqual(names, ['cl-test-known']) + self.assertEqual(len(self.unmatched), 1) + self.assertIn('some-persons-cluster', self.unmatched[0]) + self.assertIn('10.0h old', self.unmatched[0]) + + def test_a_terminated_deployment_is_not_reported_as_unrecognized(self): + obj = self._cluster('some-persons-cluster', hours=10) + obj.terminated_at = 'yes' + names, _ = self._find([obj]) + self.assertEqual(names, []) + self.assertEqual(self.unmatched, []) + + def test_zero_sweeps_everything_matched(self): + names, spared = self._find( + [self._cluster('cl-test-brand-new', hours=0)], older_than=0, + ) + self.assertEqual(names, ['cl-test-brand-new']) + self.assertEqual(spared, []) + + def test_since_inverts_the_age_filter(self): + # --since is for clearing out a recent session, so it must select what + # the age guard rejects and reject what the age guard selects. + recent = self._cluster('cl-test-today', hours=2) + stale = self._cluster('cl-test-last-week', hours=24 * 7) + + names, spared = self._find( + [recent, stale], + since=datetime.datetime.now(tz=datetime.timezone.utc) + - datetime.timedelta(hours=30), + ) + + self.assertEqual(names, ['cl-test-today']) + self.assertEqual(len(spared), 1) + self.assertIn('cl-test-last-week', spared[0]) + + def test_since_ignores_older_than(self): + # Both guards applying would leave a window nothing falls into, so a + # caller passing --since gets the calendar cutoff alone. + names, _ = self._find( + [self._cluster('cl-test-today', hours=1)], + older_than=self.mod.DEFAULT_MIN_AGE_HOURS, + since=datetime.datetime.now(tz=datetime.timezone.utc) + - datetime.timedelta(hours=30), + ) + self.assertEqual(names, ['cl-test-today']) + + def test_since_still_spares_an_unreported_creation_time(self): + # An unknown creation time cannot be shown to fall inside the window. + names, spared = self._find( + [self._cluster('cl-test-ageless')], + since=datetime.datetime.now(tz=datetime.timezone.utc), + ) + self.assertEqual(names, []) + self.assertIn('cl-test-ageless', spared[0]) + + def test_any_name_drops_the_name_gate(self): + names, _ = self._find( + [self._cluster('some-persons-cluster', hours=10)], + older_than=2, any_name=True, + ) + self.assertEqual(names, ['some-persons-cluster']) + # Nothing is unrecognized once every name counts. + self.assertEqual(self.unmatched, []) + + def test_kind_keeps_the_sweep_off_the_other_apis(self): + # This is the only guard left when --any-name and --since are both + # given, so a kind that was not asked for must not even be listed. + import singlestoredb as s2 + + clusters = MagicMock() + clusters.clusters = [self._cluster('anything', hours=10)] + clusters.starter_clusters = [] + workspaces = MagicMock() + workspaces.workspace_groups = [self._cluster('a group', hours=10)] + workspaces.starter_workspaces = [self._cluster('a starter', hours=10)] + + with patch.object( + s2, 'manage_clusters', return_value=clusters, + ) as clusters_call, patch.object( + s2, 'manage_workspaces', return_value=workspaces, + ): + found, _, _ = self.mod.find_leftovers( + older_than=2, any_name=True, kinds=['workspace-group'], + ) + + self.assertEqual([x[1].name for x in found], ['a group']) + clusters_call.assert_not_called() + + def test_since_reads_a_day_as_local_midnight(self): + for text, expected in ( + ('today', datetime.date.today()), + ( + 'yesterday', + datetime.date.today() - datetime.timedelta(days=1), + ), + ('2026-09-01', datetime.date(2026, 9, 1)), + ): + cutoff = self.mod.parse_since(text) + self.assertEqual(cutoff.date(), expected, text) + self.assertEqual(cutoff.hour, 0, text) + # Aware, or comparing it with a created_at raises. + self.assertIsNotNone(cutoff.tzinfo, text) + + def test_a_since_that_is_not_a_date_is_rejected(self): + import argparse + with self.assertRaises(argparse.ArgumentTypeError): + self.mod.parse_since('last tuesday') + + +if __name__ == '__main__': + unittest.main() diff --git a/singlestoredb/tests/test_management.py b/singlestoredb/tests/test_management_v1.py similarity index 68% rename from singlestoredb/tests/test_management.py rename to singlestoredb/tests/test_management_v1.py index f450e1f13..3328969a5 100755 --- a/singlestoredb/tests/test_management.py +++ b/singlestoredb/tests/test_management_v1.py @@ -1,12 +1,32 @@ #!/usr/bin/env python # type: ignore -"""SingleStoreDB Management API testing.""" +""" +SingleStoreDB v1 Management API testing. + +Everything here targets management API v1 -- workspaces, workspace groups and +the resources hanging off them. No test in this file may branch on version; +the v2 equivalents live in ``test_management_v2.py``, the version-neutral +helper units in ``test_management_utils.py``, and the structural cross-version +invariants in ``test_management_versioning.py``. + +The whole module carries ``@pytest.mark.management_v1`` (see ``pytestmark`` +below) so that the v1 endpoints can be switched off as a group now that +``management.version`` defaults to v2: ``-m 'not management_v1'`` for a normal +run, ``-m 'management_v1'`` for the nightly that still proves v1 works. The +marker is separate from ``management`` because this file also holds mocked +units that need no token -- those are v1-specific too, and go away with +``management/v1/``. +""" +import datetime import os import pathlib import random import re import secrets import unittest +from unittest.mock import MagicMock +from unittest.mock import patch +from unittest.mock import PropertyMock import pytest @@ -19,6 +39,9 @@ TEST_DIR = pathlib.Path(os.path.dirname(__file__)) +#: Applies to every test in this module, live or mocked. +pytestmark = pytest.mark.management_v1 + def clean_name(s): """Change all non-word characters to -.""" @@ -30,172 +53,6 @@ def shared_database_name(s): return re.sub(r'[^\w]', '', s).replace('-', '_').lower() -@pytest.mark.skip(reason='Legacy cluster Management API is going away') -@pytest.mark.management -class TestCluster(unittest.TestCase): - - manager = None - cluster = None - password = None - - @classmethod - def setUpClass(cls): - cls.manager = s2.manage_cluster() - - us_regions = [x for x in cls.manager.regions if 'US' in x.name] - cls.password = secrets.token_urlsafe(20) + '-x&$' - - cls.cluster = cls.manager.create_cluster( - clean_name('cm-test-{}'.format(secrets.token_urlsafe(20)[:20])), - region=random.choice(us_regions).id, - admin_password=cls.password, - firewall_ranges=['0.0.0.0/0'], - expires_at='1h', - size='S-00', - wait_on_active=True, - ) - - @classmethod - def tearDownClass(cls): - if cls.cluster is not None: - cls.cluster.terminate() - cls.cluster = None - cls.manager = None - cls.password = None - - def test_str(self): - assert self.cluster.name in str(self.cluster.name) - - def test_repr(self): - assert repr(self.cluster) == str(self.cluster) - - def test_region_str(self): - s = str(self.cluster.region) - assert 'Azure' in s or 'GCP' in s or 'AWS' in s, s - - def test_region_repr(self): - assert repr(self.cluster.region) == str(self.cluster.region) - - def test_regions(self): - out = self.manager.regions - providers = {x.provider for x in out} - names = [x.name for x in out] - assert 'Azure' in providers, providers - assert 'GCP' in providers, providers - assert 'AWS' in providers, providers - - objs = {} - ids = [] - for item in out: - ids.append(item.id) - objs[item.id] = item - if item.name not in objs: - objs[item.name] = item - - name = random.choice(names) - assert out[name] == objs[name] - id = random.choice(ids) - assert out[id] == objs[id] - - def test_clusters(self): - clusters = self.manager.clusters - ids = [x.id for x in clusters] - assert self.cluster.id in ids, ids - - def test_get_cluster(self): - clus = self.manager.get_cluster(self.cluster.id) - assert clus.id == self.cluster.id, clus.id - - with self.assertRaises(s2.ManagementError) as cm: - clus = self.manager.get_cluster('bad id') - - assert 'UUID' in cm.exception.msg, cm.exception.msg - - def test_update(self): - assert self.cluster.name.startswith('cm-test-') - - name = self.cluster.name.replace('cm-test-', 'cm-foo-') - self.cluster.update(name=name) - - clus = self.manager.get_cluster(self.cluster.id) - assert clus.name == name, clus.name - - def test_suspend_resume(self): - trues = ['1', 'on', 'true'] - do_test = os.environ.get('SINGLESTOREDB_TEST_SUSPEND', '0').lower() in trues - - if not do_test: - self.skipTest( - 'Suspend / resume tests skipped by default due to ' - 'being time consuming; set SINGLESTOREDB_TEST_SUSPEND=1 ' - 'to enable', - ) - - assert self.cluster.state != 'Suspended', self.cluster.state - - self.cluster.suspend(wait_on_suspended=True) - assert self.cluster.state == 'Suspended', self.cluster.state - - self.cluster.resume(wait_on_resumed=True) - assert self.cluster.state == 'Active', self.cluster.state - - def test_no_manager(self): - clus = self.manager.get_cluster(self.cluster.id) - clus._manager = None - - with self.assertRaises(s2.ManagementError) as cm: - clus.refresh() - - assert 'No cluster manager' in cm.exception.msg, cm.exception.msg - - with self.assertRaises(s2.ManagementError) as cm: - clus.update() - - assert 'No cluster manager' in cm.exception.msg, cm.exception.msg - - with self.assertRaises(s2.ManagementError) as cm: - clus.suspend() - - assert 'No cluster manager' in cm.exception.msg, cm.exception.msg - - with self.assertRaises(s2.ManagementError) as cm: - clus.resume() - - assert 'No cluster manager' in cm.exception.msg, cm.exception.msg - - with self.assertRaises(s2.ManagementError) as cm: - clus.terminate() - - assert 'No cluster manager' in cm.exception.msg, cm.exception.msg - - def test_connect(self): - trues = ['1', 'on', 'true'] - pure_python = os.environ.get('SINGLESTOREDB_PURE_PYTHON', '0').lower() in trues - - self.skipTest('Connection test is disable due to flakey server') - - if pure_python: - self.skipTest('Connections through managed service are disabled') - - try: - with self.cluster.connect(user='admin', password=self.password) as conn: - with conn.cursor() as cur: - cur.execute('show databases') - assert 'cluster' in [x[0] for x in list(cur)] - except s2.ManagementError as exc: - if 'endpoint has not been set' not in str(exc): - self.skipTest('No endpoint in response. Skipping connection test.') - - # Test missing endpoint - clus = self.manager.get_cluster(self.cluster.id) - clus.endpoint = None - - with self.assertRaises(s2.ManagementError) as cm: - clus.connect(user='admin', password=self.password) - - assert 'endpoint' in cm.exception.msg, cm.exception.msg - - @pytest.mark.management class TestWorkspace(unittest.TestCase): @@ -206,7 +63,9 @@ class TestWorkspace(unittest.TestCase): @classmethod def setUpClass(cls): - cls.manager = s2.manage_workspaces() + # Pinned: manage_workspaces() follows the management.version + # option, and this is the v1 suite. + cls.manager = s2.manage_workspaces(version='v1') us_regions = [x for x in cls.manager.regions if 'US' in x.name] cls.password = secrets.token_urlsafe(20) + '-x&$' @@ -286,8 +145,13 @@ def test_workspace_groups(self): objs = {} for item in workspace_groups: - objs[item.id] = item - objs[item.name] = item + # setdefault, and name before id, so this resolves a key the way + # NamedList._find_item does: to the *first* match. Plain assignment + # kept the last, which disagrees as soon as the listing carries two + # entries of one name -- terminated groups stay in the listing, so + # a suite that recreates a group under its old name produces that. + objs.setdefault(item.name, item) + objs.setdefault(item.id, item) name = random.choice(names) assert workspace_groups[name] == objs[name] @@ -306,8 +170,9 @@ def test_workspaces(self): objs = {} for item in spaces: - objs[item.id] = item - objs[item.name] = item + # First match wins, as in test_workspace_groups above. + objs.setdefault(item.name, item) + objs.setdefault(item.id, item) name = random.choice(names) assert spaces[name] == objs[name] @@ -379,16 +244,24 @@ class TestStarterWorkspace(unittest.TestCase): @classmethod def setUpClass(cls): - cls.manager = s2.manage_workspaces() + cls.manager = s2.manage_workspaces(version='v1') shared_tier_regions: NamedList[Region] = [ x for x in cls.manager.shared_tier_regions if 'US' in x.name ] - cls.starter_username = 'starter_user' - cls.password = secrets.token_urlsafe(20) - name = shared_database_name(secrets.token_urlsafe(20)[:20]) + # The starter-tier user name has to be unique across every starter + # deployment in the project, not just within this one: creating the + # same name in a second starter deployment fails while the first is + # live. So it is namespaced like the deployment and the database are, + # or this class collides with TestStarterCluster in test_management_v2 + # -- they run on different xdist workers -- and with any starter + # deployment an earlier failed run leaked. The API answers the + # collision with a bare 500, which names nothing. + cls.starter_username = f'starter_user_{name[:8]}' + cls.password = secrets.token_urlsafe(20) + cls.database_name = f'starter_db_{name}' shared_tier_region: Region = random.choice(shared_tier_regions) @@ -439,8 +312,9 @@ def test_starter_workspaces(self): objs = {} for item in workspaces: - objs[item.id] = item - objs[item.name] = item + # First match wins, as in test_workspace_groups above. + objs.setdefault(item.name, item) + objs.setdefault(item.id, item) name = random.choice(names) assert workspaces[name] == objs[name] @@ -489,7 +363,7 @@ class TestStage(unittest.TestCase): @classmethod def setUpClass(cls): - cls.manager = s2.manage_workspaces() + cls.manager = s2.manage_workspaces(version='v1') us_regions = [x for x in cls.manager.regions if 'US' in x.name] cls.password = secrets.token_urlsafe(20) + '-x&$' @@ -1037,32 +911,19 @@ def test_file_object(self): class TestSecrets(unittest.TestCase): manager = None - wg = None - password = None @classmethod def setUpClass(cls): - 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.wg = cls.manager.create_workspace_group( - f'wg-test-{name}', - region=random.choice(us_regions).id, - admin_password=cls.password, - firewall_ranges=['0.0.0.0/0'], - ) + # No deployment: a secret belongs to the organization, not to a + # workspace group, and test_get_secret reaches it through + # organizations.current. This used to create a group with a firewall + # and an admin password that nothing in the class ever read -- a + # provisioning wait and a teardown for an unused fixture. + cls.manager = s2.manage_workspaces(version='v1') @classmethod def tearDownClass(cls): - if cls.wg is not None: - cls.wg.terminate(force=True) - cls.wg = None cls.manager = None - cls.password = None def test_get_secret(self): # manually create secret and then get secret @@ -1101,7 +962,7 @@ class TestJob(unittest.TestCase): @classmethod def setUpClass(cls): - cls.manager = s2.manage_workspaces() + cls.manager = s2.manage_workspaces(version='v1') us_regions = [x for x in cls.manager.regions if 'US' in x.name] cls.password = secrets.token_urlsafe(20) + '-x&$' @@ -1252,7 +1113,9 @@ class TestFileSpaces(unittest.TestCase): @classmethod def setUpClass(cls): - cls.manager = s2.manage_files() + # Pinned: manage_files() follows the management.version option, and + # this is the v1 suite. + cls.manager = s2.manage_files(version='v1') cls.personal_space = cls.manager.personal_space cls.shared_space = cls.manager.shared_space @@ -1588,7 +1451,9 @@ class TestRegions(unittest.TestCase): @classmethod def setUpClass(cls): """Set up the test environment.""" - cls.manager = s2.manage_regions() + # Pinned: manage_regions() follows the management.version option, and + # this is the v1 suite. + cls.manager = s2.manage_regions(version='v1') @classmethod def tearDownClass(cls): @@ -1651,3 +1516,516 @@ def test_str_repr(self): # Test __repr__ assert repr(region) == str(region) + + +# +# v1 behavior units. These need neither a management token nor a +# container -- they drive the v1 entity classes against fake API +# payloads. Anything version-neutral belongs in +# test_management_utils.py instead. +# + +FAKE_TOKEN = 'test-token-12345' +FAKE_BASE_URL = 'https://api.example.com' +FAKE_ORG_ID = 'org-12345' + + +def _make_workspace_manager(version='v1', organization_id=FAKE_ORG_ID): + """Construct a v1 WorkspaceManager with patched token resolver.""" + from singlestoredb.management.v1.workspace import WorkspaceManager + with patch( + 'singlestoredb.management.manager.get_token', + return_value=FAKE_TOKEN, + ): + return WorkspaceManager( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + version=version, + organization_id=organization_id, + ) + + +def _make_workspace_group(manager=None, group_id='wsg-456', extra_obj=None): + """Build a v1 WorkspaceGroup from a fake API response. + + ``WorkspaceGroup.from_dict`` calls ``manager.regions`` to resolve the + region; we stub it so no network call is made. + """ + from singlestoredb.management.v1.workspace import WorkspaceGroup + from singlestoredb.management.v1.workspace import WorkspaceManager + mgr = manager or _make_workspace_manager() + obj = { + 'name': 'test-group', + 'workspaceGroupID': group_id, + 'createdAt': '2024-01-01T00:00:00Z', + 'regionID': 'region-789', + 'firewallRanges': ['0.0.0.0/0'], + } + if extra_obj: + obj.update(extra_obj) + with patch.object( + WorkspaceManager, 'regions', + new_callable=PropertyMock, return_value=[], + ): + wg = WorkspaceGroup.from_dict(obj, mgr) + return wg, mgr, obj + + +class TestTokenStorageFix(unittest.TestCase): + """Test that Manager authenticates with the resolved token.""" + + @patch('singlestoredb.management.manager.is_jwt', return_value=False) + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_none_token_resolves(self, _mock_token, _mock_jwt): + """When access_token=None, the resolved token is used.""" + from singlestoredb.management.v1.workspace import WorkspaceManager + mgr = WorkspaceManager( + access_token=None, + base_url=FAKE_BASE_URL, + version='v1', + ) + self.assertEqual( + mgr._sess.headers['Authorization'], f'Bearer {FAKE_TOKEN}', + ) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_explicit_token_used_as_is(self, _mock_token): + """When access_token is provided, it's used directly.""" + from singlestoredb.management.v1.workspace import WorkspaceManager + mgr = WorkspaceManager( + access_token='my-explicit-token', + base_url=FAKE_BASE_URL, + version='v1', + ) + self.assertEqual( + mgr._sess.headers['Authorization'], 'Bearer my-explicit-token', + ) + + +class TestWorkspaceFromDictNewFields(unittest.TestCase): + """ + Coverage for the staged additions in ``v1/workspace.py``: + ``auto_scale``, ``kai_enabled``, ``scale_factor``, plus the widened + ``cache_config`` (now float). + """ + + def _base_obj(self): + return { + 'name': 'test-ws', + 'workspaceID': 'ws-1', + 'workspaceGroupID': 'wsg-1', + 'size': 'S-00', + 'state': 'Active', + 'createdAt': '2024-01-01T00:00:00Z', + } + + def test_new_fields_present(self): + from singlestoredb.management.v1.workspace import Workspace + mgr = _make_workspace_manager() + obj = self._base_obj() + obj.update({ + 'autoScale': { + 'sensitivity': 'HIGH', + 'maxScaleFactor': 4.0, + 'changedAt': '2024-01-01T00:00:00Z', + 'lastAutoScaledAt': '2024-01-02T00:00:00Z', + }, + 'kaiEnabled': True, + 'scaleFactor': 2.5, + 'cacheConfig': 1.5, + }) + ws = Workspace.from_dict(obj, mgr) + # auto_scale keys are camel_to_snake_dict-converted + self.assertEqual(ws.auto_scale['sensitivity'], 'HIGH') + self.assertEqual(ws.auto_scale['max_scale_factor'], 4.0) + self.assertEqual(ws.auto_scale['changed_at'], '2024-01-01T00:00:00Z') + self.assertEqual( + ws.auto_scale['last_auto_scaled_at'], '2024-01-02T00:00:00Z', + ) + self.assertNotIn('maxScaleFactor', ws.auto_scale) + self.assertIs(ws.kai_enabled, True) + self.assertEqual(ws.scale_factor, 2.5) + self.assertEqual(ws.cache_config, 1.5) + + def test_new_fields_default_to_none(self): + from singlestoredb.management.v1.workspace import Workspace + mgr = _make_workspace_manager() + ws = Workspace.from_dict(self._base_obj(), mgr) + self.assertIsNone(ws.auto_scale) + self.assertIsNone(ws.kai_enabled) + self.assertIsNone(ws.scale_factor) + + +class TestWorkspaceUpdatePosting(unittest.TestCase): + """``Workspace.update`` must include the new fields in the PATCH body.""" + + def _make_workspace(self, mgr): + from singlestoredb.management.v1.workspace import Workspace + obj = { + 'name': 'test-ws', + 'workspaceID': 'ws-1', + 'workspaceGroupID': 'wsg-1', + 'size': 'S-00', + 'state': 'Active', + 'createdAt': '2024-01-01T00:00:00Z', + } + return Workspace.from_dict(obj, mgr) + + def test_update_posts_new_fields_only_when_set(self): + mgr = _make_workspace_manager() + mgr._patch = MagicMock() + ws = self._make_workspace(mgr) + ws.refresh = MagicMock() + + ws.update( + auto_scale={'sensitivity': 'HIGH'}, + enable_kai=True, + scale_factor=2.0, + cache_config=1.5, + ) + + mgr._patch.assert_called_once() + args, kwargs = mgr._patch.call_args + self.assertEqual(args[0], 'workspaces/ws-1') + body = kwargs['json'] + self.assertEqual(body['autoScale'], {'sensitivity': 'HIGH'}) + self.assertIs(body['enableKai'], True) + self.assertEqual(body['scaleFactor'], 2.0) + self.assertEqual(body['cacheConfig'], 1.5) + + def test_update_omits_keys_when_param_none(self): + mgr = _make_workspace_manager() + mgr._patch = MagicMock() + ws = self._make_workspace(mgr) + ws.refresh = MagicMock() + + ws.update(size='S-1') + + body = mgr._patch.call_args.kwargs['json'] + self.assertEqual(body, {'size': 'S-1'}) + self.assertNotIn('autoScale', body) + self.assertNotIn('enableKai', body) + self.assertNotIn('scaleFactor', body) + + +class TestWorkspaceGroupNewFields(unittest.TestCase): + """Coverage for the new staged fields on ``WorkspaceGroup.from_dict``.""" + + def _obj_with_new_fields(self): + return { + 'name': 'test-group', + 'workspaceGroupID': 'wsg-1', + 'createdAt': '2024-01-01T00:00:00Z', + 'regionID': 'region-789', + 'firewallRanges': ['0.0.0.0/0'], + 'allowAllTraffic': True, + 'deploymentType': 'PRODUCTION', + 'expiresAt': '2025-06-30T23:59:59Z', + 'highAvailabilityTwoZones': True, + 'optInPreviewFeature': False, + 'outboundAllowList': '203.0.113.0/24', + 'projectID': 'proj-1', + 'projectName': 'my-project', + 'smartDRStatus': 'ACTIVE', + 'state': 'ACTIVE', + 'updateWindow': {'day': 0, 'hour': 4}, + 'provider': 'aws', + 'regionName': 'us-east-1', + } + + def test_all_new_fields_mapped(self): + from singlestoredb.management.v1.workspace import WorkspaceGroup + mgr = _make_workspace_manager() + with patch.object( + type(mgr), 'regions', + new_callable=PropertyMock, return_value=[], + ): + wg = WorkspaceGroup.from_dict(self._obj_with_new_fields(), mgr) + self.assertEqual(wg.deployment_type, 'PRODUCTION') + self.assertIsInstance(wg.expires_at, datetime.datetime) + self.assertIs(wg.high_availability_two_zones, True) + self.assertIs(wg.opt_in_preview_feature, False) + self.assertEqual(wg.outbound_allow_list, '203.0.113.0/24') + self.assertEqual(wg.project_id, 'proj-1') + self.assertEqual(wg.project_name, 'my-project') + self.assertEqual(wg.smart_dr_status, 'ACTIVE') + self.assertEqual(wg.state, 'ACTIVE') + # update_window stays a raw dict (not snake-cased) + self.assertEqual(wg.update_window, {'day': 0, 'hour': 4}) + self.assertEqual(wg.provider, 'aws') + self.assertEqual(wg.region_name, 'us-east-1') + + def test_new_fields_default_to_none(self): + wg, _, _ = _make_workspace_group() + self.assertIsNone(wg.deployment_type) + self.assertIsNone(wg.expires_at) + self.assertIsNone(wg.high_availability_two_zones) + self.assertIsNone(wg.opt_in_preview_feature) + self.assertIsNone(wg.outbound_allow_list) + self.assertIsNone(wg.project_id) + self.assertIsNone(wg.project_name) + self.assertIsNone(wg.smart_dr_status) + self.assertIsNone(wg.state) + self.assertIsNone(wg.update_window) + self.assertIsNone(wg.provider) + self.assertIsNone(wg.region_name) + + +class TestWorkspaceGroupCreateUpdatePosting(unittest.TestCase): + """Body coverage for create_workspace_group / WorkspaceGroup.update.""" + + def test_create_workspace_group_posts_new_fields(self): + mgr = _make_workspace_manager() + # Make get_workspace_group a no-op; we only inspect the POST body. + post_response = MagicMock() + post_response.json.return_value = {'workspaceGroupID': 'wsg-new'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_workspace_group = MagicMock(return_value='sentinel') + + result = mgr.create_workspace_group( + name='wg-1', + region='region-789', + firewall_ranges=['0.0.0.0/0'], + provider='aws', + region_name='us-east-1', + deployment_type='PRODUCTION', + high_availability_two_zones=True, + opt_in_preview_feature=False, + project_id='proj-1', + ) + + self.assertEqual(result, 'sentinel') + body = mgr._post.call_args.kwargs['json'] + self.assertEqual(body['provider'], 'aws') + self.assertEqual(body['regionName'], 'us-east-1') + self.assertEqual(body['deploymentType'], 'PRODUCTION') + self.assertIs(body['highAvailabilityTwoZones'], True) + self.assertIs(body['optInPreviewFeature'], False) + self.assertEqual(body['projectID'], 'proj-1') + + def test_workspace_group_update_includes_deployment_type(self): + wg, mgr, _ = _make_workspace_group() + mgr._patch = MagicMock() + wg.refresh = MagicMock() + + wg.update(deployment_type='NON-PRODUCTION', name='renamed') + + body = mgr._patch.call_args.kwargs['json'] + self.assertEqual(body['deploymentType'], 'NON-PRODUCTION') + self.assertEqual(body['name'], 'renamed') + + def test_workspace_group_update_omits_unset_fields(self): + wg, mgr, _ = _make_workspace_group() + mgr._patch = MagicMock() + wg.refresh = MagicMock() + + wg.update(name='renamed') + + body = mgr._patch.call_args.kwargs['json'] + self.assertNotIn('deploymentType', body) + + +class TestJobsManagerScheduleDuration(unittest.TestCase): + """ + Coverage for the staged ``max_allowed_execution_duration_in_minutes`` + parameter on ``JobsManager.schedule``. + """ + + def _patch_post(self, mgr, response_obj): + post_response = MagicMock() + post_response.json.return_value = response_obj + mgr._post = MagicMock(return_value=post_response) + return post_response + + def _fake_job_response(self): + return { + 'jobID': 'job-1', + 'name': 'j', + 'description': None, + 'enqueuedBy': 'me', + 'createdAt': '2024-01-01T00:00:00Z', + 'completedExecutionsCount': 0, + 'jobMetadata': [], + 'terminatedAt': None, + 'executionConfig': { + 'createSnapshot': True, + 'notebookPath': '/x.ipynb', + }, + 'schedule': {'mode': 'Once'}, + 'targetConfig': None, + } + + def test_duration_present_when_set(self): + from singlestoredb.management.v1.job import JobsManager + from singlestoredb.management.v1.job import Mode + + ws_mgr = _make_workspace_manager() + jobs = JobsManager(ws_mgr) + self._patch_post(ws_mgr, self._fake_job_response()) + + with patch( + 'singlestoredb.management.v1.job.Job.from_dict', + return_value='sentinel', + ): + jobs.schedule( + notebook_path='/x.ipynb', + mode=Mode.ONCE, + create_snapshot=True, + max_allowed_execution_duration_in_minutes=42, + ) + + body = ws_mgr._post.call_args.kwargs['json'] + self.assertEqual( + body['executionConfig']['maxAllowedExecutionDurationInMinutes'], + 42, + ) + + def test_duration_absent_when_unset(self): + from singlestoredb.management.v1.job import JobsManager + from singlestoredb.management.v1.job import Mode + + ws_mgr = _make_workspace_manager() + jobs = JobsManager(ws_mgr) + self._patch_post(ws_mgr, self._fake_job_response()) + + with patch( + 'singlestoredb.management.v1.job.Job.from_dict', + return_value='sentinel', + ): + jobs.schedule( + notebook_path='/x.ipynb', + mode=Mode.ONCE, + create_snapshot=True, + ) + + body = ws_mgr._post.call_args.kwargs['json'] + self.assertNotIn( + 'maxAllowedExecutionDurationInMinutes', + body['executionConfig'], + ) + + +class TestWorkspaceGroupRegionResolution(unittest.TestCase): + """ + ``WorkspaceGroup.from_dict`` resolves its region through a fallback + ladder: match on ``regionID`` first, then on ``(region_name, provider)`` + for regions that carry no ID, then the payload's own fields, then + ````. + """ + + def _region_without_id(self, name, provider, region_name): + from singlestoredb.management.v1.region import Region + return Region( + name=name, provider=provider, id=None, region_name=region_name, + ) + + def _wg_payload(self, **overrides): + obj = { + 'name': 'test-group', + 'workspaceGroupID': 'wsg-1', + 'createdAt': '2024-01-01T00:00:00Z', + 'regionID': 'region-uuid-1', + 'regionName': 'us-west1', + 'provider': 'GCP', + } + obj.update(overrides) + return obj + + def test_resolves_by_region_name_and_provider_when_no_id(self): + from singlestoredb.management.v1.workspace import ( + WorkspaceGroup, WorkspaceManager, + ) + mgr = MagicMock(spec=WorkspaceManager) + mgr.regions = [ + self._region_without_id('us-west1', 'GCP', 'us-west1'), + self._region_without_id('eu-central-1', 'AWS', 'eu-central-1'), + ] + wg = WorkspaceGroup.from_dict(self._wg_payload(), mgr) + self.assertEqual(wg.region.name, 'us-west1') + self.assertEqual(wg.region.provider, 'GCP') + self.assertEqual(wg.region.region_name, 'us-west1') + + def test_match_by_id_wins(self): + from singlestoredb.management.v1.region import Region + from singlestoredb.management.v1.workspace import ( + WorkspaceGroup, WorkspaceManager, + ) + mgr = MagicMock(spec=WorkspaceManager) + mgr.regions = [ + Region( + name='us-west1', provider='GCP', + id='region-uuid-1', region_name='us-west1', + ), + ] + wg = WorkspaceGroup.from_dict(self._wg_payload(), mgr) + self.assertEqual(wg.region.id, 'region-uuid-1') + self.assertEqual(wg.region.name, 'us-west1') + + def test_no_match_falls_back_to_payload_fields(self): + from singlestoredb.management.v1.workspace import ( + WorkspaceGroup, WorkspaceManager, + ) + mgr = MagicMock(spec=WorkspaceManager) + mgr.regions = [] + wg = WorkspaceGroup.from_dict(self._wg_payload(), mgr) + self.assertEqual(wg.region.name, 'us-west1') + self.assertEqual(wg.region.provider, 'GCP') + self.assertEqual(wg.region.id, 'region-uuid-1') + self.assertEqual(wg.region.region_name, 'us-west1') + + def test_no_match_no_payload_fields_uses_unknown(self): + from singlestoredb.management.v1.workspace import ( + WorkspaceGroup, WorkspaceManager, + ) + mgr = MagicMock(spec=WorkspaceManager) + mgr.regions = [] + obj = { + 'name': 'test-group', + 'workspaceGroupID': 'wsg-1', + 'createdAt': '2024-01-01T00:00:00Z', + } + wg = WorkspaceGroup.from_dict(obj, mgr) + self.assertEqual(wg.region.name, '') + self.assertEqual(wg.region.provider, '') + self.assertIsNone(wg.region.id) + + +class TestDateTimeParsingFixes(unittest.TestCase): + """ + Regression test for commit 85faf724: ISO8601-Z timestamp parsing + on entities that go through ``to_datetime``. + """ + + def test_workspace_created_at_parsed(self): + from singlestoredb.management.v1.workspace import Workspace + mgr = _make_workspace_manager() + obj = { + 'name': 'test-ws', + 'workspaceID': 'ws-1', + 'workspaceGroupID': 'wsg-1', + 'size': 'S-00', + 'state': 'Active', + 'createdAt': '2024-03-15T12:30:45Z', + 'lastResumedAt': '2024-03-16T08:00:00.123Z', + } + ws = Workspace.from_dict(obj, mgr) + self.assertIsInstance(ws.created_at, datetime.datetime) + self.assertEqual(ws.created_at.year, 2024) + self.assertEqual(ws.created_at.month, 3) + self.assertEqual(ws.created_at.day, 15) + self.assertEqual(ws.created_at.hour, 12) + self.assertIsInstance(ws.last_resumed_at, datetime.datetime) + + def test_workspace_group_expires_at_parsed(self): + wg, _, _ = _make_workspace_group( + extra_obj={'expiresAt': '2025-06-30T23:59:59Z'}, + ) + self.assertIsInstance(wg.expires_at, datetime.datetime) + self.assertEqual(wg.expires_at.year, 2025) + + def test_workspace_group_terminated_at_zero_returns_none(self): + """The sentinel 0001-01-01 timestamp must round-trip to None.""" + wg, _, _ = _make_workspace_group( + extra_obj={'terminatedAt': '0001-01-01T00:00:00Z'}, + ) + self.assertIsNone(wg.terminated_at) diff --git a/singlestoredb/tests/test_management_v2.py b/singlestoredb/tests/test_management_v2.py new file mode 100644 index 000000000..5842bdc56 --- /dev/null +++ b/singlestoredb/tests/test_management_v2.py @@ -0,0 +1,1945 @@ +#!/usr/bin/env python +# type: ignore +""" +SingleStoreDB v2 Management API testing. + +Everything here targets management API v2 -- the flat ``Cluster`` resource and +the starter clusters, stages, secrets, jobs and regions hanging off it. No test +in this file may branch on version; the v1 equivalents live in +``test_management_v1.py``, the version-neutral helper units in +``test_management_utils.py``, and the structural cross-version invariants in +``test_management_versioning.py``. + +.. warning:: The ``@pytest.mark.management`` suites below have not been run + against a live v2 organization. They were written by translating the v1 + suites resource by resource, so every assertion that rests on a v2 response + or request *shape* rather than on SDK-internal behavior is marked with an + ``UNVERIFIED`` comment. Treat a failure in one of those as "check the API", + not automatically as "fix the test". +""" +import os +import random +import re +import secrets +import tempfile +import unittest +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest + +import singlestoredb as s2 +from singlestoredb.exceptions import ManagementError +from singlestoredb.management.job import Status +from singlestoredb.management.job import TargetType +from singlestoredb.management.project import Project +from singlestoredb.management.region import Region +from singlestoredb.management.utils import NamedList +from singlestoredb.tests import utils + + +TEST_DIR = os.path.dirname(__file__) + +FAKE_TOKEN = 'test-token-12345' +FAKE_BASE_URL = 'https://api.example.com' + +# Fake project IDs. These have to be UUID-shaped: a project can be named by +# either its name or its ID, and the wrapper tells the two apart by shape, so a +# stand-in such as 'pr-1' would be read as a name and send the manager off to +# list the organization's projects. +FAKE_PROJECT_ID = '11111111-1111-4111-8111-111111111111' +FAKE_SHARED_PROJECT_ID = '22222222-2222-4222-8222-222222222222' +FAKE_STANDARD_PROJECT_ID = '33333333-3333-4333-8333-333333333333' + +#: Fake cluster ID, for the tests that only pass one along. +FAKE_CLUSTER_ID = '44444444-4444-4444-8444-444444444444' + + +def clean_name(s): + """ + Return ``s`` as a valid v2 cluster name. + + Verified against the live API: a cluster name has to match + ``[a-z0-9]([a-z0-9-]*[a-z0-9])?`` and be 1-32 characters. Lowercase letters, + digits and hyphens only -- an uppercase letter, an underscore, a dot, a + space, or a leading or trailing hyphen all draw + ``400 name: must be in a valid format``. Repeated hyphens are fine. + """ + out = re.sub(r'[^\w]', r'-', s).replace('_', '-').lower().strip('-') + return out or 'x' + + +def shared_database_name(s): + """Return a shared database name. Cannot contain special characters except -""" + return re.sub(r'[^\w]', '', s).replace('-', '_').lower() + + +def _us_regions(manager): + """Return the US regions a v2 manager reports, or skip the test.""" + out = [x for x in manager.regions if 'US' in x.name or 'us-' in x.name] + if not out: + raise unittest.SkipTest('No US regions reported by the v2 API') + return out + + +def _project_id(manager): + """ + Return the project ID the live v2 suites deploy into, or skip the test. + + ``POST /v2/clusters`` requires ``projectID``, so a project has to be chosen + before anything can be created. ``SINGLESTOREDB_TEST_PROJECT`` wins if it is + set; otherwise the STANDARD-edition project is used, which is where every + workspace group the v1 suites create already lands. + + Not ``SINGLESTOREDB_PROJECT``: that names an inference API project, not one + of these, and pointing the suites at it deploys nothing. + """ + from_env = os.environ.get('SINGLESTOREDB_TEST_PROJECT') + if from_env: + return from_env + + standard = [x for x in manager.projects if x.edition == 'STANDARD'] + if not standard: + raise unittest.SkipTest( + 'No STANDARD project in this organization; set ' + 'SINGLESTOREDB_TEST_PROJECT to the project to deploy into', + ) + return standard[0].id + + +# +# Unit tests. These need no token and no deployment. +# + +class TestV2RegionBehavior(unittest.TestCase): + """ + ``RegionManager`` at v2: ``list_regions`` hits ``/v2/regions`` and + ``list_shared_tier_regions`` hits ``/v2/regions/sharedtier``, which + answers with the same shape. + """ + + def _make_region_manager(self): + from singlestoredb.management.v2.region import RegionManager + with patch( + 'singlestoredb.management.manager.get_token', + return_value=FAKE_TOKEN, + ): + return RegionManager( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + version='v2', + ) + + def test_list_regions_uses_regions_endpoint(self): + mgr = self._make_region_manager() + get_response = MagicMock() + # Live shape (2026-08-24): ``region`` is the display name and + # ``regionName`` the provider slug -- not the other way round. + get_response.json.return_value = [ + { + 'provider': 'AWS', + 'region': 'US East 1 (N. Virginia)', + 'regionName': 'us-east-1', + }, + { + 'provider': 'GCP', + 'region': 'US West 2 (Oregon)', + 'regionName': 'us-west2', + }, + ] + mgr._get = MagicMock(return_value=get_response) + + regions = mgr.list_regions() + mgr._get.assert_called_once_with('regions') + self.assertEqual(len(regions), 2) + # v2 region entries have id=None -- there is no regionID in the + # response, so a region is identified by (provider, region_name). + for r in regions: + self.assertIsNone(r.id) + self.assertEqual(regions[0].name, 'US East 1 (N. Virginia)') + self.assertEqual(regions[0].region_name, 'us-east-1') + + def test_list_shared_tier_regions_uses_sharedtier_endpoint(self): + mgr = self._make_region_manager() + get_response = MagicMock() + # Live shape (2026-08-24): ``GET /v2/regions/sharedtier`` returns 200 + # with exactly the same keys as ``GET /v2/regions``. + get_response.json.return_value = [ + { + 'provider': 'AWS', + 'region': 'US East 1 (N. Virginia)', + 'regionName': 'us-east-1', + }, + ] + mgr._get = MagicMock(return_value=get_response) + + regions = mgr.list_shared_tier_regions() + mgr._get.assert_called_once_with('regions/sharedtier') + self.assertEqual(len(regions), 1) + self.assertEqual(regions[0].name, 'US East 1 (N. Virginia)') + self.assertEqual(regions[0].region_name, 'us-east-1') + self.assertIsNone(regions[0].id) + + +class TestClusterManagerPosting(unittest.TestCase): + """ + Request bodies the ``ClusterManager`` sends. + + .. warning:: UNVERIFIED. Every field name asserted here comes from the + wrapper, not from a recorded v2 response, so these tests pin the + wrapper's current behavior rather than confirming the API accepts it. + ``create_cluster``'s POST body in particular -- the nested ``size`` + object and the ``provider``/``region`` pair replacing v1's + ``regionID`` -- needs checking against a live v2 organization. + """ + + def _make_cluster_manager(self): + from singlestoredb.management.v2.cluster import ClusterManager + with patch( + 'singlestoredb.management.manager.get_token', + return_value=FAKE_TOKEN, + ): + return ClusterManager( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + version='v2', + ) + + def test_create_cluster_body(self): + mgr = self._make_cluster_manager() + post_response = MagicMock() + post_response.json.return_value = {'clusterID': 'cl-1'} + mgr._post = MagicMock(return_value=post_response) + sentinel = MagicMock() + mgr.get_cluster = MagicMock(return_value=sentinel) + + out = mgr.create_cluster( + 'my-cluster', + provider='AWS', + region='us-east-1', + size='S-00', + scale_factor=1.0, + firewall_ranges=['0.0.0.0/0'], + admin_password='hunter2', + update_window={'day': 3, 'hour': 4}, + project=FAKE_PROJECT_ID, + ) + + self.assertIs(out, sentinel) + mgr.get_cluster.assert_called_once_with('cl-1') + path, kwargs = mgr._post.call_args[0][0], mgr._post.call_args[1] + self.assertEqual(path, 'clusters') + body = kwargs['json'] + self.assertEqual(body['name'], 'my-cluster') + self.assertEqual(body['provider'], 'AWS') + # v2 names the region by its provider region name; there is no + # regionID to send. + self.assertEqual(body['region'], 'us-east-1') + self.assertNotIn('regionID', body) + # Size and scale factor are nested in one object, under ``sizeConfig``. + # That rename shipped on 2026-08-26, was backed out the next morning, + # and landed again by 2026-08-28, when ``POST /v2/clusters`` began + # answering 400 "unknown field" to ``size``. + self.assertEqual(body['sizeConfig'], {'size': 'S-00', 'scaleFactor': 1.0}) + self.assertNotIn('size', body) + self.assertEqual(body['firewallRanges'], ['0.0.0.0/0']) + self.assertEqual(body['adminPassword'], 'hunter2') + self.assertEqual(body['updateWindow'], {'day': 3, 'hour': 4}) + # The API rejects a create without projectID. + self.assertEqual(body['projectID'], FAKE_PROJECT_ID) + # Unset options are dropped rather than sent as null. + self.assertNotIn('kai', body) + self.assertNotIn('autoSuspend', body) + + def test_create_cluster_always_sends_firewall_ranges(self): + """ + ``firewallRanges`` cannot be dropped when unset. + + Verified live: ``POST /v2/clusters`` answers 400 "firewallRanges cannot + be null (indicate empty list [] to disallow all inbound traffic)", so + an omitted ``firewall_ranges`` has to be sent as a deny-all ``[]`` + rather than left out with the other unset fields. + """ + mgr = self._make_cluster_manager() + post_response = MagicMock() + post_response.json.return_value = {'clusterID': 'cl-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_cluster = MagicMock(return_value=MagicMock()) + + mgr.create_cluster( + 'my-cluster', + provider='AWS', + region='us-east-1', + size='S-00', + project=FAKE_PROJECT_ID, + ) + + body = mgr._post.call_args[1]['json'] + self.assertEqual(body['firewallRanges'], []) + + def test_create_cluster_accepts_a_region_object(self): + mgr = self._make_cluster_manager() + post_response = MagicMock() + post_response.json.return_value = {'clusterID': 'cl-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_cluster = MagicMock() + + mgr.create_cluster( + 'my-cluster', + region=Region( + name='us-east-1', provider='AWS', + id=None, region_name='us-east-1', + ), + project=FAKE_PROJECT_ID, + ) + body = mgr._post.call_args[1]['json'] + self.assertEqual(body['provider'], 'AWS') + self.assertEqual(body['region'], 'us-east-1') + + def test_create_starter_cluster_body(self): + mgr = self._make_cluster_manager() + post_response = MagicMock() + # UNVERIFIED: the starter-cluster create response is expected to name + # the new deployment ``virtualClusterID``. + post_response.json.return_value = {'virtualClusterID': 'vc-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_starter_cluster = MagicMock(return_value='sentinel') + + out = mgr.create_starter_cluster( + 'my-starter', database_name='db1', + provider='AWS', region='us-east-1', + ) + self.assertEqual(out, 'sentinel') + mgr.get_starter_cluster.assert_called_once_with('vc-1') + self.assertEqual( + mgr._post.call_args[1]['json'], { + 'name': 'my-starter', + 'databaseName': 'db1', + 'provider': 'AWS', + 'regionName': 'us-east-1', + }, + ) + + def test_create_cluster_returns_the_generated_admin_password(self): + """ + The generated password is carried off the create response. + + Verified live: ``POST /v2/clusters`` generates the admin password no + matter what ``adminPassword`` is sent, returns it in the create + response, and reports it nowhere else -- ``GET /v2/clusters/{id}`` has + no such field. Losing it means losing access to the cluster. + """ + from singlestoredb.management.v2.cluster import Cluster + mgr = self._make_cluster_manager() + post_response = MagicMock() + post_response.json.return_value = { + 'clusterID': 'cl-1', 'adminPassword': 'generated-not-hunter2', + } + mgr._post = MagicMock(return_value=post_response) + cluster = Cluster(name='my-cluster', id='cl-1', state='PENDING') + mgr.get_cluster = MagicMock(return_value=cluster) + + out = mgr.create_cluster( + 'my-cluster', provider='AWS', region='us-east-1', + admin_password='hunter2', project=FAKE_PROJECT_ID, + ) + self.assertEqual(out.admin_password, 'generated-not-hunter2') + # A cluster that did not come from a create has no password to report. + self.assertIsNone(Cluster(name='x', id='cl-2', state='ACTIVE').admin_password) + # And it must not leak into the string representations. + self.assertNotIn('generated-not-hunter2', str(out)) + self.assertNotIn('generated-not-hunter2', repr(out)) + + def test_create_starter_cluster_upper_cases_the_provider(self): + """ + The shared-tier route accepts only AWS | AZURE | GCP verbatim. + + Verified live: 'Azure' -- the spelling ``GET /v2/regions`` itself + reports -- fails with ``500 Unspecified is not a valid + CloudServiceProvider``, so a region's ``provider`` cannot be passed + through as-is. ``POST /v2/clusters`` has no such restriction. + """ + mgr = self._make_cluster_manager() + post_response = MagicMock() + post_response.json.return_value = {'virtualClusterID': 'vc-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_starter_cluster = MagicMock() + + mgr.create_starter_cluster( + 'my-starter', database_name='db1', + provider='Azure', region='southcentralus', + ) + self.assertEqual(mgr._post.call_args[1]['json']['provider'], 'AZURE') + + def test_create_starter_cluster_without_an_id_raises(self): + mgr = self._make_cluster_manager() + post_response = MagicMock() + post_response.json.return_value = {} + mgr._post = MagicMock(return_value=post_response) + with self.assertRaises(ManagementError): + mgr.create_starter_cluster( + 'my-starter', database_name='db1', + provider='AWS', region='us-east-1', + ) + + def test_shared_tier_regions_uses_sharedtier_endpoint(self): + mgr = self._make_cluster_manager() + get_response = MagicMock() + get_response.json.return_value = [ + { + 'provider': 'AWS', + 'region': 'US East 1 (N. Virginia)', + 'regionName': 'us-east-1', + }, + ] + mgr._get = MagicMock(return_value=get_response) + + regions = mgr.shared_tier_regions + mgr._get.assert_called_once_with('regions/sharedtier') + self.assertEqual(len(regions), 1) + self.assertEqual(regions[0].name, 'US East 1 (N. Virginia)') + self.assertEqual(regions[0].region_name, 'us-east-1') + + +class TestClusterFirewallWaiting(unittest.TestCase): + """ + Waiting for the asynchronously-applied firewall. + + Verified live: ``POST /v2/clusters`` and ``PATCH /v2/clusters/{id}`` apply + ``firewallRanges`` outside the state machine. The cluster reaches ACTIVE + with a resolvable endpoint while ``GET /v2/clusters/{id}`` still reports + ``firewallRanges: []`` and ``allowAllTraffic: null``, which denies all + inbound traffic, so a connect attempt in that window times out at the TCP + level. + + Also verified live: a requested ``firewallRanges: ['0.0.0.0/0']`` is stored + as ``allowAllTraffic: True`` with ``firewallRanges: []`` -- and that + cluster does accept connections (port 3306 open) -- so "reachable" is + either one, not non-empty ranges. + """ + + def _make_cluster_manager(self): + from singlestoredb.management.v2.cluster import ClusterManager + with patch( + 'singlestoredb.management.manager.get_token', + return_value=FAKE_TOKEN, + ): + return ClusterManager( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + version='v2', + ) + + def _cluster( + self, firewall_ranges=None, state='ACTIVE', manager=None, + allow_all_traffic=None, + ): + from singlestoredb.management.v2.cluster import Cluster + out = Cluster( + name='my-cluster', id='cl-1', state=state, + endpoint='svc.singlestore.com', + firewall_ranges=firewall_ranges, + allow_all_traffic=allow_all_traffic, + ) + out._manager = manager + return out + + def test_wait_on_firewall_polls_until_non_empty(self): + mgr = self._make_cluster_manager() + pending = self._cluster(firewall_ranges=[]) + applied = self._cluster(firewall_ranges=['0.0.0.0/0']) + mgr.get_cluster = MagicMock( + side_effect=[self._cluster(firewall_ranges=[]), pending, applied], + ) + + with patch('singlestoredb.management.timing.time.sleep'): + out = mgr._wait_on_firewall( + self._cluster(firewall_ranges=[]), interval=1, + ) + + self.assertIs(out, applied) + self.assertEqual(mgr.get_cluster.call_count, 3) + + def test_wait_on_firewall_times_out(self): + mgr = self._make_cluster_manager() + mgr.get_cluster = MagicMock(return_value=self._cluster(firewall_ranges=[])) + + with patch('singlestoredb.management.timing.time.sleep'): + with self.assertRaises(ManagementError) as cm: + mgr._wait_on_firewall( + self._cluster(firewall_ranges=[]), interval=1, timeout=3, + ) + assert 'cl-1' in cm.exception.msg, cm.exception.msg + assert 'refuses all inbound' in cm.exception.msg, cm.exception.msg + + def test_wait_on_firewall_expected_waits_for_the_new_ranges(self): + """ + On an existing cluster, non-empty says nothing -- the pre-PATCH ranges + are already non-empty -- so the update path waits for the ranges asked + for. + """ + mgr = self._make_cluster_manager() + new = self._cluster(firewall_ranges=['192.168.0.0/16']) + mgr.get_cluster = MagicMock( + side_effect=[self._cluster(firewall_ranges=['0.0.0.0/0']), new], + ) + + with patch('singlestoredb.management.timing.time.sleep'): + out = mgr._wait_on_firewall( + self._cluster(firewall_ranges=['0.0.0.0/0']), + interval=1, expected=['192.168.0.0/16'], + ) + + self.assertIs(out, new) + self.assertEqual(mgr.get_cluster.call_count, 2) + + def _create(self, mgr, **kwargs): + post_response = MagicMock() + post_response.json.return_value = {'clusterID': 'cl-1'} + mgr._post = MagicMock(return_value=post_response) + with patch('singlestoredb.management.timing.time.sleep'): + return mgr.create_cluster( + 'my-cluster', provider='AWS', region='us-east-1', + project=FAKE_PROJECT_ID, wait_interval=1, **kwargs, + ) + + def test_create_cluster_waits_on_the_firewall(self): + mgr = self._make_cluster_manager() + applied = self._cluster(firewall_ranges=['0.0.0.0/0']) + mgr.get_cluster = MagicMock( + side_effect=[ + self._cluster(firewall_ranges=[]), + self._cluster(firewall_ranges=[]), + applied, + ], + ) + + out = self._create( + mgr, firewall_ranges=['0.0.0.0/0'], wait_on_active=True, + ) + self.assertIs(out, applied) + self.assertEqual(mgr.get_cluster.call_count, 3) + + def test_create_cluster_waits_on_the_firewall_for_allow_all_traffic(self): + mgr = self._make_cluster_manager() + applied = self._cluster(firewall_ranges=[], allow_all_traffic=True) + mgr.get_cluster = MagicMock( + side_effect=[self._cluster(firewall_ranges=[]), applied], + ) + + out = self._create(mgr, allow_all_traffic=True, wait_on_active=True) + self.assertIs(out, applied) + self.assertEqual(mgr.get_cluster.call_count, 2) + + def test_create_cluster_accepts_allow_all_traffic_as_the_applied_form(self): + """ + ``firewall_ranges=['0.0.0.0/0']`` comes back as ``allowAllTraffic``. + + Verified live: the API stores it that way and leaves ``firewallRanges`` + empty, and the endpoint accepts connections. Waiting for non-empty + ranges here would hang for the full timeout on a cluster that is + already reachable. + """ + mgr = self._make_cluster_manager() + applied = self._cluster(firewall_ranges=[], allow_all_traffic=True) + mgr.get_cluster = MagicMock( + side_effect=[self._cluster(firewall_ranges=[]), applied], + ) + + out = self._create( + mgr, firewall_ranges=['0.0.0.0/0'], wait_on_active=True, + ) + self.assertIs(out, applied) + self.assertEqual(mgr.get_cluster.call_count, 2) + + def test_wait_on_firewall_expected_accepts_allow_all_traffic(self): + """A requested 0.0.0.0/0 is satisfied by allow_all_traffic.""" + mgr = self._make_cluster_manager() + applied = self._cluster(firewall_ranges=[], allow_all_traffic=True) + mgr.get_cluster = MagicMock(side_effect=[applied]) + + with patch('singlestoredb.management.timing.time.sleep'): + out = mgr._wait_on_firewall( + self._cluster(firewall_ranges=['10.0.0.0/8']), + interval=1, expected=['0.0.0.0/0'], + ) + self.assertIs(out, applied) + + # ...but a narrower range is not. + mgr.get_cluster = MagicMock( + return_value=self._cluster( + firewall_ranges=[], allow_all_traffic=True, + ), + ) + with patch('singlestoredb.management.timing.time.sleep'): + with self.assertRaises(ManagementError): + mgr._wait_on_firewall( + self._cluster(firewall_ranges=['10.0.0.0/8']), + interval=1, timeout=3, expected=['192.168.0.0/16'], + ) + + def test_create_cluster_does_not_wait_without_a_firewall_request(self): + """ + ``firewall_ranges=[]`` is a legitimate deny-all request -- the field + must be present and an empty list disallows all inbound traffic -- so + it must not hang waiting for a non-empty value that never comes. + """ + for ranges in ([], None): + with self.subTest(firewall_ranges=ranges): + mgr = self._make_cluster_manager() + created = self._cluster(firewall_ranges=ranges) + mgr.get_cluster = MagicMock(return_value=created) + + out = self._create( + mgr, firewall_ranges=ranges, wait_on_active=True, + ) + self.assertIs(out, created) + self.assertEqual(mgr.get_cluster.call_count, 1) + + def test_create_cluster_does_not_wait_without_wait_on_active(self): + mgr = self._make_cluster_manager() + created = self._cluster(firewall_ranges=[]) + mgr.get_cluster = MagicMock(return_value=created) + + out = self._create(mgr, firewall_ranges=['0.0.0.0/0']) + self.assertIs(out, created) + self.assertEqual(mgr.get_cluster.call_count, 1) + + def test_update_waits_only_when_asked(self): + mgr = self._make_cluster_manager() + mgr._patch = MagicMock() + cluster = self._cluster(firewall_ranges=['0.0.0.0/0'], manager=mgr) + + # Without wait_on_active, only the trailing refresh() re-fetches, and + # it reports the pre-PATCH ranges. + stale = self._cluster(firewall_ranges=['0.0.0.0/0'], manager=mgr) + mgr.get_cluster = MagicMock(return_value=stale) + cluster.update(firewall_ranges=['192.168.0.0/16']) + self.assertEqual(mgr.get_cluster.call_count, 1) + self.assertEqual(cluster.firewall_ranges, ['0.0.0.0/0']) + + # With it, the new ranges are polled for. + mgr.get_cluster = MagicMock( + side_effect=[ + self._cluster(firewall_ranges=['0.0.0.0/0'], manager=mgr), + self._cluster(firewall_ranges=['192.168.0.0/16'], manager=mgr), + self._cluster(firewall_ranges=['192.168.0.0/16'], manager=mgr), + ], + ) + with patch('singlestoredb.management.timing.time.sleep'): + cluster.update( + firewall_ranges=['192.168.0.0/16'], + wait_on_active=True, wait_interval=1, + ) + self.assertEqual(mgr.get_cluster.call_count, 3) + self.assertEqual(cluster.firewall_ranges, ['192.168.0.0/16']) + + def test_update_nests_the_size_and_scale_factor(self): + """Resizing goes out as a nested object, not as a bare string.""" + mgr = self._make_cluster_manager() + mgr._patch = MagicMock() + mgr.get_cluster = MagicMock(return_value=self._cluster(manager=mgr)) + cluster = self._cluster(manager=mgr) + + cluster.update(size='S-1', scale_factor=2.0) + + body = mgr._patch.call_args[1]['json'] + self.assertEqual(body['sizeConfig'], {'size': 'S-1', 'scaleFactor': 2.0}) + self.assertNotIn('size', body) + + +class TestProjects(unittest.TestCase): + """ + Projects and the project ID ``create_cluster`` sends. + + ``POST /v2/clusters`` rejects a body without ``projectID`` -- verified + against a live v2 organization -- where ``POST /v1/workspaceGroups`` + assigned one implicitly. So a v2 create has to resolve a project first. + """ + + #: A ``GET /v2/projects`` response, as returned by the live API. + PROJECTS = [ + { + 'createdAt': '2025-10-15T11:22:33.454592Z', + 'edition': 'SHARED', + 'name': 'Shared Project', + 'projectID': FAKE_SHARED_PROJECT_ID, + }, + { + 'createdAt': '2025-10-15T11:22:33.454592Z', + 'edition': 'STANDARD', + 'name': 'Standard Project', + 'projectID': FAKE_STANDARD_PROJECT_ID, + }, + ] + + def _make_cluster_manager(self, projects=None): + from singlestoredb.management.v2.cluster import ClusterManager + with patch( + 'singlestoredb.management.manager.get_token', + return_value=FAKE_TOKEN, + ): + mgr = ClusterManager( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + version='v2', + ) + if projects is not None: + get_response = MagicMock() + get_response.json.return_value = projects + mgr._get = MagicMock(return_value=get_response) + return mgr + + def _without_env(self): + """ + Patch the environment with the deployment variables removed. + + ``SINGLESTOREDB_WORKSPACE`` is what ``_resolve_project_id`` reads, so it + has to go for the fall-through cases to be reached. ``SINGLESTOREDB_ + PROJECT`` goes too, so that a test running in a notebook cannot pass by + accident on a variable the resolver is supposed to ignore. + """ + ctx = patch.dict(os.environ) + ctx.start() + os.environ.pop('SINGLESTOREDB_PROJECT', None) + os.environ.pop('SINGLESTOREDB_WORKSPACE', None) + self.addCleanup(ctx.stop) + + def _in_deployment(self, mgr, project_id): + """Present ``mgr`` as running in a deployment in ``project_id``.""" + self._without_env() + os.environ['SINGLESTOREDB_WORKSPACE'] = FAKE_CLUSTER_ID + mgr.get_cluster = MagicMock( + return_value=MagicMock(project=Project(id=project_id, name='p')), + ) + return mgr + + def test_projects_lists_from_the_projects_endpoint(self): + mgr = self._make_cluster_manager(self.PROJECTS) + projects = mgr.projects + mgr._get.assert_called_once_with('projects') + self.assertIsInstance(projects, NamedList) + self.assertEqual( + [x.id for x in projects], + [FAKE_SHARED_PROJECT_ID, FAKE_STANDARD_PROJECT_ID], + ) + self.assertEqual([x.edition for x in projects], ['SHARED', 'STANDARD']) + # NamedList lookup works by name and by ID. + self.assertEqual(projects['Standard Project'].id, FAKE_STANDARD_PROJECT_ID) + self.assertEqual(projects[FAKE_SHARED_PROJECT_ID].name, 'Shared Project') + self.assertEqual(projects[0].created_at.year, 2025) + + def test_get_project(self): + mgr = self._make_cluster_manager(self.PROJECTS[1]) + project = mgr.get_project(FAKE_STANDARD_PROJECT_ID) + mgr._get.assert_called_once_with(f'projects/{FAKE_STANDARD_PROJECT_ID}') + self.assertEqual(project.name, 'Standard Project') + + def test_explicit_project_id_wins_over_the_current_deployment(self): + mgr = self._in_deployment( + self._make_cluster_manager(), FAKE_STANDARD_PROJECT_ID, + ) + self.assertEqual( + mgr._resolve_project_id(FAKE_PROJECT_ID), FAKE_PROJECT_ID, + ) + # The caller settled it, so the deployment is never fetched. + mgr.get_cluster.assert_not_called() + + def test_the_current_deployment_supplies_the_default_project(self): + """ + A new cluster lands in the project the current one is in. + + This is what makes ``IN PROJECT`` optional in a notebook attached to a + deployment, even in an organization with several projects. + """ + mgr = self._in_deployment( + self._make_cluster_manager(self.PROJECTS), FAKE_STANDARD_PROJECT_ID, + ) + self.assertEqual(mgr._resolve_project_id(), FAKE_STANDARD_PROJECT_ID) + mgr.get_cluster.assert_called_once_with(FAKE_CLUSTER_ID) + # The deployment reports an ID, so no project listing is needed. + mgr._get.assert_not_called() + + def test_an_unresolvable_deployment_falls_through(self): + """ + A deployment that cannot be read is not an error here. + + The variable also names starter clusters, which are not clusters, and + can go stale. Either way there are further defaults to try, so the + lookup failing must not surface. + """ + mgr = self._in_deployment( + self._make_cluster_manager(self.PROJECTS[:1]), FAKE_PROJECT_ID, + ) + mgr.get_cluster.side_effect = ManagementError( + errno=404, msg='cluster not found', + ) + self.assertEqual(mgr._resolve_project_id(), FAKE_SHARED_PROJECT_ID) + + def test_singlestoredb_project_is_not_a_management_project(self): + """ + ``SINGLESTOREDB_PROJECT`` is an inference API project and is ignored. + + The notebook environment sets it to an ID that draws ``404 project not + found`` from ``GET /v2/projects/{id}``. Reading it here made every + ``CREATE CLUSTER`` from a notebook fail, so the resolver must not look + at it at all -- not even as a hint. + """ + self._without_env() + mgr = self._make_cluster_manager(self.PROJECTS) + with patch.dict( + os.environ, {'SINGLESTOREDB_PROJECT': FAKE_STANDARD_PROJECT_ID}, + ): + with self.assertRaises(ManagementError) as cm: + mgr._resolve_project_id() + # Ignored, so this is the ordinary "more than one project" refusal. + self.assertIn('more than one', str(cm.exception)) + + def test_a_project_may_be_named_instead_of_identified(self): + mgr = self._make_cluster_manager(self.PROJECTS) + self.assertEqual( + mgr._resolve_project_id('Standard Project'), + FAKE_STANDARD_PROJECT_ID, + ) + mgr._get.assert_called_once_with('projects') + + def test_a_project_object_may_be_passed_instead_of_a_name(self): + mgr = self._make_cluster_manager(self.PROJECTS) + project = mgr.projects['Standard Project'] + mgr._get.reset_mock() + self.assertEqual( + mgr._resolve_project_id(project), FAKE_STANDARD_PROJECT_ID, + ) + # A Project carries its ID, so no lookup is needed. + mgr._get.assert_not_called() + + def test_an_unknown_project_name_raises_and_lists_the_projects(self): + mgr = self._make_cluster_manager(self.PROJECTS) + with self.assertRaises(ManagementError) as cm: + mgr._resolve_project_id('Nonexistent Project') + msg = str(cm.exception) + self.assertIn('Nonexistent Project', msg) + self.assertIn('Standard Project', msg) + self.assertIn(FAKE_SHARED_PROJECT_ID, msg) + + def test_an_ambiguous_project_name_raises(self): + # The API does not promise unique names, so two projects may share one. + twins = [ + dict(self.PROJECTS[0], name='Twin'), + dict(self.PROJECTS[1], name='Twin'), + ] + mgr = self._make_cluster_manager(twins) + with self.assertRaises(ManagementError) as cm: + mgr._resolve_project_id('Twin') + msg = str(cm.exception) + self.assertIn(FAKE_SHARED_PROJECT_ID, msg) + self.assertIn(FAKE_STANDARD_PROJECT_ID, msg) + + def test_create_starter_cluster_resolves_a_project_name(self): + mgr = self._make_cluster_manager(self.PROJECTS) + post_response = MagicMock() + post_response.json.return_value = {'virtualClusterID': 'vc-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_starter_cluster = MagicMock() + + mgr.create_starter_cluster( + 'my-starter', database_name='db1', provider='AWS', + region='us-east-1', project='Standard Project', + ) + self.assertEqual( + mgr._post.call_args[1]['json']['projectID'], + FAKE_STANDARD_PROJECT_ID, + ) + + def test_a_sole_project_is_the_default(self): + self._without_env() + mgr = self._make_cluster_manager(self.PROJECTS[:1]) + self.assertEqual(mgr._resolve_project_id(), FAKE_SHARED_PROJECT_ID) + + def test_more_than_one_project_raises_and_names_them(self): + self._without_env() + mgr = self._make_cluster_manager(self.PROJECTS) + with self.assertRaises(ManagementError) as cm: + mgr._resolve_project_id() + msg = str(cm.exception) + self.assertIn(FAKE_SHARED_PROJECT_ID, msg) + self.assertIn('Standard Project', msg) + self.assertIn('project=', msg) + # Never point the caller at a variable that names something else. + self.assertNotIn('SINGLESTOREDB_PROJECT', msg) + + def test_no_projects_raises(self): + self._without_env() + mgr = self._make_cluster_manager([]) + with self.assertRaises(ManagementError): + mgr._resolve_project_id() + + def test_create_cluster_resolves_the_project(self): + self._without_env() + mgr = self._make_cluster_manager(self.PROJECTS[:1]) + post_response = MagicMock() + post_response.json.return_value = {'clusterID': 'cl-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_cluster = MagicMock() + + mgr.create_cluster('my-cluster', provider='AWS', region='us-east-1') + self.assertEqual( + mgr._post.call_args[1]['json']['projectID'], FAKE_SHARED_PROJECT_ID, + ) + + def test_create_cluster_accepts_a_project_object(self): + self._without_env() + mgr = self._make_cluster_manager(self.PROJECTS) + project = Project(id=FAKE_STANDARD_PROJECT_ID, name='Standard Project') + post_response = MagicMock() + post_response.json.return_value = {'clusterID': 'cl-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_cluster = MagicMock() + + mgr.create_cluster( + 'my-cluster', provider='AWS', region='us-east-1', + project=project, + ) + self.assertEqual( + mgr._post.call_args[1]['json']['projectID'], + FAKE_STANDARD_PROJECT_ID, + ) + + def test_create_starter_cluster_accepts_a_project_object(self): + mgr = self._make_cluster_manager(self.PROJECTS) + project = Project(id=FAKE_STANDARD_PROJECT_ID, name='Standard Project') + post_response = MagicMock() + post_response.json.return_value = {'virtualClusterID': 'vc-1'} + mgr._post = MagicMock(return_value=post_response) + mgr.get_starter_cluster = MagicMock() + + mgr.create_starter_cluster( + 'my-starter', database_name='db1', provider='AWS', + region='us-east-1', project=project, + ) + self.assertEqual( + mgr._post.call_args[1]['json']['projectID'], + FAKE_STANDARD_PROJECT_ID, + ) + + +class TestClusterFromDict(unittest.TestCase): + """ + ``Cluster.from_dict`` against a v2 payload. + + .. warning:: UNVERIFIED response shape -- the keys below are the ones the + wrapper reads, not keys observed on the wire. + """ + + def _payload(self, **overrides): + obj = { + 'name': 'my-cluster', + 'clusterID': 'cl-1', + 'state': 'ACTIVE', + # Size is reported as an object, not a bare string. Under + # ``sizeConfig`` since the 2026-08-28 rename; ``size`` is still + # read, and tested below, because the rename was reverted once + # already. + 'sizeConfig': {'size': 'S-00', 'scaleFactor': 1.0}, + 'createdAt': '2024-03-15T12:30:45Z', + 'endpoint': 'svc.example.com', + 'provider': 'AWS', + 'region': 'us-east-1', + 'firewallRanges': ['0.0.0.0/0'], + } + obj.update(overrides) + return obj + + def test_fields_and_timestamps(self): + from singlestoredb.management.v2.cluster import Cluster + mgr = MagicMock() + c = Cluster.from_dict(self._payload(), mgr) + self.assertEqual(c.id, 'cl-1') + self.assertEqual(c.name, 'my-cluster') + self.assertEqual(c.state, 'ACTIVE') + self.assertEqual(c.provider, 'AWS') + self.assertEqual(c.size, 'S-00') + self.assertEqual(c.scale_factor, 1.0) + self.assertEqual(c.created_at.year, 2024) + self.assertEqual(c.created_at.month, 3) + self.assertEqual(c.firewall_ranges, ['0.0.0.0/0']) + + def test_either_spelling_of_the_size_object_is_read(self): + from singlestoredb.management.v2.cluster import Cluster + mgr = MagicMock() + + payload = self._payload() + payload.pop('sizeConfig') + payload['size'] = {'size': 'S-1', 'scaleFactor': 2.0} + c = Cluster.from_dict(payload, mgr) + self.assertEqual(c.size, 'S-1') + self.assertEqual(c.scale_factor, 2.0) + + # Neither: the wrapper reports no size rather than raising. + payload.pop('size') + c = Cluster.from_dict(payload, mgr) + self.assertIsNone(c.size) + self.assertIsNone(c.scale_factor) + + def test_region_falls_back_to_what_the_cluster_reports(self): + from singlestoredb.management.v2.cluster import Cluster + # A manager reporting no matching region: the cluster's own provider + # and region slug are all there is, and the display name is unknown. + mgr = MagicMock() + mgr.regions = [] + c = Cluster.from_dict(self._payload(), mgr) + self.assertIsInstance(c.region, Region) + self.assertEqual(c.region.region_name, 'us-east-1') + self.assertEqual(c.region.provider, 'AWS') + self.assertIsNone(c.region.id) + + def test_region_is_resolved_against_the_region_list(self): + from singlestoredb.management.v2.cluster import Cluster + # v2 reports the provider slug on a cluster and the display name only + # in the region list, so the two are matched on (provider, region_name). + mgr = MagicMock() + mgr.regions = [ + Region( + name='US West 2 (Oregon)', provider='AWS', + id=None, region_name='us-west-2', + ), + Region( + name='US East 1 (N. Virginia)', provider='AWS', + id=None, region_name='us-east-1', + ), + ] + c = Cluster.from_dict(self._payload(), mgr) + self.assertEqual(c.region.name, 'US East 1 (N. Virginia)') + self.assertEqual(c.region.region_name, 'us-east-1') + + def test_a_cluster_with_no_region_has_none(self): + from singlestoredb.management.v2.cluster import Cluster + payload = self._payload() + del payload['region'] + c = Cluster.from_dict(payload, MagicMock()) + self.assertIsNone(c.region) + + def test_project_is_resolved_against_the_project_list(self): + from singlestoredb.management.v2.cluster import Cluster + # A cluster reports only its projectID, so the name and edition come + # from the manager's cached project list. + mgr = MagicMock() + mgr.projects = [ + Project(id=FAKE_PROJECT_ID, name='Standard Project', edition='STANDARD'), + ] + c = Cluster.from_dict(self._payload(projectID=FAKE_PROJECT_ID), mgr) + self.assertIsInstance(c.project, Project) + self.assertEqual(c.project.id, FAKE_PROJECT_ID) + self.assertEqual(c.project.name, 'Standard Project') + self.assertEqual(c.project.edition, 'STANDARD') + + def test_project_falls_back_to_the_reported_id(self): + from singlestoredb.management.v2.cluster import Cluster + # An ID the project list does not know about still yields a Project, so + # cluster.project.id is readable either way. + mgr = MagicMock() + mgr.projects = [] + c = Cluster.from_dict(self._payload(projectID=FAKE_PROJECT_ID), mgr) + self.assertIsInstance(c.project, Project) + self.assertEqual(c.project.id, FAKE_PROJECT_ID) + self.assertIsNone(c.project.edition) + + def test_a_cluster_with_no_project_has_none(self): + from singlestoredb.management.v2.cluster import Cluster + mgr = MagicMock() + mgr.projects = [] + self.assertIsNone(Cluster.from_dict(self._payload(), mgr).project) + + def test_no_manager_raises(self): + from singlestoredb.management.v2.cluster import Cluster + mgr = MagicMock() + c = Cluster.from_dict(self._payload(), mgr) + c._manager = None + with self.assertRaises(ManagementError) as cm: + c.refresh() + self.assertIn('cluster manager', cm.exception.msg) + with self.assertRaises(ManagementError): + c.terminate() + + def test_missing_endpoint_blocks_connect(self): + from singlestoredb.management.v2.cluster import Cluster + c = Cluster.from_dict(self._payload(endpoint=None), MagicMock()) + with self.assertRaises(ManagementError) as cm: + c.connect(user='admin', password='x') + self.assertIn('endpoint', cm.exception.msg) + + def test_stage_is_nested_under_the_cluster(self): + from singlestoredb.management.v2.cluster import Cluster + c = Cluster.from_dict(self._payload(), MagicMock()) + self.assertEqual( + c.stage._fs_path('a.sql'), 'clusters/cl-1/stage/fs/a.sql', + ) + + +class TestStatementRoundTrips(unittest.TestCase): + """ + What a whole Fusion statement costs in requests. + + ``CountingManager`` pins the requests a single :class:`Stage` call makes + (``test_management_utils.py``); this pins the requests a statement makes, + deployment resolution included, which is where the redundant ones were. + """ + + def _local_file(self, tmp): + path = os.path.join(tmp, 'stats.csv') + with open(path, 'w') as f: + f.write('a,b\n1,2\n') + return path + + def _upload(self, suffix='', existing=(), clusters=None): + """Run one ``UPLOAD FILE TO STAGE ... IN ''`` and return the manager.""" + mgr = utils.counting_cluster_manager(existing=existing, clusters=clusters) + with tempfile.TemporaryDirectory() as tmp: + local = self._local_file(tmp) + utils.run_fusion_statement( + "UPLOAD FILE TO STAGE 'stats.csv' " + f"IN '{utils.COUNTING_CLUSTER_NAME}' FROM '{local}'{suffix}", + mgr, + ) + return mgr + + def test_a_plain_upload_costs_three_requests(self): + mgr = self._upload() + self.assertEqual( + mgr.calls, [ + ('GET', 'clusters'), + ('GET', 'stats.csv'), + ('PUT', 'stats.csv'), + ], + ) + + def test_an_upload_fetches_no_projects(self): + # Resolving a deployment by name reads the cluster listing, and + # nothing on that path reads a project, so a lazy Cluster.project + # keeps GET /v2/projects out of an upload entirely. + mgr = self._upload() + self.assertNotIn(('GET', 'projects'), mgr.calls) + + def test_an_overwrite_costs_four_requests(self): + # One metadata GET, not two: _upload branches on the object it already + # fetched rather than asking again through remove()'s is_dir(). + mgr = self._upload(suffix=' OVERWRITE', existing=['stats.csv']) + self.assertEqual( + mgr.calls, [ + ('GET', 'clusters'), + ('GET', 'stats.csv'), + ('DELETE', 'stats.csv'), + ('PUT', 'stats.csv'), + ], + ) + + def test_an_upload_over_a_folder_still_raises(self): + with self.assertRaises(IsADirectoryError) as cm: + self._upload(suffix=' OVERWRITE', existing=['stats.csv/']) + self.assertIn('use rmdir or removedirs', str(cm.exception)) + + def test_a_conflict_without_overwrite_still_raises(self): + with self.assertRaises(OSError) as cm: + self._upload(existing=['stats.csv']) + self.assertIn('stage path already exists', str(cm.exception)) + + def test_a_region_on_the_payload_costs_nothing(self): + # A cluster payload carrying a region used to make from_dict match it + # against ClusterManager.regions, so a realistic listing paid for a + # GET /v2/regions the upload never looked at. Cluster.region is lazy + # for the same reason Cluster.project is. + mgr = self._upload( + clusters=[ + utils.cluster_payload( + utils.COUNTING_CLUSTER_NAME, utils.COUNTING_CLUSTER_ID, + project_id=utils.COUNTING_PROJECT_ID, region='us-east-1', + ), + ], + ) + self.assertNotIn(('GET', 'regions'), mgr.calls) + self.assertEqual( + mgr.calls, [ + ('GET', 'clusters'), + ('GET', 'stats.csv'), + ('PUT', 'stats.csv'), + ], + ) + + def _three_clusters(self): + return utils.counting_cluster_manager( + clusters=[ + utils.cluster_payload( + f'c{i}', f'{utils.COUNTING_CLUSTER_ID[:-1]}{i}', + project_id=utils.COUNTING_PROJECT_ID, region='us-east-1', + ) + for i in range(3) + ], + ) + + def test_show_clusters_extended_reports_the_project_once(self): + # .project is lazy now, so EXTENDED reads it per row -- and the + # one-hour ttl_property on ClusterManager.projects is what keeps that + # at one GET /v2/projects however many rows there are. + mgr = self._three_clusters() + res = utils.run_fusion_statement('SHOW CLUSTERS EXTENDED', mgr) + columns = [x[0] for x in res.description] + rows = [dict(zip(columns, row)) for row in res.rows] + self.assertEqual(len(rows), 3) + self.assertEqual( + [x['ProjectName'] for x in rows], ['Test Project'] * 3, + ) + self.assertEqual(mgr.calls.count(('GET', 'projects')), 1) + + def test_show_clusters_extended_reports_the_region_once(self): + # Same shape for the lazy region: read per row, fetched once. + mgr = self._three_clusters() + res = utils.run_fusion_statement('SHOW CLUSTERS EXTENDED', mgr) + columns = [x[0] for x in res.description] + rows = [dict(zip(columns, row)) for row in res.rows] + self.assertEqual([x['Region'] for x in rows], ['us-east-1'] * 3) + self.assertEqual(mgr.calls.count(('GET', 'regions')), 1) + + def test_printing_a_cluster_costs_nothing(self): + # vars_to_str skips the underscored attributes the lazy properties are + # stored in, so Cluster.__str__ hands it the unresolved ID / name. + # Printing a cluster must not turn into two requests. + mgr = self._three_clusters() + cluster = mgr.clusters[0] + before = list(mgr.calls) + text = str(cluster) + self.assertEqual(mgr.calls, before) + self.assertIn(f'project={utils.COUNTING_PROJECT_ID!r}', text) + self.assertIn("region='us-east-1'", text) + + def test_printing_a_cluster_shows_what_is_resolved(self): + # Once something has read the property, the resolved object is what + # gets reported. + mgr = self._three_clusters() + cluster = mgr.clusters[0] + self.assertEqual(cluster.project.name, 'Test Project') + self.assertIn("project=Project(name='Test Project'", str(cluster)) + + +class TestDeploymentEnvVars(unittest.TestCase): + """ + The environment-variable contract the notebook environment publishes. + + There is no ``SINGLESTOREDB_CLUSTER``: the current deployment arrives as + ``SINGLESTOREDB_WORKSPACE`` at every API version, and its value is a + cluster ID at v2. ``SINGLESTOREDB_WORKSPACE_GROUP`` is published too, but + it holds a group ID, which v2 reports only as :attr:`Cluster.group` and + offers no route to look up. + """ + + def _clean_env(self, **values): + ctx = patch.dict(os.environ) + ctx.start() + self.addCleanup(ctx.stop) + for name in ( + 'SINGLESTOREDB_WORKSPACE', + 'SINGLESTOREDB_WORKSPACE_GROUP', + 'SINGLESTOREDB_VIRTUAL_WORKSPACE', + 'SINGLESTOREDB_DEFAULT_DATABASE', + ): + os.environ.pop(name, None) + os.environ.update(values) + + def test_get_cluster_reads_the_workspace_variable(self): + from singlestoredb.management.cluster import get_cluster + self._clean_env(SINGLESTOREDB_WORKSPACE=FAKE_CLUSTER_ID) + mgr = MagicMock() + with patch( + 'singlestoredb.management.cluster.manage_clusters', + return_value=mgr, + ): + get_cluster() + mgr.clusters.__getitem__.assert_called_once_with(FAKE_CLUSTER_ID) + + def test_get_cluster_ignores_the_group_variable(self): + from singlestoredb.management.cluster import get_cluster + # A group ID is not a cluster ID and there is no group route to turn + # one into the other, so this is left unresolved rather than guessed at. + self._clean_env(SINGLESTOREDB_WORKSPACE_GROUP=FAKE_CLUSTER_ID) + with patch( + 'singlestoredb.management.cluster.manage_clusters', + return_value=MagicMock(), + ): + with self.assertRaises(RuntimeError): + get_cluster() + + def test_cluster_id_is_the_workspace_variable(self): + from singlestoredb.management.utils import get_cluster_id + from singlestoredb.management.utils import get_workspace_id + self._clean_env(SINGLESTOREDB_WORKSPACE=FAKE_CLUSTER_ID) + self.assertEqual(get_cluster_id(), FAKE_CLUSTER_ID) + self.assertEqual(get_workspace_id(), FAKE_CLUSTER_ID) + + def test_no_deployment_variable_leaves_the_id_unset(self): + from singlestoredb.management.utils import get_cluster_id + self._clean_env() + self.assertIsNone(get_cluster_id()) + + def test_job_target_comes_from_the_workspace_variable(self): + from singlestoredb.management.job import TargetType + from singlestoredb.management.v1.job import JobsManager as V1JobsManager + from singlestoredb.management.v2.job import JobsManager as V2JobsManager + self._clean_env(SINGLESTOREDB_WORKSPACE=FAKE_CLUSTER_ID) + + for manager_cls, target_type in ( + (V2JobsManager, TargetType.CLUSTER), + (V1JobsManager, TargetType.WORKSPACE), + ): + target_config = {} + manager_cls(MagicMock())._resolve_target(target_config) + self.assertEqual( + target_config, + dict(targetID=FAKE_CLUSTER_ID, targetType=target_type.value), + ) + + def test_group_variable_is_not_a_job_target(self): + from singlestoredb.management.v2.job import JobsManager + self._clean_env(SINGLESTOREDB_WORKSPACE_GROUP=FAKE_CLUSTER_ID) + target_config = {} + JobsManager(MagicMock())._resolve_target(target_config) + self.assertEqual(target_config, {}) + + +# +# Live suites. These need SINGLESTOREDB_MANAGEMENT_TOKEN and an organization +# with v2 access, and they create and destroy real deployments. +# + +@pytest.mark.management +class TestCluster(unittest.TestCase): + + manager = None + cluster = None + password = None + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_clusters(version='v2') + + us_regions = _us_regions(cls.manager) + + name = clean_name(secrets.token_urlsafe(20)[:20]) + region = random.choice(us_regions) + + # v2 has no workspace group: the cluster is created in one call, with + # the firewall settings passed alongside the compute settings. + cls.cluster = cls.manager.create_cluster( + f'cl-test-{name}', + region=region, + size='S-00', + firewall_ranges=['0.0.0.0/0'], + project=_project_id(cls.manager), + wait_on_active=True, + ) + + # v2 generates the admin password and reports it only in the create + # response; anything passed as admin_password= is ignored. So the + # password has to be read back rather than chosen here. + cls.password = cls.cluster.admin_password + + # The firewall is applied asynchronously, after the cluster is already + # ACTIVE with a resolvable endpoint; until it lands the cluster admits + # nothing and refuses every inbound connection, so test_connect would + # time out at the TCP level. wait_on_active covers that, and this + # asserts it did -- no polling needed here. + # + # Verified live: a requested firewall_ranges=['0.0.0.0/0'] is stored as + # allow_all_traffic=True with firewall_ranges == [], so either one + # means reachable. + assert cls.cluster.allow_all_traffic or cls.cluster.firewall_ranges, ( + 'create_cluster(wait_on_active=True) returned a cluster whose ' + 'firewall still admits nothing; every inbound connection would be ' + 'refused' + ) + + @classmethod + def tearDownClass(cls): + if cls.cluster is not None: + cls.cluster.terminate(force=True) + cls.cluster = None + cls.manager = None + cls.password = None + + def test_str(self): + assert self.cluster.name in str(self.cluster) + + def test_repr(self): + assert repr(self.cluster) == str(self.cluster) + + def test_regions(self): + out = self.manager.regions + providers = {x.provider for x in out} + assert any( + p in providers for p in ('Azure', 'GCP', 'AWS', 'azure', 'gcp', 'aws') + ), providers + # v2 regions carry no ID, so they are addressable by name only. + for region in out: + assert region.id is None, region + + def test_clusters(self): + clusters = self.manager.clusters + ids = [x.id for x in clusters] + names = [x.name for x in clusters] + assert self.cluster.id in ids + assert self.cluster.name in names + + assert clusters.ids() == ids + assert clusters.names() == names + + objs = {} + for item in clusters: + # setdefault, and name before id, so this resolves a key the way + # NamedList._find_item does: to the *first* match. Plain assignment + # kept the last, which disagrees as soon as the listing carries two + # entries of one name -- GET /v2/clusters reports a terminated + # cluster alongside its live replacement, so that happens. + objs.setdefault(item.name, item) + objs.setdefault(item.id, item) + + name = random.choice(names) + assert clusters[name] == objs[name] + id = random.choice(ids) + assert clusters[id] == objs[id] + + def test_get_cluster(self): + cluster = self.manager.get_cluster(self.cluster.id) + assert cluster.id == self.cluster.id, cluster.id + + with self.assertRaises(s2.ManagementError): + self.manager.get_cluster('bad id') + + def test_update(self): + """ + Update the firewall, and show that ``name`` is not updatable. + + Both halves live in one test because each ``PATCH /v2/clusters/{id}`` + cycles the cluster back through PENDING, and a second test issuing its + own PATCH while that is in flight is asking for trouble. + + On the name: verified live that v2 clusters cannot be renamed, unlike + v1 workspace groups. ``name`` is a *known* field on the PATCH route -- + an unknown field draws ``400 request body contains an unknown field`` + and ``name`` does not -- and the request succeeds, even cycling the + cluster through PENDING, but the name never changes in either + ``GET /v2/clusters/{id}`` or ``GET /v2/clusters`` (polled for two + minutes). Pinned here so the API growing real rename support is + noticed rather than assumed. + """ + # setUpClass asked for ['0.0.0.0/0'], which the API may store either + # verbatim or as allow_all_traffic with the ranges left empty. + opened = self.cluster.allow_all_traffic \ + or self.cluster.firewall_ranges == ['0.0.0.0/0'] + assert opened, ( + self.cluster.allow_all_traffic, self.cluster.firewall_ranges, + ) + + # The PATCH is applied asynchronously: without wait_on_active the + # refresh() inside update() still reports the old ranges. + self.cluster.update( + firewall_ranges=['192.168.0.0/16'], wait_on_active=True, + ) + + cluster = self.cluster + assert cluster.firewall_ranges == ['192.168.0.0/16'], \ + cluster.firewall_ranges + + name = cluster.name.replace('cl-test-', 'cl-foo-') + assert name != cluster.name + cluster.update(name=name) + + assert cluster.name != name, cluster.name + assert self.manager.get_cluster(cluster.id).name != name + + def test_no_manager(self): + cluster = self.manager.get_cluster(self.cluster.id) + cluster._manager = None + + with self.assertRaises(s2.ManagementError) as cm: + cluster.refresh() + assert 'cluster manager' in cm.exception.msg, cm.exception.msg + + with self.assertRaises(s2.ManagementError) as cm: + cluster.terminate() + assert 'cluster manager' in cm.exception.msg, cm.exception.msg + + def test_connect(self): + with self.cluster.connect(user='admin', password=self.password) as conn: + with conn.cursor() as cur: + cur.execute('show databases') + assert 'cluster' in [x[0] for x in list(cur)] + + # Test missing endpoint + cluster = self.manager.get_cluster(self.cluster.id) + cluster.endpoint = None + + with self.assertRaises(s2.ManagementError) as cm: + cluster.connect(user='admin', password=self.password) + assert 'endpoint' in cm.exception.msg, cm.exception.msg + + +@pytest.mark.management +class TestStarterCluster(unittest.TestCase): + + manager = None + starter_cluster = None + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_clusters(version='v2') + + # Starter regions come from GET /v2/regions/sharedtier, which answers + # at v2 with the same shape as GET /v2/regions. Only regions on that + # list work -- anything else gets a 500 'no shared tier region found + # for provider X and region Y' out of POST /v2/sharedtier/ + # virtualClusters -- so discover rather than sampling all regions. + # US-only where possible, matching the v1 starter test: non-US regions + # are likelier to answer a creation with a control-plane 500. Fall back + # to the full list rather than skipping if the org has no US region. + all_regions = list(cls.manager.shared_tier_regions) + regions = [x for x in all_regions if 'US' in x.name] or all_regions + if not regions: + raise unittest.SkipTest( + 'no shared-tier capable region is available to this ' + 'organization', + ) + + name = shared_database_name(secrets.token_urlsafe(20)[:20]) + + # Namespaced for the same reason as the database: the starter-tier user + # name has to be unique across the project's starter deployments, not + # just within this one, so a fixed name collides with + # TestStarterWorkspace in test_management_v1 -- which runs on another + # xdist worker -- and with anything an earlier failed run leaked. The + # API reports the collision as a bare 500. + cls.starter_username = f'starter_user_{name[:8]}' + cls.password = secrets.token_urlsafe(20) + + cls.database_name = f'starter_db_{name}' + + region = random.choice(regions) + + cls.starter_cluster = cls.manager.create_starter_cluster( + f'starter-cl-test-{name}', + database_name=cls.database_name, + region=region, + ) + + cls.starter_cluster.create_user( + username=cls.starter_username, + password=cls.password, + ) + + @classmethod + def tearDownClass(cls): + if cls.starter_cluster is not None: + cls.starter_cluster.terminate() + cls.starter_cluster = None + cls.manager = None + cls.password = None + + def test_str(self): + assert self.starter_cluster.name in str(self.starter_cluster) + + def test_repr(self): + assert repr(self.starter_cluster) == str(self.starter_cluster) + + def test_get_starter_cluster(self): + cluster = self.manager.get_starter_cluster(self.starter_cluster.id) + assert cluster.id == self.starter_cluster.id, cluster.id + + with self.assertRaises(s2.ManagementError): + self.manager.get_starter_cluster('bad id') + + def test_starter_clusters(self): + clusters = self.manager.starter_clusters + ids = [x.id for x in clusters] + names = [x.name for x in clusters] + assert self.starter_cluster.id in ids + assert self.starter_cluster.name in names + + objs = {} + for item in clusters: + # setdefault, and name before id, so this resolves a key the way + # NamedList._find_item does: to the *first* match. Plain assignment + # kept the last, which disagrees as soon as the listing carries two + # entries of one name -- GET /v2/clusters reports a terminated + # cluster alongside its live replacement, so that happens. + objs.setdefault(item.name, item) + objs.setdefault(item.id, item) + + name = random.choice(names) + assert clusters[name] == objs[name] + id = random.choice(ids) + assert clusters[id] == objs[id] + + def test_no_manager(self): + cluster = self.manager.get_starter_cluster(self.starter_cluster.id) + cluster._manager = None + + with self.assertRaises(s2.ManagementError) as cm: + cluster.refresh() + assert 'cluster manager' in cm.exception.msg, cm.exception.msg + + with self.assertRaises(s2.ManagementError) as cm: + cluster.terminate() + assert 'cluster manager' in cm.exception.msg, cm.exception.msg + + def test_connect(self): + with self.starter_cluster.connect( + user=self.starter_username, + password=self.password, + ) as conn: + with conn.cursor() as cur: + cur.execute('show databases') + assert self.database_name in [x[0] for x in list(cur)] + + # Test missing endpoint + cluster = self.manager.get_starter_cluster(self.starter_cluster.id) + cluster.endpoint = None + + with self.assertRaises(s2.ManagementError) as cm: + cluster.connect(user=self.starter_username, password=self.password) + assert 'endpoint' in cm.exception.msg, cm.exception.msg + + +@pytest.mark.management +@pytest.mark.xdist_group(utils.SHARED_CLUSTER_STAGE_GROUP) +class TestStage(unittest.TestCase): + """ + Stage at v2 hangs off the cluster (``clusters/{id}/stage/fs/``) rather + than being a top-level resource keyed by workspace group. + """ + + manager = None + cluster = None + password = None + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_clusters(version='v2') + + # UNVERIFIED: v1 could reach a stage from a workspace group without + # ever starting a workspace. At v2 there is no group, so a cluster has + # to exist for its stage to be addressable. + # + # A shared one: every assertion below is scoped to one path, and every + # path is namespaced with id(self), so what another class left in this + # cluster's stage is invisible here. See utils.shared_clusters. + cls.cluster = utils.shared_clusters(1)[0] + + # v2 generates the admin password; see TestCluster.setUpClass. + cls.password = cls.cluster.admin_password + + @classmethod + def tearDownClass(cls): + # The cluster is the shared pool's: it stays live for the classes that + # follow and is terminated once, at the end of the session. + cls.cluster = None + cls.manager = None + cls.password = None + + def test_root_info(self): + st = self.cluster.stage + root = st.info('/') + assert str(root.path) == '/' + assert root.type == 'directory' + + def test_upload_file(self): + st = self.cluster.stage + + upload_test_sql = f'upload_test_{id(self)}.sql' + upload_test2_sql = f'upload_test2_{id(self)}.sql' + + f = st.upload_file(f'{TEST_DIR}/test.sql', upload_test_sql) + assert str(f.path) == upload_test_sql + assert f.type == 'file' + + txt = f.download(encoding='utf-8') + assert txt == open(f'{TEST_DIR}/test.sql').read() + + # No silent overwrite + with self.assertRaises(OSError): + st.upload_file(f'{TEST_DIR}/test.sql', upload_test_sql) + + f = st.upload_file( + open(f'{TEST_DIR}/test2.sql', 'r'), + upload_test_sql, + overwrite=True, + ) + txt = f.download(encoding='utf-8') + assert txt == open(f'{TEST_DIR}/test2.sql').read() + + with self.assertRaises(IsADirectoryError): + st.upload_file(TEST_DIR, 'test3.sql') + + lib = st.mkdir(f'/lib_{id(self)}/') + assert lib.type == 'directory' + + with self.assertRaises(IsADirectoryError): + st.upload_file(f'{TEST_DIR}/test2.sql', lib.path, overwrite=True) + + f = st.upload_file( + f'{TEST_DIR}/test2.sql', + os.path.join(lib.path, upload_test2_sql), + ) + assert str(f.path) == f'{lib.path}{upload_test2_sql}' + assert f.type == 'file' + + def test_open(self): + st = self.cluster.stage + open_test_sql = f'open_test_{id(self)}.sql' + + with st.open(open_test_sql, 'w') as f: + f.write('create table foo (id int);') + + with st.open(open_test_sql, 'r') as f: + assert f.read() == 'create table foo (id int);' + + # Reading a missing object fails. Note that this raises + # ManagementError rather than the FileNotFoundError the rest of + # Stage.open's builtin-open emulation would suggest -- the 404 from + # the download comes straight back out. Verified live; asserted here + # so a change to it is deliberate rather than accidental. + with self.assertRaises(s2.ManagementError) as cm: + st.open(f'missing_{id(self)}.sql', 'r') + assert cm.exception.errno == 404, cm.exception.errno + + def test_listdir_and_remove(self): + st = self.cluster.stage + name = f'listdir_test_{id(self)}.sql' + + st.upload_file(f'{TEST_DIR}/test.sql', name) + assert name in [str(x) for x in st.listdir('/')] + assert st.exists(name) + assert st.is_file(name) + assert not st.is_dir(name) + + st.remove(name) + assert not st.exists(name) + + def test_rename(self): + st = self.cluster.stage + src = f'rename_src_{id(self)}.sql' + dst = f'rename_dst_{id(self)}.sql' + + st.upload_file(f'{TEST_DIR}/test.sql', src) + st.rename(src, dst) + assert not st.exists(src) + assert st.exists(dst) + st.remove(dst) + + def test_mkdir_and_rmdir(self): + st = self.cluster.stage + d = f'dir_{id(self)}' + + # mkdir() and rmdir() append the trailing slash themselves, but + # exists()/is_dir()/info() do not: without it the metadata GET 404s + # and is_dir() reports False. The v1 suite passes the slash + # explicitly for the same reason. + st.mkdir(d) + assert st.is_dir(f'{d}/') + assert not st.is_file(f'{d}/') + st.rmdir(d) + assert not st.exists(f'{d}/') + + +@pytest.mark.management +class TestSecrets(unittest.TestCase): + """ + Secrets are organization-scoped, so unlike v1 this needs no deployment. + """ + + manager = None + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_clusters(version='v2') + + @classmethod + def tearDownClass(cls): + cls.manager = None + + def test_get_secret(self): + # A fixed name, deliberately not one built from id(self): that is a + # process-local address, so a name built from it can never match what + # an interrupted run left behind, which makes the cleanup below dead + # code. A secret is org-scoped and permanent and nothing sweeps them, + # so a leaked one is leaked for good. Distinct from the v1 suite's + # 'secret_name' so the two suites do not delete each other's. + name = 'secret_v2_test' + + # Clear a leftover secret from a previous run + try: + leftover = self.manager.organizations.current.get_secret(name) + self.manager._delete(f'secrets/{leftover.id}') + except s2.ManagementError: + pass + + created = self.manager._post( + 'secrets', + json=dict(name=name, value='secret_value'), + ).json() + + # The ID comes from the create response rather than from the lookup + # under test: binding it inside the try would leave the cleanup raising + # UnboundLocalError over whatever the lookup actually failed with. + secret_id = created['secret']['secretID'] + try: + secret = self.manager.organizations.current.get_secret(name) + assert secret.name == name + assert secret.value == 'secret_value' + finally: + self.manager._delete(f'secrets/{secret_id}') + + +@pytest.mark.management +@pytest.mark.xdist_group(utils.SHARED_CLUSTER_JOBS_GROUP) +class TestJob(unittest.TestCase): + """ + Scheduled notebook jobs at v2. + + The one v2-visible difference is the ``targetType`` the SDK sends for a + deployment: v1 called it ``Workspace``, v2 calls it ``Cluster``. + """ + + manager = None + cluster = None + password = None + job_ids = [] + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_clusters(version='v2') + + # A shared cluster: a job only needs a live deployment to name as its + # target, and each assertion here is about the job it just created. + cls.cluster = utils.shared_clusters(1)[0] + + # v2 generates the admin password; see TestCluster.setUpClass. + cls.password = cls.cluster.admin_password + + @classmethod + def tearDownClass(cls): + for job_id in cls.job_ids: + try: + cls.manager.organizations.current.jobs.delete(job_id) + except Exception: + pass + # The cluster is the shared pool's; see TestStage.tearDownClass. + cls.cluster = None + cls.manager = None + cls.password = None + os.environ.pop('SINGLESTOREDB_WORKSPACE', None) + os.environ.pop('SINGLESTOREDB_DEFAULT_DATABASE', None) + + def test_job_without_database_target(self): + os.environ.pop('SINGLESTOREDB_WORKSPACE', None) + os.environ.pop('SINGLESTOREDB_DEFAULT_DATABASE', None) + + job_manager = self.manager.organizations.current.jobs + job = job_manager.run( + 'Scheduling Test.ipynb', + 'notebooks-cpu-small', + {'strParam': 'string', 'intParam': 1, 'floatParam': 1.0, 'boolParam': True}, + ) + self.job_ids.append(job.job_id) + assert job.execution_config.notebook_path == 'Scheduling Test.ipynb' + assert job.schedule.mode == job_manager.modes().ONCE + assert not job.execution_config.create_snapshot + assert job.completed_executions_count == 0 + assert job.target_config is None + job.wait() + job = job_manager.get(job.job_id) + assert job.completed_executions_count == 1 + assert len(job.job_metadata) == 1 + assert job.job_metadata[0].status == Status.COMPLETED + assert job.target_config is None + assert job.delete() + job = job_manager.get(job.job_id) + assert job.terminated_at is not None + + def test_job_with_database_target(self): + os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = 'information_schema' + os.environ['SINGLESTOREDB_WORKSPACE'] = self.cluster.id + + job_manager = self.manager.organizations.current.jobs + job = job_manager.run( + 'Scheduling Test.ipynb', + 'notebooks-cpu-small', + {'strParam': 'string', 'intParam': 1, 'floatParam': 1.0, 'boolParam': True}, + ) + self.job_ids.append(job.job_id) + assert job.target_config is not None + assert job.target_config.database_name == 'information_schema' + assert job.target_config.target_id == self.cluster.id + # The v2 name for a deployment target. + assert job.target_config.target_type == TargetType.CLUSTER + assert not job.target_config.resume_target + job.wait() + job = job_manager.get(job.job_id) + assert job.completed_executions_count == 1 + assert job.job_metadata[0].status == Status.COMPLETED + assert job.target_config.target_type == TargetType.CLUSTER + assert job.delete() + job = job_manager.get(job.job_id) + assert job.terminated_at is not None + + +@pytest.mark.management +class TestRegions(unittest.TestCase): + """Region listing through the standalone region manager.""" + + manager = None + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_regions(version='v2') + + @classmethod + def tearDownClass(cls): + cls.manager = None + + def test_list_regions(self): + regions = self.manager.list_regions() + assert isinstance(regions, NamedList) + assert len(regions) > 0 + + region = regions[0] + assert isinstance(region, Region) + # v2 responses carry no regionID. + assert region.id is None + assert region.name + assert region.provider + # ``region`` is the display name, ``regionName`` the provider slug. + assert region.region_name + + def test_list_shared_tier_regions(self): + regions = self.manager.list_shared_tier_regions() + assert isinstance(regions, NamedList) + assert len(regions) > 0 + + region = regions[0] + assert isinstance(region, Region) + assert region.id is None + assert region.name + assert region.provider + assert region.region_name + + def test_str_repr(self): + regions = self.manager.list_regions() + if not regions: + self.skipTest('No regions available for testing') + + region = regions[0] + s = str(region) + assert region.name in s + assert region.provider in s + assert repr(region) == s + + +if __name__ == '__main__': + unittest.main() diff --git a/singlestoredb/tests/test_management_versioning.py b/singlestoredb/tests/test_management_versioning.py new file mode 100644 index 000000000..97d2f2c19 --- /dev/null +++ b/singlestoredb/tests/test_management_versioning.py @@ -0,0 +1,894 @@ +#!/usr/bin/env python +# type: ignore +""" +Structural tests for the management API's version split. + +These are the only versioning tests worth keeping now that the cross-version +bridge is gone: that the version-module importer reports failures usefully, +that the ``manage_*`` factories route to the right version package, and that +``management/v1/`` and ``management/v2/`` do not import each other -- the +invariant that makes deleting either one an ``rm -rf``. +""" +import ast +import contextlib +import importlib +import os +import subprocess +import sys +import unittest +import warnings +from unittest.mock import patch + +from singlestoredb.exceptions import ManagementError +from singlestoredb.management._version_import import _import_versioned_module + + +FAKE_TOKEN = 'test-token-12345' +FAKE_BASE_URL = 'https://api.example.com' + + +@contextlib.contextmanager +def management_version(value): + """Set the ``management.version`` option, restoring the exact original. + + ``conftest.py``'s ``protect_singlestoredb_url`` does not cover this + option, and restoring with ``original or 'v1'`` would silently rewrite a + ``None``/``''`` original into ``'v1'``. + """ + from singlestoredb import config + original = config.get_option('management.version') + try: + config.set_option('management.version', value) + yield + finally: + config.set_option('management.version', original) + + +class TestImportVersionedModule(unittest.TestCase): + """Test dynamic module import.""" + + def test_import_v1_workspace(self): + mod = _import_versioned_module('v1', 'workspace') + self.assertTrue(hasattr(mod, 'Workspace')) + self.assertTrue(hasattr(mod, 'WorkspaceManager')) + + def test_import_v2_cluster(self): + """v2 has clusters, not workspaces.""" + mod = _import_versioned_module('v2', 'cluster') + self.assertTrue(hasattr(mod, 'Cluster')) + self.assertTrue(hasattr(mod, 'ClusterManager')) + + def test_v2_has_no_workspace_module(self): + with self.assertRaises(ManagementError) as ctx: + _import_versioned_module('v2', 'workspace') + msg = str(ctx.exception) + self.assertIn('workspace', msg) + self.assertIn('v2', msg) + + def test_import_nonexistent_version_raises(self): + with self.assertRaises(ManagementError) as ctx: + _import_versioned_module('v99', 'workspace') + self.assertIn('v99', str(ctx.exception)) + + def test_import_nonexistent_module_raises(self): + with self.assertRaises(ManagementError) as ctx: + _import_versioned_module('v1', 'nonexistent_module') + msg = str(ctx.exception) + # Should NOT claim the version is unsupported when the version + # package itself imports cleanly; should name the missing module. + self.assertNotIn('Unsupported API version', msg) + self.assertIn('nonexistent_module', msg) + self.assertIn('v1', msg) + + +class TestConfigOption(unittest.TestCase): + """Test that management.version config option exists and works.""" + + def test_config_option_exists(self): + from singlestoredb import config + val = config.get_option('management.version') + self.assertIn(val, ('v1', 'v2', None, '')) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_config_option_routes_manage_regions(self, _mock_token): + """Setting management.version to v2 routes to v2.""" + from singlestoredb.management.region import manage_regions + from singlestoredb.management.v2.region import RegionManager as V2RM + + with management_version('v2'): + mgr = manage_regions( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + ) + self.assertIsInstance(mgr, V2RM) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_the_option_does_not_reach_manage_workspaces(self, _mock_token): + """ + Neither workspace factory consults the option -- v1 keeps working. + + Workspaces exist only at v1, so there is nothing for the option to + select between. Flipping the default to v2 must not turn a bare + ``manage_workspaces()`` into an error: that would be v1 ceasing to work + rather than v1 being deprecated. The deprecation warning is what steers + callers to clusters. + """ + from singlestoredb.management.workspace import manage_workspaces + from singlestoredb.management.workspace import _manage_workspaces_v1 + from singlestoredb.management.v1.workspace import ( + WorkspaceManager as V1WM, + ) + + for option in ('v1', 'v2', None): + with self.subTest(option=option), management_version(option): + for label, factory in ( + ('public', manage_workspaces), + ('internal', _manage_workspaces_v1), + ): + with self.subTest(factory=label): + mgr = factory( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + ) + self.assertIsInstance(mgr, V1WM) + self.assertIn('/v1/', mgr._base_url) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_manage_workspaces_still_rejects_an_explicit_other_version( + self, _mock_token, + ): + """Pinning to v1 is not the same as ignoring the argument.""" + from singlestoredb.management.workspace import manage_workspaces + with self.assertRaises(ManagementError) as ctx: + manage_workspaces( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v2', + ) + self.assertIn('manage_clusters', str(ctx.exception)) + + def test_default_version_is_a_literal_not_the_config_option(self): + """ + ``default_version`` must not be frozen from the config option at + import time -- that let a v1 class declare itself to be v2. + + ``Manager`` takes the shared ``DEFAULT_VERSION`` and ``FilesManager`` + inherits it; ``WorkspaceManager`` is a v1 class and pins itself. + Setting the option must move none of them. + """ + from singlestoredb.management.manager import Manager + from singlestoredb.management.v1.workspace import WorkspaceManager + from singlestoredb.management.files import FilesManager + expected = {Manager: 'v2', FilesManager: 'v2', WorkspaceManager: 'v1'} + for value in ('v1', 'v2', None): + with management_version(value): + for cls, want in expected.items(): + self.assertEqual(cls.default_version, want, cls.__name__) + + def test_default_version_ignores_the_environment_variable(self): + """ + The same guard for ``SINGLESTOREDB_MANAGEMENT_VERSION``, which the + in-process check above cannot reach: the option's *registered default* + absorbs the environment variable at import + (``utils/config.py``, ``Option.__init__``), so resolving + ``default_version`` through ``config.get_default()`` would hand a v2 + class a v1 URL whenever the variable was set. A fresh interpreter is + the only way to see it. + """ + script = ( + 'from singlestoredb.management.manager import Manager;' + 'from singlestoredb.management.files import FilesManager;' + 'from singlestoredb.management import _version_import as vi;' + 'from singlestoredb import config;' + 'print(Manager.default_version, FilesManager.default_version,' + ' vi.DEFAULT_VERSION, config.get_option("management.version"))' + ) + env = dict(os.environ, SINGLESTOREDB_MANAGEMENT_VERSION='v1') + out = subprocess.run( + [sys.executable, '-c', script], + env=env, capture_output=True, text=True, check=True, + ).stdout.split() + # The option follows the variable; the class attributes do not. + self.assertEqual(out, ['v2', 'v2', 'v2', 'v1']) + + +class TestManageRoutingForAllFactories(unittest.TestCase): + """ + ``manage_*`` factories must route to the correct version module: + ``version='v2'`` returns a v2 manager, default returns a v1 manager. + """ + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_manage_workspaces(self, _mock_token): + """Workspaces are v1-only; an explicit v2 is refused, v1 still works.""" + from singlestoredb.management.workspace import manage_workspaces + from singlestoredb.management.v1.workspace import ( + WorkspaceManager as V1WM, + ) + + with self.assertRaises(ManagementError): + manage_workspaces( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v2', + ) + v1 = manage_workspaces( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v1', + ) + self.assertIsInstance(v1, V1WM) + # A bare call is pinned to v1 rather than resolved through the option, + # so flipping the default to v2 left it working. Covered in full by + # TestConfigOption.test_the_option_does_not_reach_manage_workspaces. + for option in ('v1', None): + with management_version(option): + self.assertIsInstance( + manage_workspaces( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, + ), + V1WM, + ) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_manage_clusters(self, _mock_token): + """ + ``manage_clusters`` follows ``management.version``. + + Clusters are v2-only, so a resolved ``v1`` raises whether it came from + the caller or from the option. The option defaults to ``v2`` now; the + live v2 suites still pass ``version='v2'`` explicitly so that they do + not start testing v1 if the option is ever pointed back. + """ + from singlestoredb.management.cluster import manage_clusters + from singlestoredb.management.cluster import DEFAULT_CLUSTER_VERSION + from singlestoredb.management.v2.cluster import ClusterManager as V2CM + + v2 = manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v2', + ) + self.assertIsInstance(v2, V2CM) + self.assertIn('/v2/', v2._base_url) + + # The option is followed, and beats DEFAULT_CLUSTER_VERSION. + with management_version('v2'): + default = manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, + ) + self.assertIsInstance(default, V2CM) + self.assertIn('/v2/', default._base_url) + + # Unset option: DEFAULT_CLUSTER_VERSION is the fallback. + self.assertEqual(DEFAULT_CLUSTER_VERSION, 'v2') + with management_version(None): + self.assertIsInstance( + manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, + ), + V2CM, + ) + + # v1 raises, whether asked for outright... + with self.assertRaises(ManagementError): + manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v1', + ) + # ...or inherited from the option. + with management_version('v1'): + with self.assertRaises(ManagementError): + manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, + ) + # An explicit v2 still overrides it. + self.assertIsInstance( + manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, + version='v2', + ), + V2CM, + ) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_manage_regions(self, _mock_token): + from singlestoredb.management.region import manage_regions + from singlestoredb.management.v1.region import RegionManager as V1RM + from singlestoredb.management.v2.region import RegionManager as V2RM + + self.assertIsInstance( + manage_regions( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v2', + ), + V2RM, + ) + self.assertIsInstance( + manage_regions( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v1', + ), + V1RM, + ) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_manage_files(self, _mock_token): + from singlestoredb.management.files import manage_files + + # The Files API is unchanged at v2, so both versions share one + # ``FilesManager`` class; the version shows up only in the base URL. + for ver in ('v1', 'v2'): + mgr = manage_files( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version=ver, + ) + self.assertTrue( + mgr._base_url.endswith(f'/{ver}/'), + f'expected base URL to end with /{ver}/, got {mgr._base_url}', + ) + + +class TestVersionNeutralHelpers(unittest.TestCase): + """ + ``get_organization``/``get_secret``/``get_stage`` exported from + ``singlestoredb.management`` follow ``management.version`` like the + factories do, instead of being the v1 implementations under a neutral name. + The version-locked ones remain reachable through the shim modules. + """ + + def test_top_level_names_are_the_neutral_ones(self): + import singlestoredb.management as m + self.assertEqual( + m.get_organization.__module__, + 'singlestoredb.management.organization', + ) + self.assertEqual( + m.get_secret.__module__, + 'singlestoredb.management.organization', + ) + self.assertEqual( + m.get_stage.__module__, 'singlestoredb.management.stage', + ) + + def test_shims_still_expose_their_own_version(self): + from singlestoredb.management import cluster, workspace + for name in ('get_organization', 'get_secret', 'get_stage'): + self.assertEqual( + getattr(workspace, name).__module__, + 'singlestoredb.management.v1.workspace', name, + ) + self.assertEqual( + getattr(cluster, name).__module__, + 'singlestoredb.management.v2.cluster', name, + ) + + def test_helpers_dispatch_on_the_option(self): + from singlestoredb.management import get_organization + from singlestoredb.management import get_secret + from singlestoredb.management import get_stage + calls = [] + for ver in ('v1', 'v2'): + with management_version(ver): + for name, call, expected in ( + ('get_organization', lambda: get_organization(), ()), + ('get_secret', lambda: get_secret('s'), ('s',)), + ('get_stage', lambda: get_stage('d'), ('d',)), + ): + target = f'singlestoredb.management.{ver}.{name}' + + def record(*args, _n=name, _v=ver): + calls.append((_v, _n, args)) + return 'ok' + + with patch(target, record): + self.assertEqual(call(), 'ok') + self.assertEqual(calls[-1], (ver, name, expected)) + self.assertEqual(len(calls), 6) + + def test_explicit_version_beats_the_option(self): + from singlestoredb.management import get_organization + with management_version('v1'): + with patch( + 'singlestoredb.management.v2.get_organization', + lambda: 'from-v2', + ): + self.assertEqual(get_organization(version='v2'), 'from-v2') + + def test_unknown_version_raises(self): + from singlestoredb.management import get_organization + with self.assertRaises(ManagementError) as ctx: + get_organization(version='v99') + self.assertIn('v99', str(ctx.exception)) + + def test_version_without_the_helper_raises(self): + from singlestoredb.management._version_import import _versioned_attr + with self.assertRaises(ManagementError) as ctx: + _versioned_attr('get_nothing', 'v1') + msg = str(ctx.exception) + self.assertIn('get_nothing', msg) + self.assertIn('v1', msg) + + +class TestManageWorkspacesDeprecation(unittest.TestCase): + """ + ``manage_workspaces()`` warns, but the internal v1-only path does not. + + Fusion, the UDF ``stage://`` handling and the AI helpers are v1-only by + design, so they go through ``_manage_workspaces_v1`` -- warning there would + be noise the caller can do nothing about. + """ + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_public_factory_warns(self, _mock_token): + from singlestoredb.management.workspace import manage_workspaces + with self.assertWarns(DeprecationWarning) as ctx: + # Pinned so the assertion is about the warning rather than about + # whatever version the ambient option happens to name. + manage_workspaces( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v1', + ) + self.assertIn('manage_clusters', str(ctx.warning)) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_internal_path_is_silent(self, _mock_token): + from singlestoredb.management.workspace import _manage_workspaces_v1 + with warnings.catch_warnings(): + warnings.simplefilter('error', DeprecationWarning) + _manage_workspaces_v1( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, + ) + + +class TestDeprecatedVersionWarning(unittest.TestCase): + """ + Every public version-neutral entry point warns when it resolves to v1. + + v1 is being wound down, so a caller who lands on it -- whether by passing + ``version='v1'`` or by inheriting it from the ``management.version`` + option -- has to be told. The warning fires after resolution rather than in + ``_resolve_version``, so both routes are covered and the internal v1-only + paths stay silent (see :class:`TestManageWorkspacesDeprecation`). + """ + + # (label, callable taking a version kwarg). Each is a public entry point + # that can resolve to v1; ``manage_clusters`` is absent because v1 has no + # clusters and it raises instead, and ``manage_workspaces`` because it + # raises its own more specific warning, asserted separately below. + def _entry_points(self): + import singlestoredb as s2 + from singlestoredb.management import get_organization + from singlestoredb.management import get_secret + from singlestoredb.management import get_stage + return [ + ( + 'manage_files', lambda **kw: s2.manage_files( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, **kw, + ), + ), + ( + 'manage_regions', lambda **kw: s2.manage_regions( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, **kw, + ), + ), + # The three helpers dispatch through _versioned_attr, so they are + # patched out: the assertion is about the warning, not the route. + ('get_organization', lambda **kw: get_organization(**kw)), + ('get_secret', lambda **kw: get_secret('s', **kw)), + ('get_stage', lambda **kw: get_stage('d', **kw)), + ] + + @contextlib.contextmanager + def _stubbed_helpers(self): + """Stub the three version-package helpers at both versions.""" + with contextlib.ExitStack() as stack: + for ver in ('v1', 'v2'): + for name in ('get_organization', 'get_secret', 'get_stage'): + stack.enter_context( + patch( + f'singlestoredb.management.{ver}.{name}', + lambda *a: 'ok', + create=True, + ), + ) + yield + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_explicit_v1_warns(self, _mock_token): + with self._stubbed_helpers(): + for label, call in self._entry_points(): + with self.subTest(entry_point=label): + with self.assertWarns(DeprecationWarning) as ctx: + call(version='v1') + msg = str(ctx.warning) + self.assertIn('v1', msg) + self.assertIn('deprecated', msg) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_v1_inherited_from_the_option_warns(self, _mock_token): + """A caller who never names a version still gets told.""" + with self._stubbed_helpers(), management_version('v1'): + for label, call in self._entry_points(): + with self.subTest(entry_point=label): + with self.assertWarns(DeprecationWarning) as ctx: + call() + self.assertIn('deprecated', str(ctx.warning)) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_v2_is_silent(self, _mock_token): + """The default version must not warn -- otherwise nobody reads any of them.""" + with self._stubbed_helpers(), management_version('v2'): + for label, call in self._entry_points() + [ + ( + 'manage_clusters', lambda **kw: __import__( + 'singlestoredb', + ).manage_clusters( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, **kw, + ), + ), + ]: + with self.subTest(entry_point=label): + with warnings.catch_warnings(): + warnings.simplefilter('error', DeprecationWarning) + call() + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_v1_still_works(self, _mock_token): + """ + Deprecated must not mean broken. This is the point of the whole set. + + v2 is the default, but v1 is still a supported version: every entry + point must return a working v1 object, and none may raise merely + because the default moved. Warnings are the only consequence. + """ + import singlestoredb as s2 + from singlestoredb.management.workspace import manage_workspaces + with self._stubbed_helpers(), warnings.catch_warnings(): + warnings.simplefilter('ignore', DeprecationWarning) + for label, call in self._entry_points(): + with self.subTest(entry_point=label): + self.assertIsNotNone(call(version='v1')) + # The v1 routes really are v1 routes, not v2 ones relabelled. + for label, factory in ( + ('manage_files', s2.manage_files), + ('manage_regions', s2.manage_regions), + ('manage_workspaces', manage_workspaces), + ): + with self.subTest(factory=label): + mgr = factory( + access_token=FAKE_TOKEN, + base_url=FAKE_BASE_URL, + version='v1', + ) + self.assertIn('/v1/', mgr._base_url) + + def test_the_deprecated_version_is_not_the_default(self): + """Guards the pair: whatever DEPRECATED_VERSION names cannot be the default.""" + from singlestoredb import config + from singlestoredb.management import _version_import as vi + self.assertNotEqual(vi.DEPRECATED_VERSION, vi.DEFAULT_VERSION) + self.assertEqual(vi.DEFAULT_VERSION, 'v2') + self.assertNotEqual( + config.get_default('management.version'), vi.DEPRECATED_VERSION, + ) + + @patch('singlestoredb.management.manager.get_token', return_value=FAKE_TOKEN) + def test_manage_workspaces_warns_once_not_twice(self, _mock_token): + """ + ``manage_workspaces()`` is the one v1 entry point with its own message. + + It reaches v1 through ``_manage_workspaces_v1``, which is deliberately + silent, so the caller gets exactly one warning -- the specific one + naming ``manage_clusters`` -- rather than that plus the generic + "v1 is deprecated". + """ + from singlestoredb.management.workspace import manage_workspaces + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + manage_workspaces( + access_token=FAKE_TOKEN, base_url=FAKE_BASE_URL, version='v1', + ) + deprecations = [ + w for w in caught if issubclass(w.category, DeprecationWarning) + ] + self.assertEqual(len(deprecations), 1, [str(w.message) for w in deprecations]) + self.assertIn('manage_clusters', str(deprecations[0].message)) + + +class TestV1IsDocumentedAsDeprecated(unittest.TestCase): + """ + Every module under ``management/v1/`` carries a deprecation note. + + A docstring check rather than a runtime one because most of these are + classes built by ``from_dict`` deep in the library, where a warning would + be noise the caller cannot act on. The note is what a reader of the API + docs and of an IDE tooltip actually sees. + """ + + #: ``inference/*`` has no v2 counterpart, so there is nowhere to send + #: callers and deprecating it would be a lie. Its docstring says so + #: explicitly, which the test below checks instead. + NOT_DEPRECATED = {'inference_api'} + + def _v1_modules(self): + import singlestoredb.management.v1 as v1 + directory = os.path.dirname(v1.__file__) + return sorted( + name[:-3] for name in os.listdir(directory) + if name.endswith('.py') and name != '__init__.py' + ) + + def test_every_v1_module_says_it_is_deprecated(self): + modules = self._v1_modules() + self.assertTrue(modules, 'found no modules under management/v1/') + for name in modules: + if name in self.NOT_DEPRECATED: + continue + with self.subTest(module=name): + mod = importlib.import_module(f'singlestoredb.management.v1.{name}') + self.assertIsNotNone(mod.__doc__, f'v1/{name}.py has no docstring') + self.assertIn('deprecated', mod.__doc__.lower()) + + def test_the_v1_package_itself_says_it_is_deprecated(self): + import singlestoredb.management.v1 as v1 + self.assertIn('deprecated', v1.__doc__.lower()) + + def test_the_workspace_shim_says_it_is_deprecated(self): + from singlestoredb.management import workspace + self.assertIn('deprecated', workspace.__doc__.lower()) + + def test_the_inference_api_explains_why_it_is_exempt(self): + """The exemption must be justified in the module, not just in this test.""" + from singlestoredb.management.v1 import inference_api + doc = inference_api.__doc__.lower() + self.assertIn('not** deprecated', doc) + # The reason: these routes are served nowhere else, so there is no + # replacement to send callers to. + self.assertIn('nowhere else', doc) + + def test_v1_only_classes_name_their_v2_replacement(self): + """ + The v1 classes that v2 genuinely replaced carry their own note. + + Restricted to classes actually defined under ``v1/``: the modules that + only re-export a shared implementation (``files``, ``region``, + ``billing_usage``) must *not* grow a class-level note, because that + note would show up on the v2 class too. + """ + from singlestoredb.management.v1 import export + from singlestoredb.management.v1 import job + from singlestoredb.management.v1 import organization + from singlestoredb.management.v1 import stage + from singlestoredb.management.v1 import workspace + expected = [ + (workspace.Workspace, 'cluster.Cluster'), + (workspace.WorkspaceGroup, 'cluster.Cluster'), + (workspace.StarterWorkspace, 'cluster.StarterCluster'), + (workspace.WorkspaceManager, 'cluster.ClusterManager'), + (stage.Stage, 'management.stage.Stage'), + (job.JobsManager, 'management.job.JobsManager'), + (organization.Organization, 'management.organization.Organization'), + (organization.Organizations, 'management.organization.Organizations'), + (export.ExportService, 'management.export.ExportService'), + (export.ExportStatus, 'management.export.ExportStatus'), + ] + for cls, replacement in expected: + with self.subTest(cls=cls.__name__): + doc = cls.__doc__ or '' + self.assertIn('.. deprecated::', doc) + self.assertIn(replacement, doc) + + def test_shared_classes_are_not_marked_deprecated(self): + """ + ``v1/files.py`` and friends re-export the shared classes. + + Marking those classes deprecated would tell v2 users their own classes + are going away, so only the v1 *module path* carries the note. + """ + from singlestoredb.management.v1 import billing_usage as v1_billing + from singlestoredb.management.v1 import files as v1_files + from singlestoredb.management.v1 import region as v1_region + for mod, names in ( + (v1_files, ('FilesManager', 'FilesObject')), + (v1_region, ('Region', 'RegionManager')), + (v1_billing, ('BillingUsageItem', 'UsageItem')), + ): + for name in names: + with self.subTest(cls=f'{mod.__name__}.{name}'): + cls = getattr(mod, name) + self.assertNotIn('.. deprecated::', cls.__doc__ or '') + # ...and it really is the shared class, not a v1 subclass. + self.assertFalse( + cls.__module__.startswith('singlestoredb.management.v1'), + f'{name} is defined under v1/, so the note above ' + 'would be correct and this test is wrong', + ) + + +class TestFactoriesAreNotDuplicated(unittest.TestCase): + """ + The ``manage_*`` factories must live in exactly one place. + + They are version-neutral -- they take ``version`` as an argument and + dispatch -- so duplicating them into ``v1/`` (as an earlier layout did) + both invites the two copies to drift and makes ``v1/`` un-deletable. + """ + + def test_factories_defined_only_at_top_level(self): + factories = { + 'manage_files': 'files', + 'manage_regions': 'region', + 'manage_workspaces': 'workspace', + 'manage_clusters': 'cluster', + } + for func, mod_name in factories.items(): + shared = importlib.import_module(f'singlestoredb.management.{mod_name}') + self.assertTrue( + callable(getattr(shared, func, None)), + f'{func} should be defined in management/{mod_name}.py', + ) + for ver in ('v1', 'v2'): + try: + mod = importlib.import_module( + f'singlestoredb.management.{ver}.{mod_name}', + ) + except ModuleNotFoundError: + # Not every resource exists at every version; e.g. there + # is no v2 ``workspace`` module. + continue + self.assertNotIn( + func, vars(mod), + f'{func} must not be duplicated into ' + f'management/{ver}/{mod_name}.py', + ) + + +class TestVersionPackagesAreIndependent(unittest.TestCase): + """ + Guard the invariant that makes either version package removable. + + The v1 endpoints will eventually be abandoned, at which point + ``management/v1/`` should be deletable by ``rm -rf`` plus removal of the + back-compat shims. That only holds while ``management/v1/`` and + ``management/v2/`` do not import each other, in either direction: + version-neutral code belongs in the shared top-level ``management/`` + modules, which both version packages import sideways. + + If this test fails, the fix is to move the shared code up to + ``management/`` -- not to add a cross-version import. + """ + + def test_version_packages_extend_the_shared_base(self): + """ + Inheritance runs shared base -> version subclass, never v1 -> v2. + + ``Organization`` is the representative case: the base carries the v2 + behavior and ``v1/`` holds the backward override -- repointing the job + and inference sub-managers -- so ``v2/`` is a plain re-export. + + ``RegionManager`` used to play this role, but no longer can: once + ``regions/sharedtier`` was found to answer at both versions the v1 + override collapsed into a re-export, making ``V1 is V2 is Base`` and + the ``issubclass(V2, V1)`` assertion vacuously wrong. + """ + from singlestoredb.management.organization import ( + Organization as Base, + ) + from singlestoredb.management.v1.organization import ( + Organization as V1, + ) + from singlestoredb.management.v2.organization import ( + Organization as V2, + ) + self.assertTrue(issubclass(V1, Base)) + self.assertIsNot(V1, Base) + self.assertIs(V2, Base) + self.assertFalse(issubclass(V2, V1)) + + def test_a_version_package_that_only_re_exports_shares_the_base(self): + """ + A version with no behavioral difference re-exports, not subclasses. + + Both ``region`` modules are now pure re-exports, so all three names + are the same object. Asserted explicitly so that reintroducing a + subclass on one side has to be a deliberate edit to this test. + """ + from singlestoredb.management.region import RegionManager as Base + from singlestoredb.management.v1.region import RegionManager as V1 + from singlestoredb.management.v2.region import RegionManager as V2 + self.assertIs(V1, Base) + self.assertIs(V2, Base) + + def _module_paths(self, version): + pkg = importlib.import_module(f'singlestoredb.management.{version}') + pkg_dir = os.path.dirname(pkg.__file__) + return sorted( + os.path.join(pkg_dir, f) + for f in os.listdir(pkg_dir) + if f.endswith('.py') + ) + + def _cross_version_imports(self, version, other): + """Return every import of ``other`` found in ``version``'s modules.""" + offenders = [] + for path in self._module_paths(version): + with open(path) as f: + tree = ast.parse(f.read(), filename=path) + for node in ast.walk(tree): + # Relative ``from ..v2.x import y`` shows up as level=2 with + # module='v2.x'; absolute imports show up with the full path. + if isinstance(node, ast.ImportFrom): + mod = node.module or '' + if mod == other or mod.startswith(f'{other}.') or \ + f'management.{other}' in mod: + offenders.append( + f'{os.path.basename(path)}:{node.lineno}: ' + f'from {"." * node.level}{mod}', + ) + elif isinstance(node, ast.Import): + for alias in node.names: + if f'management.{other}' in alias.name: + offenders.append( + f'{os.path.basename(path)}:{node.lineno}: ' + f'import {alias.name}', + ) + return offenders + + def test_no_v2_module_imports_from_v1(self): + """No module under management/v2/ may import from management/v1/.""" + offenders = self._cross_version_imports('v2', 'v1') + self.assertEqual( + offenders, [], + 'management/v2/ must not import from management/v1/; move the ' + 'shared code up to management/ instead:\n ' + + '\n '.join(offenders), + ) + + def test_no_v1_module_imports_from_v2(self): + """No module under management/v1/ may import from management/v2/.""" + offenders = self._cross_version_imports('v1', 'v2') + self.assertEqual( + offenders, [], + 'management/v1/ must not import from management/v2/; move the ' + 'shared code up to management/ instead:\n ' + + '\n '.join(offenders), + ) + + def _assert_imports_survive_removal(self, version, other): + """Import every module of ``version`` with ``other`` blocked.""" + names = [ + f'singlestoredb.management.{version}.' + os.path.basename(p)[:-3] + for p in self._module_paths(version) + if not os.path.basename(p).startswith('__') + ] + blocked_prefix = f'singlestoredb.management.{other}' + + # Drop anything already imported so the blocker actually gets + # consulted, then forbid the other version package outright. + saved = { + k: v for k, v in sys.modules.items() + if k.startswith(blocked_prefix) or k in names + } + for k in saved: + del sys.modules[k] + + class _Blocker: + def find_module(self, fullname, path=None): + return self.find_spec(fullname, path) + + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith(blocked_prefix): + raise AssertionError( + f'{version} import chain reached {fullname}; ' + f'{other} is supposed to be removable', + ) + return None + + blocker = _Blocker() + sys.meta_path.insert(0, blocker) + try: + for name in names: + importlib.import_module(name) + finally: + sys.meta_path.remove(blocker) + sys.modules.update(saved) + + def test_v2_imports_survive_v1_removal(self): + """Importing every v2 module works with management.v1 blocked.""" + self._assert_imports_survive_removal('v2', 'v1') + + def test_v1_imports_survive_v2_removal(self): + """Importing every v1 module works with management.v2 blocked.""" + self._assert_imports_survive_removal('v1', 'v2') + + +if __name__ == '__main__': + unittest.main() diff --git a/singlestoredb/tests/utils.py b/singlestoredb/tests/utils.py index c7cba6808..94d754c47 100644 --- a/singlestoredb/tests/utils.py +++ b/singlestoredb/tests/utils.py @@ -4,16 +4,24 @@ import glob import logging import os +import random import re +import secrets +import unittest import uuid +from types import SimpleNamespace from typing import Any from typing import Dict from typing import List +from typing import Optional from typing import Tuple +from unittest import mock from urllib.parse import urlparse import singlestoredb as s2 from singlestoredb.connection import build_params +from singlestoredb.exceptions import ManagementError +from singlestoredb.management.v2.cluster import ClusterManager as _ClusterManager logger = logging.getLogger(__name__) @@ -276,3 +284,963 @@ def drop_user(name: str) -> None: with s2.connect(**args) as conn: with conn.cursor() as cur: cur.execute(f'DROP USER IF EXISTS {name};') + + +# +# Live deployment tracking +# +# Every workspace group, workspace, cluster and starter cluster a test creates +# costs money until it is terminated, and the usual `tearDownClass` is not +# enough on its own: +# +# * unittest does not call `tearDownClass` at all if `setUpClass` raises, so +# a fixture that dies partway through -- two of three clusters created, +# then a dropped connection -- leaks everything it had made so far; +# * a test that creates a deployment in its body and then fails before its +# own cleanup line leaks it too. +# +# So creations are registered here as well, and `cleanup_tracked()` sweeps +# whatever is left: per test class as the run moves on to the next one, and +# again for everything at the end of the session (see conftest.py). +# Terminating twice is harmless -- the second attempt finds it gone and is +# ignored -- so tracked objects do not have to be untracked by the tests that +# clean up after themselves. +# + +#: (owner, label, object) for every deployment created so far and not yet +#: swept. The owner is the test class that was running at creation time, so +#: a class's leftovers can be dropped when the run leaves that class rather +#: than idling -- and billing -- until the session ends. +_tracked: List[Tuple[str, str, Any]] = [] + +#: (receiver, finder, args, kwargs) for every creation call currently +#: executing. A creator POSTs and only then waits for the deployment to come +#: up, so for the whole ``wait_on_active`` window -- twenty minutes for a +#: cluster -- something billable exists that nothing has registered yet: +#: ``_tracking_wrapper`` tracks on return and recovers in its ``except``, and +#: neither runs if the process is killed. See :func:`recover_in_flight`. +_in_flight: List[Tuple[Any, Any, Tuple[Any, ...], Dict[str, Any]]] = [] + +#: Test class currently running, as set by conftest. +_owner = '' + + +def get_owner() -> str: + """Return the test class creations are currently attributed to.""" + return _owner + + +def set_owner(owner: str) -> None: + """Record which test class subsequent creations belong to.""" + global _owner + _owner = owner + + +def _is_mocked(obj: Any) -> bool: + """ + Did this object come out of a mocked manager? + + The unit tests call the same creation methods with ``_post`` patched, and + the objects they get back name deployments that do not exist. Sweeping + those would be a round trip per fake object and a warning apiece. + + An unrecognisable object counts as real, including one whose ``_manager`` + is ``None``: a fake deployment swept is a round trip and a warning, whereas + a real one skipped is a cluster left running and billing. That bias lives + in :func:`_creator_is_mocked`, which this defers to for everything but the + receiver itself. + """ + from unittest.mock import NonCallableMock + + if isinstance(obj, NonCallableMock): + return True + manager = getattr(obj, '_manager', None) + if isinstance(manager, NonCallableMock): + return True + return _creator_is_mocked(manager) + + +def track(obj: Any, label: str = '') -> Any: + """ + Register a live deployment for end-of-session cleanup. + + Returns the object, so it can wrap a creation call in place:: + + cls.cluster = utils.track(mgr.create_cluster(...)) + + """ + if obj is not None and not _is_mocked(obj): + _tracked.append(( + _owner, + label or '{} {!r}'.format( + type(obj).__name__, getattr(obj, 'name', None) or + getattr(obj, 'id', '?'), + ), + obj, + )) + return obj + + +def _recover_orphan( + receiver: Any, + finder: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> None: + """ + Track the deployment a *failed* creation call left running. + + Every creator brings the deployment into existence and only then waits for + it: ``create_cluster`` has its ``get_cluster`` before ``_wait_on_state`` + (``management/v2/cluster.py:1426``). So a wait that times out, hits a + transient error, or is interrupted raises *after* the server has a live, + billable deployment -- and since tracking wraps the return value, nothing + is ever registered. That leak is silent: no per-class sweep, no + end-of-session sweep, and no mention in the summary. + + The name is the first argument to every creator, so the orphan can be + found by listing and matching on it. Failures here are logged, not raised: + this runs while another exception is propagating, and replacing the + caller's error with a cleanup error would hide the real failure. + """ + name = kwargs.get('name') or (args[0] if args else None) + if not isinstance(name, str): + return + + try: + for obj in finder(receiver): + if getattr(obj, 'name', None) != name: + continue + track( + obj, + '{} {!r} (left behind by a failed create)'.format( + type(obj).__name__, name, + ), + ) + return + except Exception as exc: + logger.warning( + f'Could not look for a deployment named {name!r} left behind by ' + f'a failed create; it may still be running: {exc}', + ) + + +def untrack(obj: Any) -> None: + """Forget a deployment that has been terminated.""" + for i, entry in reversed(list(enumerate(_tracked))): + if entry[2] is obj: + _tracked.pop(i) + + +def terminate(obj: Any) -> None: + """ + Terminate a deployment, whatever kind it is. + + ``force=True`` is what makes a workspace group with live workspaces in it + go away; the starter variants take no arguments at all. + """ + try: + obj.terminate(force=True) + except TypeError: + obj.terminate() + + +def _creator_is_mocked(target: Any) -> bool: + """ + Is this creation call going through a mocked manager? + + The unit tests call the creation methods with ``_post`` patched, and the + objects they get back name deployments that do not exist, so they must not + be tracked. ``target`` is the creation call's receiver -- a manager, or the + ``WorkspaceGroup`` of ``WorkspaceGroup.create_workspace`` -- or, through + :func:`_is_mocked`, whatever a created object holds in ``_manager``. + """ + from unittest.mock import NonCallableMock + + if isinstance(target, NonCallableMock): + return True + manager = target if hasattr(target, '_post') else getattr( + target, '_manager', None, + ) + if isinstance(manager, NonCallableMock): + return True + # An unrecognisable receiver counts as real: a fake deployment swept is a + # round trip and a warning, whereas a real one skipped is a cluster left + # running and billing. + return any( + isinstance(getattr(manager, x, None), NonCallableMock) + for x in ('_get', '_post', '_delete') + ) + + +#: (module, class, method, finder) tuples for the calls that bring a billable +#: deployment into existence. Wrapping them is what makes tracking automatic, +#: so a new test cannot leak a cluster by forgetting to register it. +#: +#: ``finder`` takes the receiver -- the manager, or the group for +#: ``WorkspaceGroup.create_workspace`` -- and returns the collection to search +#: for a deployment the call created but did not return. See +#: :func:`_recover_orphan`. +_CREATORS = [ + ( + 'singlestoredb.management.v1.workspace', 'WorkspaceManager', + 'create_workspace_group', + lambda recv: recv.workspace_groups, + ), + ( + 'singlestoredb.management.v1.workspace', 'WorkspaceManager', + 'create_workspace', + # WorkspaceManager has no `workspaces` of its own, so the search goes + # group by group. Only ever walked on the failure path. + lambda recv: [w for g in recv.workspace_groups for w in g.workspaces], + ), + ( + 'singlestoredb.management.v1.workspace', 'WorkspaceManager', + 'create_starter_workspace', + lambda recv: recv.starter_workspaces, + ), + ( + 'singlestoredb.management.v1.workspace', 'WorkspaceGroup', + 'create_workspace', + lambda recv: recv.workspaces, + ), + ( + 'singlestoredb.management.v2.cluster', 'ClusterManager', + 'create_cluster', + lambda recv: recv.clusters, + ), + ( + 'singlestoredb.management.v2.cluster', 'ClusterManager', + 'create_starter_cluster', + lambda recv: recv.starter_clusters, + ), +] + +_tracking_installed = False + + +def _tracking_wrapper(func: Any, finder: Any) -> Any: + """ + Wrap a creation method so its result -- or its orphan -- gets tracked. + + On success the returned deployment is registered. On failure the + deployment the call already brought into existence is looked up and + registered instead; see :func:`_recover_orphan` for why one exists. + + The call is also listed in ``_in_flight`` for its duration, so a sweep + that runs while it is still waiting -- SIGTERM, atexit -- can recover the + orphan itself rather than being killed before the ``except`` below gets to + (see :func:`recover_in_flight`). + + ``_creator_is_mocked``, not ``_is_mocked``: the receiver is the manager (or + the workspace group), and ``_is_mocked`` looks for a ``_manager`` + attribute, which a manager does not have -- so a real manager with a + patched ``_post`` would read as live and the recovery would fire a real + API call from a unit test. ``_creator_is_mocked`` inspects the receiver's + own transport and handles both receiver shapes. + + That same verdict also decides whether the *result* is tracked, rather than + leaving it to ``track()``. ``track()`` can only judge what it is handed, + and it is deliberately biased toward "real" for anything it cannot place -- + including an object whose ``_manager`` is ``None``, which is exactly what a + unit test's stubbed ``get_cluster`` returns. Nothing a mocked creator + returns names a deployment that exists, so the receiver's verdict is the + authoritative one and it is the one used here. + """ + import functools + + @functools.wraps(func) + def wrapper(receiver: Any, *args: Any, **kwargs: Any) -> Any: + mocked = _creator_is_mocked(receiver) + entry = (receiver, finder, args, kwargs) + if not mocked: + _in_flight.append(entry) + try: + out = func(receiver, *args, **kwargs) + return out if mocked else track(out) + except BaseException: + # BaseException, not Exception: a KeyboardInterrupt during the + # twenty-minute wait_on_active wait leaves the same live + # deployment behind as a timeout does. + # + # Only if the entry is still listed: claiming it is what keeps this + # from tracking the orphan a second time when a sweep already + # recovered it mid-wait and then let the call unwind. + if not mocked and _drop_in_flight(entry): + _recover_orphan(receiver, finder, args, kwargs) + raise + finally: + if not mocked: + _drop_in_flight(entry) + + return wrapper + + +def _drop_in_flight(entry: Tuple[Any, Any, Tuple[Any, ...], Any]) -> bool: + """ + Remove one in-flight entry, and say whether it was still there. + + By identity, and only this entry: two creations with equal arguments -- + a retried create, say -- would otherwise pop each other's. + """ + for i, other in reversed(list(enumerate(_in_flight))): + if other is entry: + _in_flight.pop(i) + return True + return False + + +def recover_in_flight() -> None: + """ + Track the deployments that creation calls still in progress have created. + + A creator POSTs, then waits. Everything that registers a deployment runs + after that wait -- ``track()`` on return, ``_recover_orphan()`` in the + wrapper's ``except`` -- so a sweep triggered from outside the call, by + SIGTERM or atexit, sees nothing in ``_tracked`` and the deployment is left + running. A cancelled CI job during a shared-pool build is exactly that + case. + + So each in-flight call is looked up the same way a failed one is, putting + whatever the server already created into ``_tracked`` before the sweep + walks it. Entries are popped as they are drained, so a handler that + returns and lets the wrapper's own ``except`` run cannot recover twice. + Never raises: ``_recover_orphan`` logs its own failures, and this runs on + the way out. + """ + while _in_flight: + receiver, finder, args, kwargs = _in_flight.pop() + _recover_orphan(receiver, finder, args, kwargs) + + +def install_deployment_tracking() -> None: + """ + Wrap the deployment creation methods so their results are tracked. + + Called from ``pytest_configure`` rather than a fixture: it has to be in + place before any test module is imported, since a ``setUpClass`` can run + creations that a later fixture would never see. + """ + global _tracking_installed + if _tracking_installed: + return + _tracking_installed = True + + import importlib + + for module_name, class_name, method_name, finder in _CREATORS: + try: + klass = getattr(importlib.import_module(module_name), class_name) + setattr( + klass, method_name, + _tracking_wrapper(getattr(klass, method_name), finder), + ) + except AttributeError as exc: + # A renamed method must not silently stop being tracked. + logger.warning( + f'Cannot track {module_name}.{class_name}.' + f'{method_name}: {exc}', + ) + + +def _is_gone(obj: Any) -> bool: + """ + Has this deployment already been terminated? + + The local copy is stale -- a test that terminated in its own teardown + still holds an object whose ``terminated_at`` is None -- so ask the + server. + + Only a 404 counts as gone. Any other refresh failure reports "still + there": answering "gone" on a transient 5xx or a dropped connection + skips the termination below, and a cluster left running costs money, + whereas a redundant terminate on something already gone is one wasted + round trip. + """ + if hasattr(obj, 'refresh'): + try: + obj.refresh() + except ManagementError as exc: + if exc.errno == 404: + return True + logger.warning( + f'Could not refresh {obj!r} to see whether it is already ' + f'gone; assuming it is still live: {exc}', + ) + return False + except Exception as exc: + logger.warning( + f'Could not refresh {obj!r} to see whether it is already ' + f'gone; assuming it is still live: {exc}', + ) + return False + if getattr(obj, 'terminated_at', None) is not None: + return True + return str(getattr(obj, 'state', '') or '').upper() in ( + 'TERMINATED', 'TERMINATING', + ) + + +def cleanup_tracked(owner: Optional[str] = None) -> List[str]: + """ + Terminate tracked deployments that are still live. + + Parameters + ---------- + owner : str, optional + Only sweep what this test class created. The default sweeps + everything, which is what the end of the session wants. + + Returns + ------- + List[str] + Labels of the deployments this call terminated. Failures are logged + rather than raised: this runs outside any test, where an exception + would be reported against whatever happens to run next. + + """ + # Last created, first terminated: a workspace goes before the group that + # holds it. + entries = [x for x in reversed(_tracked) if owner is None or x[0] == owner] + + removed = [] + for entry in entries: + _, label, obj = entry + if _is_gone(obj): + _tracked.remove(entry) + continue + try: + terminate(obj) + except Exception as exc: + # Deliberately left in ``_tracked``, so the end-of-session sweep + # tries again. Dropping the entry first -- as this used to -- meant + # one transient error was enough to leak the deployment for good, + # and it did not even appear in the summary below. + logger.warning(f'Could not terminate {label}: {exc}') + else: + _tracked.remove(entry) + removed.append(label) + return removed + + +def tracked_labels() -> List[str]: + """ + Labels of every deployment still tracked, i.e. not yet swept. + + After the end-of-session sweep this should be empty; anything left is a + deployment that is still live and still costing money, so conftest + reports it rather than letting the run end quietly. + """ + return [label for _, label, _ in _tracked] + + +# +# Shared deployment pool +# +# Several classes need nothing from a deployment but that it is live: the +# Stage and Job suites read and write through the management API against +# whatever cluster they are handed. Deploying one apiece cost 2190s of the +# 8915s a traced run took, and an S-00 cluster reaching ACTIVE is ~460s that +# cannot be made faster -- so the only lever is deploying fewer of them. +# +# The pool is built on first use and reused for the rest of the process. A +# class must not mutate what it borrows, so anything whose subject *is* the +# deployment keeps deploying its own: ``TestCluster`` and ``TestWorkspace`` +# (``test_update`` PATCHes the cluster and cycles it back through PENDING), +# ``TestClusterFusionCreateDrop`` and ``TestClusterFusionSuspendResume``. So +# does ``TestWorkspaceFusion``, whose workspace groups are the subject of its +# ``SHOW WORKSPACE GROUPS`` assertions and cost 40s to deploy unwaited anyway. +# +# What makes the four borrowers safe is that each scopes its assertions to +# itself: every Stage path is namespaced with the class's ``cls.id``, job +# listings filter by job id rather than listing a deployment's jobs, and none +# of them asserts a row count over an org-wide listing. +# +# The pool is process-wide, so under ``pytest-xdist`` every worker that gets a +# borrowing class builds a pool of its own. The ``xdist_group`` marks below +# keep the borrowers together on a worker; see ``SHARED_CLUSTER_*_GROUP``. +# + +#: ``xdist_group`` names for the classes that borrow from the pool, so +#: ``--dist loadgroup`` puts each set on one worker and each set builds one +#: pool. Two groups rather than one: a single group serialises all four classes +#: behind one pool build, and the groups run concurrently on separate workers, +#: so splitting costs one extra cluster and halves that chain. +#: +#: Stage wants two clusters (``TestStageFusion`` names a second one in +#: ``IN GROUP``) and jobs want one, so the split follows what they borrow: +#: +#: * ``SHARED_CLUSTER_STAGE_GROUP`` -- ``TestStageFusion``, v2 ``TestStage`` +#: * ``SHARED_CLUSTER_JOBS_GROUP`` -- ``TestJobsFusion``, v2 ``TestJob`` +#: +#: Without ``-n``/``--dist loadgroup`` the marks do nothing: one process, one +#: pool of two, which is the serial behaviour they were added on top of. +SHARED_CLUSTER_STAGE_GROUP = 'shared-cluster-stage' +SHARED_CLUSTER_JOBS_GROUP = 'shared-cluster-jobs' + +#: Live clusters shared by the classes that need only *a* deployment. +_pool: List[Any] = [] + +#: Why the pool cannot be built in this organization, once that is known. +#: Cached so the second class to ask skips without repeating the lookups. +_pool_skip: Optional[str] = None + +#: Suffix for the pool's cluster names, so a run's clusters are distinguishable +#: from a concurrent run's. Matches the ``cl-test-*`` pattern the maintenance +#: sweep in ``cleanup_deployments.py`` looks for. +_pool_id = secrets.token_hex(4) + + +def shared_clusters(count: int = 1) -> List[Any]: + """ + Return ``count`` live v2 clusters shared by the whole test session. + + The pool grows to fit the largest request and is never rebuilt, so every + caller gets the same objects:: + + @classmethod + def setUpClass(cls): + cls.cluster, cls.cluster_2 = utils.shared_clusters(2) + + Raises ``unittest.SkipTest`` for the same reasons the per-class fixtures + did -- no US regions, or no project to deploy into -- so a class that + borrows from the pool skips where it used to skip. + + Terminating a pool cluster is not this module's business beyond the + end-of-session sweep: a class that borrows one must leave it live and + usable, since the classes after it get the same object. + """ + global _pool_skip + + if _pool_skip: + raise unittest.SkipTest(_pool_skip) + + if len(_pool) >= count: + return _pool[:count] + + # Pinned to v2: the pool's consumers are v2 suites, so the fixture must + # not follow the management.version option out of v2 either. + mgr = s2.manage_clusters(version='v2') + + us_regions = [ + x for x in mgr.regions + if 'US' in x.name or 'us-' in (x.region_name or '') + ] + if not us_regions: + _pool_skip = 'No US regions reported by the v2 API' + raise unittest.SkipTest(_pool_skip) + + project_id = os.environ.get('SINGLESTOREDB_TEST_PROJECT') + if not project_id: + standard = [x for x in mgr.projects if x.edition == 'STANDARD'] + if not standard: + _pool_skip = ( + 'No STANDARD project in this organization; set ' + 'SINGLESTOREDB_TEST_PROJECT to the project to deploy into' + ) + raise unittest.SkipTest(_pool_skip) + project_id = standard[0].id + + # Tracked under the empty owner rather than under whichever class happened + # to ask first. conftest.pytest_runtest_setup sweeps the previous owner's + # deployments as soon as the run moves to the next class, so a pool + # attributed to a class would be terminated after its first consumer; + # ``''`` matches no per-class sweep and is swept exactly once, by + # pytest_unconfigure, which passes owner=None and so matches everything. + prev = get_owner() + set_owner('') + try: + while len(_pool) < count: + _pool.append( + mgr.create_cluster( + f'cl-test-shared-{len(_pool)}-{_pool_id}', + region=random.choice(us_regions), + size='S-00', + # The v2 suites that deploy their own ask for this, and a + # pool cluster stands in for those, so it has to be at + # least as reachable as what it replaces. + firewall_ranges=['0.0.0.0/0'], + project=project_id, + wait_on_active=True, + wait_timeout=1200, + ), + ) + finally: + set_owner(prev) + + return _pool[:count] + + +class CountingManager: + """ + Stand-in for a :class:`Manager` that records every request. + + Enough of the management API's filesystem behaviour is simulated for a + :class:`Stage` or :class:`FileSpace` to be driven end to end without a + deployment: paths listed in ``existing`` answer metadata requests, and + anything else raises the 404 ``ManagementError`` that ``exists`` reads. + Writes and deletes update that set, so a sequence of operations sees the + effect of the ones before it. + + ``calls`` holds one ``(method, path)`` pair per request, in order, which + is what makes the round-trip count of an operation assertable. The paths + are the remote path the caller asked for, with the route prefix + (``clusters//stage/fs/``, ``files/fs//``) and any query string + removed, so the same expectations read the same for Stage and for a file + space. + + Parameters + ---------- + existing : iterable of str, optional + Remote paths that already exist. A path ending in ``/`` is a folder. + + """ + + def __init__(self, existing: Any = ()): + self.existing = {self._key(x) for x in existing} + self.calls: List[Tuple[str, str]] = [] + + @staticmethod + def _key(path: Any) -> str: + """Reduce a request path to the remote path it addresses.""" + path = str(path).split('?')[0] + # 'files/fs//' for a file space, '/fs/' + # for a Stage at either version + path = re.sub(r'^files/fs/[^/]+/', r'', path) + path = re.split(r'/fs/', path, maxsplit=1)[-1] + # A trailing '/' marks a folder, but the routes collapse runs of them + return re.sub(r'/+$', r'/', path).lstrip('/') + + def _response(self, key: str) -> Any: + """Return a metadata response for an existing path.""" + is_dir = key.endswith('/') + return SimpleNamespace( + json=lambda: dict( + name=key.rstrip('/').rsplit('/', 1)[-1], + path=key, + size=0 if is_dir else 8, + type='directory' if is_dir else 'file', + format='', + mimetype='' if is_dir else 'text/plain', + writable=True, + content=[] if is_dir else None, + ), + content=b'' if is_dir else b'contents', + ) + + def _get(self, path: Any, params: Any = None, **kwargs: Any) -> Any: + key = self._key(path) + self.calls.append(('GET', key)) + if key not in self.existing: + # A folder resolves whether or not the caller asked for it with a + # trailing '/', the way the routes behave + if not key.endswith('/') and f'{key}/' in self.existing: + return self._response(f'{key}/') + raise ManagementError(errno=404, msg=f'path does not exist: {key}') + return self._response(key) + + def _put(self, path: Any, **kwargs: Any) -> Any: + key = self._key(path) + if 'isFile=false' in str(path): + key = re.sub(r'/*$', r'/', key) + self.calls.append(('PUT', key)) + self.existing.add(key) + return SimpleNamespace( + json=lambda: dict(name=key.rsplit('/', 1)[-1], path=key), + content=b'', + ) + + def _patch(self, path: Any, json: Any = None, **kwargs: Any) -> Any: + key = self._key(path) + self.calls.append(('PATCH', key)) + self.existing.discard(key) + self.existing.add(self._key((json or {}).get('newPath', key))) + return SimpleNamespace(json=lambda: {}, content=b'') + + def _delete(self, path: Any, **kwargs: Any) -> Any: + key = self._key(path) + self.calls.append(('DELETE', key)) + self.existing.discard(key) + return SimpleNamespace(json=lambda: {}, content=b'') + + def counts(self) -> Dict[str, int]: + """Return the number of recorded requests per method.""" + out: Dict[str, int] = {} + for method, _ in self.calls: + out[method] = out.get(method, 0) + 1 + return out + + +def counting_stage(existing: Any = (), stage_cls: Any = None) -> Tuple[Any, Any]: + """ + Return a ``(Stage, CountingManager)`` pair wired to no deployment. + + Parameters + ---------- + existing : iterable of str, optional + Stage paths that already exist + stage_cls : type, optional + ``Stage`` class to instantiate. Defaults to the version-neutral one; + pass ``v1.stage.Stage`` to exercise the v1 route prefix, which the + recorded paths have stripped either way. + + """ + if stage_cls is None: + from singlestoredb.management.stage import Stage as stage_cls + manager = CountingManager(existing) + stage = stage_cls.__new__(stage_cls) + stage._deployment_id = 'deployment-id' + stage._manager = manager + return stage, manager + + +def counting_file_space(existing: Any = ()) -> Tuple[Any, Any]: + """ + Return a ``(FileSpace, CountingManager)`` pair wired to no organization. + + Parameters + ---------- + existing : iterable of str, optional + File paths that already exist + + """ + from singlestoredb.management.files import FileSpace + manager = CountingManager(existing) + space = FileSpace.__new__(FileSpace) + space._location = 'personal' + space._manager = manager + return space, manager + + +#: IDs for the fixtures :func:`counting_cluster_manager` builds by default. +COUNTING_CLUSTER_NAME = 'counting-cluster' +COUNTING_CLUSTER_ID = 'ffffffff-0000-0000-0000-000000000001' +COUNTING_PROJECT_ID = 'ffffffff-0000-0000-0000-000000000002' + + +def cluster_payload( + name: str, + id: str, + project_id: Optional[str] = None, + region: Optional[str] = None, + **extra: Any, +) -> Dict[str, Any]: + """ + Return one item of a ``GET /v2/clusters`` response. + + ``region`` is omitted unless asked for, so a caller can say which of the + lazy properties it is exercising: reading ``Cluster.region`` resolves the + name against ``ClusterManager.regions``, which is a request, and only a + payload carrying a region has anything to resolve. + + Parameters + ---------- + name : str + Name of the cluster + id : str + Cluster ID + project_id : str, optional + Value for ``projectID`` + region : str, optional + Value for ``region``, the provider region name + **extra : keyword arguments, optional + Further response keys, in the API's own spelling + + """ + out: Dict[str, Any] = dict( + name=name, clusterID=id, state='ACTIVE', + sizeConfig=dict(size='S-00', scaleFactor=1.0), + ) + if project_id is not None: + out['projectID'] = project_id + if region is not None: + out['region'] = region + out.update(extra) + return out + + +def project_payload( + id: str, + name: str, + edition: str = 'STANDARD', +) -> Dict[str, Any]: + """Return one item of a ``GET /v2/projects`` response.""" + return dict(projectID=id, name=name, edition=edition) + + +class CountingClusterManager(_ClusterManager): + """ + A :class:`ClusterManager` that answers from fixtures and records requests. + + This is :class:`CountingManager` widened to a whole Fusion statement: the + management routes a statement resolves its deployment through + (``clusters``, ``clusters/``, ``projects``, ``regions``, + ``sharedtier/virtualClusters``) are served from the lists given here, and + Stage's own filesystem routes are delegated to a :class:`CountingManager` + sharing this object's ``calls`` list, so one ordered record covers both. + + Any other route raises, so a request nobody accounted for cannot slip + through as a mock's default return value. + + Parameters + ---------- + clusters : list of dict, optional + ``GET /v2/clusters`` items; see :func:`cluster_payload` + projects : list of dict, optional + ``GET /v2/projects`` items; see :func:`project_payload` + starter_clusters : list of dict, optional + ``GET /v2/sharedtier/virtualClusters`` items + regions : list of dict, optional + ``GET /v2/regions`` items + existing : iterable of str, optional + Stage paths that already exist; a path ending in ``/`` is a folder + + """ + + def __init__( + self, + clusters: Any = None, + projects: Any = None, + starter_clusters: Any = (), + regions: Any = (), + existing: Any = (), + ): + # Deliberately not calling ClusterManager.__init__: it wants an access + # token and a base URL, and nothing here makes a request. + if clusters is None: + clusters = [ + cluster_payload( + COUNTING_CLUSTER_NAME, COUNTING_CLUSTER_ID, + project_id=COUNTING_PROJECT_ID, + ), + ] + if projects is None: + projects = [project_payload(COUNTING_PROJECT_ID, 'Test Project')] + + self._cluster_payloads = list(clusters) + self._project_payloads = list(projects) + self._starter_cluster_payloads = list(starter_clusters) + self._region_payloads = list(regions) + + #: Serves the Stage filesystem routes + self.files = CountingManager(existing) + + #: One ``(method, path)`` pair per request, in order + self.calls = self.files.calls + + def _get(self, path: Any, params: Any = None, **kwargs: Any) -> Any: + if '/fs/' in str(path): + return self.files._get(path, params=params, **kwargs) + + key = str(path).split('?')[0] + self.calls.append(('GET', key)) + + if key == 'clusters': + return SimpleNamespace(json=lambda: self._cluster_payloads) + if key == 'projects': + return SimpleNamespace(json=lambda: self._project_payloads) + if key == 'regions': + return SimpleNamespace(json=lambda: self._region_payloads) + if key == 'sharedtier/virtualClusters': + return SimpleNamespace(json=lambda: self._starter_cluster_payloads) + + if key.startswith('clusters/'): + wanted = key.split('/', 1)[1] + for item in self._cluster_payloads: + if item['clusterID'] == wanted: + return SimpleNamespace(json=lambda item=item: item) + raise ManagementError(errno=404, msg=f'cluster not found: {wanted}') + + raise AssertionError(f'unexpected request: GET {key}') + + def _put(self, path: Any, **kwargs: Any) -> Any: + if '/fs/' in str(path): + return self.files._put(path, **kwargs) + raise AssertionError(f'unexpected request: PUT {path}') + + def _patch(self, path: Any, **kwargs: Any) -> Any: + if '/fs/' in str(path): + return self.files._patch(path, **kwargs) + raise AssertionError(f'unexpected request: PATCH {path}') + + def _delete(self, path: Any, **kwargs: Any) -> Any: + if '/fs/' in str(path): + return self.files._delete(path, **kwargs) + raise AssertionError(f'unexpected request: DELETE {path}') + + def _post(self, path: Any, **kwargs: Any) -> Any: + raise AssertionError(f'unexpected request: POST {path}') + + def counts(self) -> Dict[str, int]: + """Return the number of recorded requests per method.""" + return self.files.counts() + + +def counting_cluster_manager(**kwargs: Any) -> CountingClusterManager: + """Return a :class:`CountingClusterManager`; see it for the arguments.""" + return CountingClusterManager(**kwargs) + + +def run_fusion_statement(sql: str, manager: Any) -> Any: + """ + Execute one Fusion statement against a counting cluster manager. + + The statement is parsed and run the way a cursor would run it, so the + requests recorded on ``manager.calls`` are the ones the whole statement + costs -- deployment resolution included -- rather than the ones a single + :class:`Stage` call makes. + + Parameters + ---------- + sql : str + The Fusion statement + manager : CountingClusterManager + The manager every handler in the statement resolves through + + Returns + ------- + FusionSQLResult + + """ + from singlestoredb.fusion import registry + from singlestoredb.fusion.handlers import cluster as cluster_handlers + from singlestoredb.fusion.handlers import utils as handler_utils + + # The results are formatted against the connection's decoders; there is no + # connection here and nothing to decode. + conn = SimpleNamespace(decoders={}, _results_type='tuples') + + with mock.patch.dict(os.environ, {'SINGLESTOREDB_FUSION_ENABLED': '1'}): + handler = registry.get_handler(sql) + if handler is None: + raise ValueError(f'no Fusion handler for statement: {sql}') + with mock.patch.object( + handler_utils, 'get_cluster_manager', return_value=manager, + ), mock.patch.object( + cluster_handlers, 'get_cluster_manager', return_value=manager, + ): + return handler(conn).execute(sql) + + +def clear_stage(deployment: Any) -> None: + """ + Empty a deployment's stage. + + A pool cluster carries whatever the class before left in its stage, and + ``TestStageFusion`` asserts exact listings of the stage root, so it starts + from a known-empty one rather than from whatever ran first. Failures are + logged rather than raised: this runs in a fixture, where the interesting + failure is the test's, not the cleanup's. + """ + stage = deployment.stage + + # The root listing is enough: a folder goes recursively, so there is no + # reason to enumerate what is inside it. + for obj in stage.listdir('/', return_objects=True): + try: + if obj.type == 'directory': + stage.removedirs(obj.path) + else: + stage.remove(obj.path) + except Exception as exc: + logger.warning(f'Could not clear stage path {obj.path}: {exc}') diff --git a/singlestoredb/warnings.py b/singlestoredb/warnings.py index 10edef75b..a00d15351 100644 --- a/singlestoredb/warnings.py +++ b/singlestoredb/warnings.py @@ -3,3 +3,19 @@ class PreviewFeatureWarning(UserWarning): """Warning for experimental preview features.""" pass + + +class DeprecatedFeatureWarning(UserWarning): + """ + Warning for deprecated features that still work. + + Deliberately a ``UserWarning`` rather than a ``DeprecationWarning``, for the + same reason :class:`PreviewFeatureWarning` is: Python ignores + ``DeprecationWarning`` by default outside ``__main__``, and these fire from + library frames several calls below the notebook cell that triggered them, so + a ``DeprecationWarning`` would reach almost nobody. The Python-level + management API uses the builtin there (see + :func:`singlestoredb.manage_workspaces`), where the caller's own frame is + close enough for the default filter to do the right thing. + """ + pass