From 1e51780172330f8eb22531b20a0d73f77a3061fd Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Wed, 3 Jun 2026 10:09:08 -0700 Subject: [PATCH] feat: richer Astra database creation (database_type and pcu_group_id) Address #408 by lifting the two hard-coded constraints in `AstraDBAdmin.create_database` / `async_create_database`: - `database_type` (default "vector") selects the database flavor instead of always sending `dbType: vector`. Passing None omits the field so the DevOps API applies its own default, enabling non-vector databases. The value is not coerced/validated, so future database types need no astrapy upgrade. - `pcu_group_id` assigns the new database to a PCU group (`pcuGroupUUID`) at creation time. Both parameters are optional and keyword-only; existing callers keep getting a vector database with no PCU group, so the change is fully backward compatible. Adds unit tests asserting the DevOps payload for the default, non-vector and fully-specified cases (sync and async) against a mock DevOps API. --- CHANGES | 9 + astrapy/admin/admin.py | 30 ++- tests/base/unit/test_admin_create_database.py | 199 ++++++++++++++++++ 3 files changed, 236 insertions(+), 2 deletions(-) create mode 100644 tests/base/unit/test_admin_create_database.py diff --git a/CHANGES b/CHANGES index 7e4f0ca1..637e3d88 100644 --- a/CHANGES +++ b/CHANGES @@ -1,3 +1,12 @@ +main +==== + +Astra DB admin, richer database creation through `[async_]create_database`: + - `database_type` parameter to choose the database flavor (default "vector"; pass None to omit and let the DevOps API default apply); + - `pcu_group_id` parameter to assign the new database to a PCU group (`pcuGroupUUID`) at creation time. + + + v 2.2.1 ======= diff --git a/astrapy/admin/admin.py b/astrapy/admin/admin.py index f43255b5..97a8456b 100644 --- a/astrapy/admin/admin.py +++ b/astrapy/admin/admin.py @@ -1064,6 +1064,8 @@ def create_database( cloud_provider: str, region: str, keyspace: str | None = None, + database_type: str | None = "vector", + pcu_group_id: str | None = None, wait_until_active: bool = True, database_admin_timeout_ms: int | None = None, request_timeout_ms: int | None = None, @@ -1080,6 +1082,16 @@ def create_database( region: any of the available cloud regions. keyspace: name for the one keyspace the database starts with. If omitted, DevOps API will use its default. + database_type: the kind ("flavor") of database to create. The default + is "vector", which provisions a vector-enabled database. Passing + a different value selects another database flavor, while passing + None omits the setting altogether and lets the DevOps API apply + its own default. No coercion or validation is enforced on this + value, so that newly-introduced database types can be requested + without requiring an astrapy upgrade. + pcu_group_id: if provided, the UUID of the PCU (Provisioned Capacity + Unit) group the new database should be assigned to at creation + time. If omitted, the database is not associated to any PCU group. wait_until_active: if True (default), the method returns only after the newly-created database is in ACTIVE state (a few minutes, usually). If False, it will return right after issuing the @@ -1153,8 +1165,9 @@ def create_database( "cloudProvider": cloud_provider, "region": region, "capacityUnits": 1, - "dbType": "vector", + "dbType": database_type, "keyspace": keyspace, + "pcuGroupUUID": pcu_group_id, }.items() if v is not None } @@ -1243,6 +1256,8 @@ async def async_create_database( cloud_provider: str, region: str, keyspace: str | None = None, + database_type: str | None = "vector", + pcu_group_id: str | None = None, wait_until_active: bool = True, database_admin_timeout_ms: int | None = None, request_timeout_ms: int | None = None, @@ -1260,6 +1275,16 @@ async def async_create_database( region: any of the available cloud regions. keyspace: name for the one keyspace the database starts with. If omitted, DevOps API will use its default. + database_type: the kind ("flavor") of database to create. The default + is "vector", which provisions a vector-enabled database. Passing + a different value selects another database flavor, while passing + None omits the setting altogether and lets the DevOps API apply + its own default. No coercion or validation is enforced on this + value, so that newly-introduced database types can be requested + without requiring an astrapy upgrade. + pcu_group_id: if provided, the UUID of the PCU (Provisioned Capacity + Unit) group the new database should be assigned to at creation + time. If omitted, the database is not associated to any PCU group. wait_until_active: if True (default), the method returns only after the newly-created database is in ACTIVE state (a few minutes, usually). If False, it will return right after issuing the @@ -1326,8 +1351,9 @@ async def async_create_database( "cloudProvider": cloud_provider, "region": region, "capacityUnits": 1, - "dbType": "vector", + "dbType": database_type, "keyspace": keyspace, + "pcuGroupUUID": pcu_group_id, }.items() if v is not None } diff --git a/tests/base/unit/test_admin_create_database.py b/tests/base/unit/test_admin_create_database.py new file mode 100644 index 00000000..656175e3 --- /dev/null +++ b/tests/base/unit/test_admin_create_database.py @@ -0,0 +1,199 @@ +# Copyright DataStax, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unit tests for the request payload built by the database-creation admin methods. + +These do not hit Astra DB: a mock DevOps API (``pytest_httpserver``) is pointed +at by the admin object, and the ``create_database`` payload is asserted through +the ``json`` request matcher (a mismatch yields a non-201 response, which in turn +surfaces as a ``DevOpsAPIException``). ``wait_until_active=False`` keeps each call +to a single DevOps request, with the new database id taken from the +``Location`` response header. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from pytest_httpserver import HTTPServer + +from astrapy.admin.admin import AstraDBAdmin, AstraDBDatabaseAdmin +from astrapy.utils.api_options import ( + APIOptions, + DevOpsAPIURLOptions, + defaultAPIOptions, +) +from astrapy.utils.request_tools import HttpMethod + +DEV_OPS_API_VERSION = "v2" +CREATE_DB_PATH = f"/{DEV_OPS_API_VERSION}/databases" + +DATABASE_ID = "01234567-89ab-cdef-0123-456789abcdef" +DATABASE_NAME = "test_database" +CLOUD_PROVIDER = "aws" +REGION = "us-east-1" +KEYSPACE = "the_keyspace" +PCU_GROUP_ID = "f5e6d7c8-1234-5678-9abc-def012345678" + +# Default invocation: a vector database, no keyspace, no PCU group. +VECTOR_DEFAULT_PAYLOAD = { + "name": DATABASE_NAME, + "tier": "serverless", + "cloudProvider": CLOUD_PROVIDER, + "region": REGION, + "capacityUnits": 1, + "dbType": "vector", +} +# database_type=None: the "dbType" key is omitted entirely. +NON_VECTOR_PAYLOAD = { + "name": DATABASE_NAME, + "tier": "serverless", + "cloudProvider": CLOUD_PROVIDER, + "region": REGION, + "capacityUnits": 1, +} +# All the new knobs supplied at once. +FULL_PAYLOAD = { + "name": DATABASE_NAME, + "tier": "serverless", + "cloudProvider": CLOUD_PROVIDER, + "region": REGION, + "capacityUnits": 1, + "dbType": "tabular", + "keyspace": KEYSPACE, + "pcuGroupUUID": PCU_GROUP_ID, +} + + +@pytest.fixture +def mock_astra_db_admin(httpserver: HTTPServer) -> AstraDBAdmin: + base_endpoint = httpserver.url_for("/") + api_options = defaultAPIOptions(environment="prod").with_override( + APIOptions( + token="t1", + dev_ops_api_url_options=DevOpsAPIURLOptions( + dev_ops_url=base_endpoint, + dev_ops_api_version=DEV_OPS_API_VERSION, + ), + ), + ) + return AstraDBAdmin(api_options=api_options) + + +def _expect_create_database( + httpserver: HTTPServer, expected_payload: dict[str, Any] +) -> None: + httpserver.expect_oneshot_request( + CREATE_DB_PATH, + method=HttpMethod.POST, + json=expected_payload, + ).respond_with_data("", status=201, headers={"Location": DATABASE_ID}) + + +class TestAdminCreateDatabaseDryMethods: + @pytest.mark.describe("create_database requests a vector database by default, sync") + def test_create_database_default_dbtype_sync( + self, httpserver: HTTPServer, mock_astra_db_admin: AstraDBAdmin + ) -> None: + _expect_create_database(httpserver, VECTOR_DEFAULT_PAYLOAD) + created = mock_astra_db_admin.create_database( + DATABASE_NAME, + cloud_provider=CLOUD_PROVIDER, + region=REGION, + wait_until_active=False, + ) + assert isinstance(created, AstraDBDatabaseAdmin) + assert created.id == DATABASE_ID + + @pytest.mark.describe( + "create_database requests a vector database by default, async" + ) + async def test_create_database_default_dbtype_async( + self, httpserver: HTTPServer, mock_astra_db_admin: AstraDBAdmin + ) -> None: + _expect_create_database(httpserver, VECTOR_DEFAULT_PAYLOAD) + created = await mock_astra_db_admin.async_create_database( + DATABASE_NAME, + cloud_provider=CLOUD_PROVIDER, + region=REGION, + wait_until_active=False, + ) + assert isinstance(created, AstraDBDatabaseAdmin) + assert created.id == DATABASE_ID + + @pytest.mark.describe("create_database with database_type=None omits dbType, sync") + def test_create_database_non_vector_sync( + self, httpserver: HTTPServer, mock_astra_db_admin: AstraDBAdmin + ) -> None: + _expect_create_database(httpserver, NON_VECTOR_PAYLOAD) + created = mock_astra_db_admin.create_database( + DATABASE_NAME, + cloud_provider=CLOUD_PROVIDER, + region=REGION, + database_type=None, + wait_until_active=False, + ) + assert created.id == DATABASE_ID + + @pytest.mark.describe("create_database with database_type=None omits dbType, async") + async def test_create_database_non_vector_async( + self, httpserver: HTTPServer, mock_astra_db_admin: AstraDBAdmin + ) -> None: + _expect_create_database(httpserver, NON_VECTOR_PAYLOAD) + created = await mock_astra_db_admin.async_create_database( + DATABASE_NAME, + cloud_provider=CLOUD_PROVIDER, + region=REGION, + database_type=None, + wait_until_active=False, + ) + assert created.id == DATABASE_ID + + @pytest.mark.describe( + "create_database forwards database_type, keyspace and pcu_group_id, sync" + ) + def test_create_database_full_payload_sync( + self, httpserver: HTTPServer, mock_astra_db_admin: AstraDBAdmin + ) -> None: + _expect_create_database(httpserver, FULL_PAYLOAD) + created = mock_astra_db_admin.create_database( + DATABASE_NAME, + cloud_provider=CLOUD_PROVIDER, + region=REGION, + keyspace=KEYSPACE, + database_type="tabular", + pcu_group_id=PCU_GROUP_ID, + wait_until_active=False, + ) + assert created.id == DATABASE_ID + + @pytest.mark.describe( + "create_database forwards database_type, keyspace and pcu_group_id, async" + ) + async def test_create_database_full_payload_async( + self, httpserver: HTTPServer, mock_astra_db_admin: AstraDBAdmin + ) -> None: + _expect_create_database(httpserver, FULL_PAYLOAD) + created = await mock_astra_db_admin.async_create_database( + DATABASE_NAME, + cloud_provider=CLOUD_PROVIDER, + region=REGION, + keyspace=KEYSPACE, + database_type="tabular", + pcu_group_id=PCU_GROUP_ID, + wait_until_active=False, + ) + assert created.id == DATABASE_ID