diff --git a/src/firebolt/model/V2/engine.py b/src/firebolt/model/V2/engine.py index 967b5278171..58e4c433d0b 100644 --- a/src/firebolt/model/V2/engine.py +++ b/src/firebolt/model/V2/engine.py @@ -1,9 +1,10 @@ from __future__ import annotations import logging +import re import time from dataclasses import dataclass, field -from typing import TYPE_CHECKING, ClassVar, Optional, Tuple, Union +from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Union from firebolt.db import Connection, connect from firebolt.model.V2 import FireboltBaseModel @@ -26,15 +27,17 @@ class Engine(FireboltBaseModel): START_SQL: ClassVar[str] = 'START ENGINE "{}"' STOP_SQL: ClassVar[str] = 'STOP ENGINE "{}"' - ALTER_PREFIX_SQL: ClassVar[str] = 'ALTER ENGINE "{}" SET ' + ALTER_PREFIX_SQL: ClassVar[str] = 'ALTER ENGINE "{}" ' ALTER_PARAMETER_NAMES: ClassVar[Tuple] = ( "NODES", "TYPE", "AUTO_STOP", - "RENAME_TO", ) DROP_SQL: ClassVar[str] = 'DROP ENGINE "{}"' + # Engine names can only contain alphanumeric characters and underscores + _engine_name_re = re.compile(r"^[a-zA-Z0-9_]+$") + _service: EngineService = field(repr=False, compare=False) name: str = field(metadata={"db_name": "engine_name"}) @@ -69,11 +72,11 @@ def database(self) -> Optional[Database]: pass return None - def refresh(self) -> None: + def refresh(self, name: Optional[str] = None) -> None: """Update attributes of the instance from Firebolt.""" field_name_overrides = self._get_field_overrides() - for name, value in self._service._get_dict(self.name).items(): - setattr(self, field_name_overrides.get(name, name), value) + for field_name, value in self._service._get_dict(name or self.name).items(): + setattr(self, field_name_overrides.get(field_name, field_name), value) self.__post_init__() @@ -194,6 +197,9 @@ def update( # Nothing to be updated return self + if name is not None and any(x is not None for x in (scale, spec, auto_stop)): + raise ValueError("Cannot update name and other parameters at the same time") + self.refresh() self._wait_for_start_stop() if self.current_status in (EngineStatus.DROPPING, EngineStatus.REPAIRING): @@ -203,21 +209,31 @@ def update( ) sql = self.ALTER_PREFIX_SQL.format(self.name) - parameters = [] - for param, value in zip( - self.ALTER_PARAMETER_NAMES, - (scale, spec, auto_stop, name), - ): - if value is not None: - sql_part, new_params = self._service._format_engine_attribute_sql( - param, value + parameters: List[Union[str, int]] = [] + if name is not None: + if not self._engine_name_re.match(name): + raise ValueError( + f"Engine name {name} is invalid, " + "it must only contain alphanumeric characters and underscores." ) - sql += sql_part - parameters.extend(new_params) + sql += f" RENAME TO {name}" + else: + sql += " SET " + parameters = [] + for param, value in zip( + self.ALTER_PARAMETER_NAMES, + (scale, spec, auto_stop), + ): + if value is not None: + sql_part, new_params = self._service._format_engine_attribute_sql( + param, value + ) + sql += sql_part + parameters.extend(new_params) with self._service._connection.cursor() as c: c.execute(sql, parameters) - self.refresh() + self.refresh(name) return self def delete(self) -> None: diff --git a/tests/integration/resource_manager/V2/test_engine.py b/tests/integration/resource_manager/V2/test_engine.py index 06a653f09e1..96f43bd2139 100644 --- a/tests/integration/resource_manager/V2/test_engine.py +++ b/tests/integration/resource_manager/V2/test_engine.py @@ -1,6 +1,6 @@ from collections import namedtuple -from pytest import mark +from pytest import mark, raises from firebolt.client.auth import Auth from firebolt.service.manager import ResourceManager @@ -139,3 +139,31 @@ def test_engine_update_multiple_parameters( finally: engine.stop() engine.delete() + + +def test_engine_rename( + auth: Auth, + account_name: str, + api_endpoint: str, + database_name: str, + engine_name: str, +): + rm = ResourceManager( + auth=auth, account_name=account_name, api_endpoint=api_endpoint + ) + name = f"{engine_name}_to_rename" + new_name = f"{engine_name}_renamed" + engine = rm.engines.create(name=name) + + try: + with raises(ValueError): + engine.update(name="name; drop database users") + + engine.update(name=new_name) + assert engine.name == new_name + + new_engine = rm.engines.get(new_name) + assert new_engine.name == new_name + finally: + engine.stop() + engine.delete()