Skip to content

Commit 07184b3

Browse files
jvstmejvstme
authored
Exports API and CLI (#3647)
1 parent 3e17881 commit 07184b3

15 files changed

Lines changed: 1649 additions & 18 deletions

File tree

‎src/dstack/_internal/cli/commands/export.py‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
import argparse
2+
from typing import Any, Union
3+
4+
from rich.table import Table
5+
6+
from dstack._internal.cli.commands import APIBaseCommand
7+
from dstack._internal.cli.services.completion import ExportNameCompleter
8+
from dstack._internal.cli.utils.common import add_row_from_dict, confirm_ask, console
9+
from dstack._internal.core.models.exports import Export
10+
11+
12+
class ExportCommand(APIBaseCommand):
13+
NAME = "export"
14+
DESCRIPTION = "Manage exports"
15+
16+
def _register(self):
17+
super()._register()
18+
self._parser.set_defaults(subfunc=self._list)
19+
subparsers = self._parser.add_subparsers(dest="action")
20+
21+
list_parser = subparsers.add_parser(
22+
"list", help="List exports", formatter_class=self._parser.formatter_class
23+
)
24+
list_parser.set_defaults(subfunc=self._list)
25+
26+
create_parser = subparsers.add_parser(
27+
"create", help="Create an export", formatter_class=self._parser.formatter_class
28+
)
29+
create_parser.add_argument(
30+
"name",
31+
help="The name of the export",
32+
)
33+
create_parser.add_argument(
34+
"--importer",
35+
action="append",
36+
dest="importers",
37+
help="Importer project name (can be specified multiple times)",
38+
default=[],
39+
)
40+
create_parser.add_argument(
41+
"--fleet",
42+
action="append",
43+
dest="fleets",
44+
help="Fleet name to export (can be specified multiple times)",
45+
default=[],
46+
)
47+
create_parser.set_defaults(subfunc=self._create)
48+
49+
update_parser = subparsers.add_parser(
50+
"update", help="Update an export", formatter_class=self._parser.formatter_class
51+
)
52+
update_parser.add_argument(
53+
"name",
54+
help="The name of the export",
55+
).completer = ExportNameCompleter() # type: ignore[attr-defined]
56+
update_parser.add_argument(
57+
"--add-importer",
58+
action="append",
59+
dest="add_importers",
60+
help="Importer project name to add (can be specified multiple times)",
61+
default=[],
62+
)
63+
update_parser.add_argument(
64+
"--remove-importer",
65+
action="append",
66+
dest="remove_importers",
67+
help="Importer project name to remove (can be specified multiple times)",
68+
default=[],
69+
)
70+
update_parser.add_argument(
71+
"--add-fleet",
72+
action="append",
73+
dest="add_fleets",
74+
help="Fleet name to add (can be specified multiple times)",
75+
default=[],
76+
)
77+
update_parser.add_argument(
78+
"--remove-fleet",
79+
action="append",
80+
dest="remove_fleets",
81+
help="Fleet name to remove (can be specified multiple times)",
82+
default=[],
83+
)
84+
update_parser.set_defaults(subfunc=self._update)
85+
86+
delete_parser = subparsers.add_parser(
87+
"delete", help="Delete an export", formatter_class=self._parser.formatter_class
88+
)
89+
delete_parser.add_argument(
90+
"name",
91+
help="The name of the export",
92+
).completer = ExportNameCompleter() # type: ignore[attr-defined]
93+
delete_parser.add_argument(
94+
"-y", "--yes", help="Don't ask for confirmation", action="store_true"
95+
)
96+
delete_parser.set_defaults(subfunc=self._delete)
97+
98+
def _command(self, args: argparse.Namespace):
99+
super()._command(args)
100+
args.subfunc(args)
101+
102+
def _list(self, args: argparse.Namespace):
103+
exports = self.api.client.exports.list(self.api.project)
104+
print_exports_table(exports)
105+
106+
def _create(self, args: argparse.Namespace):
107+
with console.status("Creating export..."):
108+
export = self.api.client.exports.create(
109+
project_name=self.api.project,
110+
name=args.name,
111+
importer_projects=args.importers,
112+
exported_fleets=args.fleets,
113+
)
114+
print_exports_table([export])
115+
116+
def _update(self, args: argparse.Namespace):
117+
with console.status("Updating export..."):
118+
export = self.api.client.exports.update(
119+
project_name=self.api.project,
120+
name=args.name,
121+
add_importer_projects=args.add_importers,
122+
remove_importer_projects=args.remove_importers,
123+
add_exported_fleets=args.add_fleets,
124+
remove_exported_fleets=args.remove_fleets,
125+
)
126+
print_exports_table([export])
127+
128+
def _delete(self, args: argparse.Namespace):
129+
if not args.yes and not confirm_ask(f"Delete the export [code]{args.name}[/]?"):
130+
console.print("\nExiting...")
131+
return
132+
133+
with console.status("Deleting export..."):
134+
self.api.client.exports.delete(project_name=self.api.project, name=args.name)
135+
136+
console.print(f"Export [code]{args.name}[/] deleted")
137+
138+
139+
def print_exports_table(exports: list[Export]):
140+
table = Table(box=None)
141+
table.add_column("NAME", no_wrap=True)
142+
table.add_column("FLEETS")
143+
table.add_column("IMPORTERS")
144+
145+
for export in exports:
146+
fleets = (
147+
", ".join([f.name for f in export.exported_fleets]) if export.exported_fleets else "-"
148+
)
149+
importers = ", ".join([i.project_name for i in export.imports]) if export.imports else "-"
150+
151+
row: dict[Union[str, int], Any] = {
152+
"NAME": export.name,
153+
"FLEETS": fleets,
154+
"IMPORTERS": importers,
155+
}
156+
add_row_from_dict(table, row)
157+
158+
console.print(table)
159+
console.print()

‎src/dstack/_internal/cli/main.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from dstack._internal.cli.commands.completion import CompletionCommand
1010
from dstack._internal.cli.commands.delete import DeleteCommand
1111
from dstack._internal.cli.commands.event import EventCommand
12+
from dstack._internal.cli.commands.export import ExportCommand
1213
from dstack._internal.cli.commands.fleet import FleetCommand
1314
from dstack._internal.cli.commands.gateway import GatewayCommand
1415
from dstack._internal.cli.commands.init import InitCommand
@@ -66,6 +67,7 @@ def main():
6667
AttachCommand.register(subparsers)
6768
DeleteCommand.register(subparsers)
6869
EventCommand.register(subparsers)
70+
ExportCommand.register(subparsers)
6971
FleetCommand.register(subparsers)
7072
GatewayCommand.register(subparsers)
7173
InitCommand.register(subparsers)

‎src/dstack/_internal/cli/services/completion.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,11 @@ def fetch_resource_names(self, api: Client) -> Iterable[str]:
8080
return [r.name for r in api.client.secrets.list(api.project)]
8181

8282

83+
class ExportNameCompleter(BaseAPINameCompleter):
84+
def fetch_resource_names(self, api: Client) -> Iterable[str]:
85+
return [r.name for r in api.client.exports.list(api.project)]
86+
87+
8388
class ProjectNameCompleter(BaseCompleter):
8489
"""
8590
Completer for local project names.

‎src/dstack/_internal/core/models/exports.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import uuid
2+
3+
from dstack._internal.core.models.common import CoreModel
4+
5+
6+
class ExportImport(CoreModel):
7+
project_name: str
8+
9+
10+
class ExportedFleet(CoreModel):
11+
id: uuid.UUID
12+
name: str
13+
14+
15+
class Export(CoreModel):
16+
id: uuid.UUID
17+
name: str
18+
imports: list[ExportImport]
19+
exported_fleets: list[ExportedFleet]

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
auth,
3232
backends,
3333
events,
34+
exports,
3435
files,
3536
fleets,
3637
gateways,
@@ -253,6 +254,7 @@ def register_routes(app: FastAPI, ui: bool = True):
253254
app.include_router(files.router)
254255
app.include_router(events.root_router)
255256
app.include_router(templates.router)
257+
app.include_router(exports.project_router)
256258

257259
@app.exception_handler(ForbiddenError)
258260
async def forbidden_error_handler(request: Request, exc: ForbiddenError):

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

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,9 @@
4949
from dstack._internal.utils.logging import get_logger
5050

5151
logger = get_logger(__name__)
52+
# Default options (save-update, merge) + delete-orphan + delete (required by delete-orphan)
53+
# delete-orphan allows to automatically delete entities removed from the relationship
54+
CASCADE_DEFAULT_WITH_DELETE_ORPHAN = "save-update, merge, delete-orphan, delete"
5255

5356

5457
class NaiveDateTime(TypeDecorator):
@@ -769,10 +772,7 @@ class InstanceModel(PipelineModelMixin, BaseModel):
769772

770773
volume_attachments: Mapped[List["VolumeAttachmentModel"]] = relationship(
771774
back_populates="instance",
772-
# Add delete-orphan option so that removing entries from volume_attachments
773-
# automatically marks them for deletion.
774-
# SQLAlchemy requires delete when using delete-orphan.
775-
cascade="save-update, merge, delete-orphan, delete",
775+
cascade=CASCADE_DEFAULT_WITH_DELETE_ORPHAN,
776776
)
777777

778778
__table_args__ = (
@@ -1055,8 +1055,14 @@ class ExportModel(BaseModel):
10551055
)
10561056
project: Mapped["ProjectModel"] = relationship()
10571057
created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime)
1058-
imports: Mapped[List["ImportModel"]] = relationship(back_populates="export")
1059-
exported_fleets: Mapped[List["ExportedFleetModel"]] = relationship(back_populates="export")
1058+
imports: Mapped[List["ImportModel"]] = relationship(
1059+
back_populates="export",
1060+
cascade=CASCADE_DEFAULT_WITH_DELETE_ORPHAN,
1061+
)
1062+
exported_fleets: Mapped[List["ExportedFleetModel"]] = relationship(
1063+
back_populates="export",
1064+
cascade=CASCADE_DEFAULT_WITH_DELETE_ORPHAN,
1065+
)
10601066

10611067

10621068
class ImportModel(BaseModel):

‎src/dstack/_internal/server/routers/exports.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
from typing import Annotated
2+
3+
from fastapi import APIRouter, Depends
4+
from sqlalchemy.ext.asyncio import AsyncSession
5+
6+
from dstack._internal.core.models.exports import Export
7+
from dstack._internal.server.db import get_session
8+
from dstack._internal.server.models import ProjectModel, UserModel
9+
from dstack._internal.server.schemas.exports import (
10+
CreateExportRequest,
11+
DeleteExportRequest,
12+
UpdateExportRequest,
13+
)
14+
from dstack._internal.server.security.permissions import ProjectAdmin, ProjectMember
15+
from dstack._internal.server.services import exports as exports_services
16+
from dstack._internal.server.utils.routers import get_base_api_additional_responses
17+
18+
project_router = APIRouter(
19+
prefix="/api/project/{project_name}/exports",
20+
tags=["exports"],
21+
responses=get_base_api_additional_responses(),
22+
)
23+
24+
25+
@project_router.post("/create", response_model=Export)
26+
async def create_export(
27+
body: CreateExportRequest,
28+
session: Annotated[AsyncSession, Depends(get_session)],
29+
user_project: Annotated[tuple[UserModel, ProjectModel], Depends(ProjectAdmin())],
30+
):
31+
user, project = user_project
32+
return await exports_services.create_export(
33+
session=session,
34+
project=project,
35+
user=user,
36+
name=body.name,
37+
importer_project_names=body.importer_projects,
38+
exported_fleet_names=body.exported_fleets,
39+
)
40+
41+
42+
@project_router.post("/update", response_model=Export)
43+
async def update_export(
44+
body: UpdateExportRequest,
45+
session: Annotated[AsyncSession, Depends(get_session)],
46+
user_project: Annotated[tuple[UserModel, ProjectModel], Depends(ProjectAdmin())],
47+
):
48+
user, project = user_project
49+
return await exports_services.update_export(
50+
session=session,
51+
project=project,
52+
user=user,
53+
name=body.name,
54+
add_importer_project_names=body.add_importer_projects,
55+
remove_importer_project_names=body.remove_importer_projects,
56+
add_exported_fleet_names=body.add_exported_fleets,
57+
remove_exported_fleet_names=body.remove_exported_fleets,
58+
)
59+
60+
61+
@project_router.post("/delete")
62+
async def delete_export(
63+
body: DeleteExportRequest,
64+
session: Annotated[AsyncSession, Depends(get_session)],
65+
user_project: Annotated[tuple[UserModel, ProjectModel], Depends(ProjectAdmin())],
66+
):
67+
_, project = user_project
68+
await exports_services.delete_export(
69+
session=session,
70+
project=project,
71+
name=body.name,
72+
)
73+
74+
75+
@project_router.post("/list", response_model=list[Export])
76+
async def list_exports(
77+
session: Annotated[AsyncSession, Depends(get_session)],
78+
user_project: Annotated[tuple[UserModel, ProjectModel], Depends(ProjectMember())],
79+
):
80+
_, project = user_project
81+
return await exports_services.list_exports(
82+
session=session,
83+
project=project,
84+
)

‎src/dstack/_internal/server/schemas/exports.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from dstack._internal.core.models.common import CoreModel
2+
3+
4+
class CreateExportRequest(CoreModel):
5+
name: str
6+
importer_projects: list[str] = []
7+
exported_fleets: list[str] = []
8+
9+
10+
class UpdateExportRequest(CoreModel):
11+
name: str
12+
add_importer_projects: list[str] = []
13+
remove_importer_projects: list[str] = []
14+
add_exported_fleets: list[str] = []
15+
remove_exported_fleets: list[str] = []
16+
17+
18+
class DeleteExportRequest(CoreModel):
19+
name: str

0 commit comments

Comments
 (0)