Skip to content

Commit b2a31ca

Browse files
r4victorr4victor
authored
Store BackendType as string instead of enum in the DB (#2393)
* Store BackendType as EnumAsString * Update backend guide * Fix type annotation
1 parent 3c0aa03 commit b2a31ca

File tree

4 files changed

+203
-6
lines changed

4 files changed

+203
-6
lines changed

‎contributing/BACKENDS.md‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,8 +97,6 @@ Add any dependencies required by your cloud provider to `setup.py`. Create a sep
9797
Add a new enumeration member for your provider to `BackendType` (`src/dstack/_internal/core/models/backends/base.py`).
9898
Use the name of the provider.
9999

100-
Then create a database [migration](MIGRATIONS.md) to reflect the new enum member.
101-
102100
##### 2.4.2. Create the backend directory
103101

104102
Create a new directory under `src/dstack/_internal/core/backends` with the name of the backend type.

‎src/dstack/_internal/server/migrations/versions/bc8ca4a505c6_store_backendtype_as_string.py‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
"""Store BackendType as string
2+
3+
Revision ID: bc8ca4a505c6
4+
Revises: 98d1b92988bc
5+
Create Date: 2025-03-10 14:49:06.837118
6+
7+
"""
8+
9+
import sqlalchemy as sa
10+
from alembic import op
11+
from sqlalchemy.dialects import postgresql
12+
13+
# revision identifiers, used by Alembic.
14+
revision = "bc8ca4a505c6"
15+
down_revision = "98d1b92988bc"
16+
branch_labels = None
17+
depends_on = None
18+
19+
20+
def upgrade() -> None:
21+
# ### commands auto generated by Alembic - please adjust! ###
22+
with op.batch_alter_table("backends", schema=None) as batch_op:
23+
batch_op.alter_column(
24+
"type",
25+
existing_type=postgresql.ENUM(
26+
"AWS",
27+
"AZURE",
28+
"CUDO",
29+
"DATACRUNCH",
30+
"DSTACK",
31+
"GCP",
32+
"KUBERNETES",
33+
"LAMBDA",
34+
"LOCAL",
35+
"REMOTE",
36+
"NEBIUS",
37+
"OCI",
38+
"RUNPOD",
39+
"TENSORDOCK",
40+
"VASTAI",
41+
"VULTR",
42+
name="backendtype",
43+
),
44+
type_=sa.String(length=100),
45+
existing_nullable=False,
46+
)
47+
48+
with op.batch_alter_table("instances", schema=None) as batch_op:
49+
batch_op.alter_column(
50+
"backend",
51+
existing_type=postgresql.ENUM(
52+
"AWS",
53+
"AZURE",
54+
"CUDO",
55+
"DATACRUNCH",
56+
"DSTACK",
57+
"GCP",
58+
"KUBERNETES",
59+
"LAMBDA",
60+
"LOCAL",
61+
"REMOTE",
62+
"NEBIUS",
63+
"OCI",
64+
"RUNPOD",
65+
"TENSORDOCK",
66+
"VASTAI",
67+
"VULTR",
68+
name="backendtype",
69+
),
70+
type_=sa.String(length=100),
71+
existing_nullable=True,
72+
)
73+
74+
sa.Enum(
75+
"AWS",
76+
"AZURE",
77+
"CUDO",
78+
"DATACRUNCH",
79+
"DSTACK",
80+
"GCP",
81+
"KUBERNETES",
82+
"LAMBDA",
83+
"LOCAL",
84+
"REMOTE",
85+
"NEBIUS",
86+
"OCI",
87+
"RUNPOD",
88+
"TENSORDOCK",
89+
"VASTAI",
90+
"VULTR",
91+
name="backendtype",
92+
).drop(op.get_bind())
93+
# ### end Alembic commands ###
94+
95+
96+
def downgrade() -> None:
97+
# ### commands auto generated by Alembic - please adjust! ###
98+
sa.Enum(
99+
"AWS",
100+
"AZURE",
101+
"CUDO",
102+
"DATACRUNCH",
103+
"DSTACK",
104+
"GCP",
105+
"KUBERNETES",
106+
"LAMBDA",
107+
"LOCAL",
108+
"REMOTE",
109+
"NEBIUS",
110+
"OCI",
111+
"RUNPOD",
112+
"TENSORDOCK",
113+
"VASTAI",
114+
"VULTR",
115+
name="backendtype",
116+
).create(op.get_bind())
117+
with op.batch_alter_table("instances", schema=None) as batch_op:
118+
batch_op.alter_column(
119+
"backend",
120+
existing_type=sa.String(length=100),
121+
type_=postgresql.ENUM(
122+
"AWS",
123+
"AZURE",
124+
"CUDO",
125+
"DATACRUNCH",
126+
"DSTACK",
127+
"GCP",
128+
"KUBERNETES",
129+
"LAMBDA",
130+
"LOCAL",
131+
"REMOTE",
132+
"NEBIUS",
133+
"OCI",
134+
"RUNPOD",
135+
"TENSORDOCK",
136+
"VASTAI",
137+
"VULTR",
138+
name="backendtype",
139+
),
140+
existing_nullable=True,
141+
postgresql_using="backend::VARCHAR::backendtype",
142+
)
143+
144+
with op.batch_alter_table("backends", schema=None) as batch_op:
145+
batch_op.alter_column(
146+
"type",
147+
existing_type=sa.String(length=100),
148+
type_=postgresql.ENUM(
149+
"AWS",
150+
"AZURE",
151+
"CUDO",
152+
"DATACRUNCH",
153+
"DSTACK",
154+
"GCP",
155+
"KUBERNETES",
156+
"LAMBDA",
157+
"LOCAL",
158+
"REMOTE",
159+
"NEBIUS",
160+
"OCI",
161+
"RUNPOD",
162+
"TENSORDOCK",
163+
"VASTAI",
164+
"VULTR",
165+
name="backendtype",
166+
),
167+
existing_nullable=False,
168+
postgresql_using="type::VARCHAR::backendtype",
169+
)
170+
171+
# ### end Alembic commands ###

‎src/dstack/_internal/server/models.py‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import enum
12
import uuid
23
from datetime import datetime
34
from typing import Callable, List, Optional, Union
@@ -112,7 +113,11 @@ def set_encrypt_decrypt(
112113
cls._encrypt_func = encrypt_func
113114
cls._decrypt_func = decrypt_func
114115

115-
def process_bind_param(self, value: Union[DecryptedString, str], dialect):
116+
def process_bind_param(
117+
self, value: Optional[Union[DecryptedString, str]], dialect
118+
) -> Optional[str]:
119+
if value is None:
120+
return None
116121
if isinstance(value, str):
117122
# Passing string allows binding an encrypted value directly
118123
# e.g. for comparisons
@@ -130,6 +135,29 @@ def process_result_value(self, value: Optional[str], dialect) -> Optional[Decryp
130135
return DecryptedString(plaintext=None, decrypted=False, exc=e)
131136

132137

138+
class EnumAsString(TypeDecorator):
139+
"""
140+
A custom type decorator that stores enums as strings in the DB.
141+
"""
142+
143+
impl = String
144+
cache_ok = True
145+
146+
def __init__(self, enum_class: type[enum.Enum], *args, **kwargs):
147+
self.enum_class = enum_class
148+
super().__init__(*args, **kwargs)
149+
150+
def process_bind_param(self, value: Optional[enum.Enum], dialect) -> Optional[str]:
151+
if value is None:
152+
return None
153+
return value.name
154+
155+
def process_result_value(self, value: Optional[str], dialect) -> Optional[enum.Enum]:
156+
if value is None:
157+
return None
158+
return self.enum_class[value]
159+
160+
133161
constraint_naming_convention = {
134162
"ix": "ix_%(column_0_label)s",
135163
"uq": "uq_%(table_name)s_%(column_0_name)s",
@@ -222,7 +250,7 @@ class BackendModel(BaseModel):
222250
)
223251
project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE"))
224252
project: Mapped["ProjectModel"] = relationship()
225-
type: Mapped[BackendType] = mapped_column(Enum(BackendType))
253+
type: Mapped[BackendType] = mapped_column(EnumAsString(BackendType, 100))
226254

227255
config: Mapped[str] = mapped_column(String(20000))
228256
auth: Mapped[DecryptedString] = mapped_column(EncryptedString(20000))
@@ -533,7 +561,7 @@ class InstanceModel(BaseModel):
533561
last_termination_retry_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime)
534562

535563
# backend
536-
backend: Mapped[Optional[BackendType]] = mapped_column(Enum(BackendType))
564+
backend: Mapped[Optional[BackendType]] = mapped_column(EnumAsString(BackendType, 100))
537565
backend_data: Mapped[Optional[str]] = mapped_column(Text)
538566

539567
# offer

‎src/dstack/_internal/server/services/backends/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ async def validate_and_create_backend_model(
102102
)
103103
return BackendModel(
104104
project_id=project.id,
105-
type=configurator.TYPE.value,
105+
type=configurator.TYPE,
106106
config=backend_record.config,
107107
auth=DecryptedString(plaintext=backend_record.auth),
108108
)

0 commit comments

Comments
 (0)