Skip to content

Commit

Permalink
fix: Fir 35673 python sdk engine update fix renaming (#399)
Browse files Browse the repository at this point in the history
  • Loading branch information
stepansergeevitch authored Aug 13, 2024
1 parent 8568f9c commit 70d41ec
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 18 deletions.
50 changes: 33 additions & 17 deletions src/firebolt/model/V2/engine.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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"})
Expand Down Expand Up @@ -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__()

Expand Down Expand Up @@ -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):
Expand All @@ -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:
Expand Down
30 changes: 29 additions & 1 deletion tests/integration/resource_manager/V2/test_engine.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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()

0 comments on commit 70d41ec

Please sign in to comment.