From 49084dd9639f43a30647ac020b0ff369f22b6f36 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:30:05 -0400 Subject: [PATCH 01/17] update volumes model --- src/dstack/_internal/server/models.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index 705d1e5b97..a57b1fc63f 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -535,6 +535,9 @@ class VolumeModel(BaseModel): ) name: Mapped[str] = mapped_column(String(100)) + user_id: Mapped["UserModel"] = mapped_column(ForeignKey("users.id", ondelete="CASCADE")) + user: Mapped["UserModel"] = relationship(lazy="joined") + project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) project: Mapped["ProjectModel"] = relationship(foreign_keys=[project_id]) From 6d032a462292469dd62a7e770f8e19f88cd03e6e Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:31:34 -0400 Subject: [PATCH 02/17] generate db migration --- ...added_user_id_and_user_to_volumes_table.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py diff --git a/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py b/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py new file mode 100644 index 0000000000..5534471370 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py @@ -0,0 +1,45 @@ +"""Added user_id and user to Volumes table + +Revision ID: 9be642a33d08 +Revises: afbc600ff2b2 +Create Date: 2024-11-02 01:46:48.989331 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +# revision identifiers, used by Alembic. +revision = "9be642a33d08" +down_revision = "afbc600ff2b2" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("volumes", schema=None) as batch_op: + batch_op.add_column( + sa.Column( + "user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False + ) + ) + batch_op.create_foreign_key( + batch_op.f("fk_volumes_user_id_users"), + "users", + ["user_id"], + ["id"], + ondelete="CASCADE", + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("volumes", schema=None) as batch_op: + batch_op.drop_constraint(batch_op.f("fk_volumes_user_id_users"), type_="foreignkey") + batch_op.drop_column("user_id") + + # ### end Alembic commands ### From da2e26620f4d92cdde8afd65843a208a447bd785 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:32:46 -0400 Subject: [PATCH 03/17] add user field to Volume class --- src/dstack/_internal/core/models/volumes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/dstack/_internal/core/models/volumes.py b/src/dstack/_internal/core/models/volumes.py index b25aff5484..5ff5decb4a 100644 --- a/src/dstack/_internal/core/models/volumes.py +++ b/src/dstack/_internal/core/models/volumes.py @@ -64,6 +64,7 @@ class VolumeAttachmentData(CoreModel): class Volume(CoreModel): id: uuid.UUID name: str + user: str project_name: str configuration: VolumeConfiguration external: bool From c2246a2752062e2135b771763f4b41e6f4178f14 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:39:32 -0400 Subject: [PATCH 04/17] update volume service - create volume now requires a user of type UserModel to be passed in - db for volumes should have column for user_id -only passing back the user name when creating the volume --- src/dstack/_internal/server/services/volumes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/dstack/_internal/server/services/volumes.py b/src/dstack/_internal/server/services/volumes.py index 9a0359d233..3267e0fccc 100644 --- a/src/dstack/_internal/server/services/volumes.py +++ b/src/dstack/_internal/server/services/volumes.py @@ -164,6 +164,7 @@ async def get_project_volume_model_by_name( async def create_volume( session: AsyncSession, project: ProjectModel, + user: UserModel, configuration: VolumeConfiguration, ) -> Volume: _validate_volume_configuration(configuration) @@ -193,6 +194,7 @@ async def create_volume( volume_model = VolumeModel( id=uuid.uuid4(), name=configuration.name, + user_id=user.id, project=project, status=VolumeStatus.SUBMITTED, configuration=configuration.json(), @@ -263,6 +265,7 @@ def volume_model_to_volume(volume_model: VolumeModel) -> Volume: return Volume( name=volume_model.name, project_name=volume_model.project.name, + user=volume_model.user.name, configuration=configuration, external=configuration.volume_id is not None, created_at=volume_model.created_at.replace(tzinfo=timezone.utc), From b5d166d567d0a01234bbe6073b5b26cf49ba68ca Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:40:05 -0400 Subject: [PATCH 05/17] update router --- src/dstack/_internal/server/routers/volumes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/dstack/_internal/server/routers/volumes.py b/src/dstack/_internal/server/routers/volumes.py index 9004f7c25b..c04dc12b28 100644 --- a/src/dstack/_internal/server/routers/volumes.py +++ b/src/dstack/_internal/server/routers/volumes.py @@ -66,12 +66,14 @@ async def get_volume( async def create_volume( body: CreateVolumeRequest, session: AsyncSession = Depends(get_session), + user: UserModel = Depends(Authenticated()), user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), ) -> Volume: _, project = user_project return await volumes_services.create_volume( session=session, project=project, + user=user, configuration=body.configuration, ) From cf356b0f03e349a9265ed8821c75943251cd19d1 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:40:20 -0400 Subject: [PATCH 06/17] tag aws volume with dstack user that created it --- src/dstack/_internal/core/backends/aws/compute.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index b52f25f0a6..61544e8bfb 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -489,6 +489,7 @@ def create_volume(self, volume: Volume) -> VolumeProvisioningData: tags = { "Name": volume.configuration.name, "owner": "dstack", + "dstack_user": volume.user, "dstack_project": volume.project_name, } tags = merge_tags(tags=tags, backend_tags=self.config.tags) From 4006f2b8f37ab7af492b0eae0320dac2dd6b9764 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sat, 2 Nov 2024 21:43:14 -0400 Subject: [PATCH 07/17] add datack user to gcp volume label --- src/dstack/_internal/core/backends/gcp/compute.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 1f263d14a4..2133ccc833 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -500,6 +500,7 @@ def create_volume(self, volume: Volume) -> VolumeProvisioningData: labels = { "owner": "dstack", "dstack_project": volume.project_name.lower(), + "dstack_user": volume.user, } labels = {k: v for k, v in labels.items() if gcp_resources.is_valid_label_value(v)} labels = merge_tags(tags=labels, backend_tags=self.config.tags) From 99534473410ed54c22936cca4e02b54d2b66190e Mon Sep 17 00:00:00 2001 From: james-boydell Date: Sun, 3 Nov 2024 19:33:19 -0500 Subject: [PATCH 08/17] updating tests --- src/dstack/_internal/server/testing/common.py | 2 ++ .../tasks/test_process_submitted_volumes.py | 7 +++++-- src/tests/_internal/server/routers/test_volumes.py | 14 ++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 1a42df0949..6c6d9deb20 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -520,6 +520,7 @@ async def create_instance( async def create_volume( session: AsyncSession, project: ProjectModel, + user: UserModel, status: VolumeStatus = VolumeStatus.SUBMITTED, created_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), configuration: Optional[VolumeConfiguration] = None, @@ -532,6 +533,7 @@ async def create_volume( configuration = get_volume_configuration(backend=backend, region=region) vm = VolumeModel( project=project, + user_id=user.id, name=configuration.name, status=status, created_at=created_at, diff --git a/src/tests/_internal/server/background/tasks/test_process_submitted_volumes.py b/src/tests/_internal/server/background/tasks/test_process_submitted_volumes.py index 970c674b4a..86fbed603c 100644 --- a/src/tests/_internal/server/background/tasks/test_process_submitted_volumes.py +++ b/src/tests/_internal/server/background/tasks/test_process_submitted_volumes.py @@ -8,6 +8,7 @@ from dstack._internal.server.background.tasks.process_volumes import process_submitted_volumes from dstack._internal.server.testing.common import ( create_project, + create_user, create_volume, ) @@ -17,8 +18,9 @@ class TestProcessSubmittedVolumes: @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_fails_job_when_no_backends(self, test_db, session: AsyncSession): project = await create_project(session=session) + user = await create_user(session=session) volume = await create_volume( - session=session, project=project, status=VolumeStatus.SUBMITTED + session=session, project=project, user=user, status=VolumeStatus.SUBMITTED ) await process_submitted_volumes() await session.refresh(volume) @@ -29,8 +31,9 @@ async def test_fails_job_when_no_backends(self, test_db, session: AsyncSession): @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_provisiones_volumes(self, test_db, session: AsyncSession): project = await create_project(session=session) + user = await create_user(session=session) volume = await create_volume( - session=session, project=project, status=VolumeStatus.SUBMITTED + session=session, project=project, user=user, status=VolumeStatus.SUBMITTED ) with patch( "dstack._internal.server.services.backends.get_project_backend_by_type_or_error" diff --git a/src/tests/_internal/server/routers/test_volumes.py b/src/tests/_internal/server/routers/test_volumes.py index a39c6e16b6..7ef5154420 100644 --- a/src/tests/_internal/server/routers/test_volumes.py +++ b/src/tests/_internal/server/routers/test_volumes.py @@ -44,6 +44,7 @@ async def test_lists_volumes_across_projects( volume1 = await create_volume( session=session, project=project1, + user=user, created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), configuration=get_volume_configuration(name="volume1"), ) @@ -51,6 +52,7 @@ async def test_lists_volumes_across_projects( volume2 = await create_volume( session=session, project=project2, + user=user, created_at=datetime(2023, 1, 2, 3, 5, tzinfo=timezone.utc), configuration=get_volume_configuration(name="volume2"), ) @@ -65,6 +67,7 @@ async def test_lists_volumes_across_projects( "id": str(volume2.id), "name": volume2.name, "project_name": project2.name, + "user": user.name, "configuration": json.loads(volume2.configuration), "external": False, "created_at": "2023-01-02T03:05:00+00:00", @@ -79,6 +82,7 @@ async def test_lists_volumes_across_projects( "id": str(volume1.id), "name": volume1.name, "project_name": project1.name, + "user": user.name, "configuration": json.loads(volume1.configuration), "external": False, "created_at": "2023-01-02T03:04:00+00:00", @@ -134,12 +138,14 @@ async def test_non_admin_cannot_see_others_projects( volume1 = await create_volume( session=session, project=project1, + user=user1, created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), configuration=get_volume_configuration(name="volume1"), ) await create_volume( session=session, project=project2, + user=user2, created_at=datetime(2023, 1, 2, 3, 5, tzinfo=timezone.utc), configuration=get_volume_configuration(name="volume2"), ) @@ -154,6 +160,7 @@ async def test_non_admin_cannot_see_others_projects( "id": str(volume1.id), "name": volume1.name, "project_name": project1.name, + "user": user1.name, "configuration": json.loads(volume1.configuration), "external": False, "created_at": "2023-01-02T03:04:00+00:00", @@ -187,6 +194,7 @@ async def test_lists_volumes(self, test_db, session: AsyncSession, client: Async volume = await create_volume( session=session, project=project, + user=user, created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), ) response = await client.post( @@ -199,6 +207,7 @@ async def test_lists_volumes(self, test_db, session: AsyncSession, client: Async "id": str(volume.id), "name": volume.name, "project_name": project.name, + "user": user.name, "configuration": json.loads(volume.configuration), "external": False, "created_at": "2023-01-02T03:04:00+00:00", @@ -232,6 +241,7 @@ async def test_returns_volume(self, test_db, session: AsyncSession, client: Asyn volume = await create_volume( session=session, project=project, + user=user, created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), ) response = await client.post( @@ -244,6 +254,7 @@ async def test_returns_volume(self, test_db, session: AsyncSession, client: Asyn "id": str(volume.id), "name": volume.name, "project_name": project.name, + "user": user.name, "configuration": json.loads(volume.configuration), "external": False, "created_at": "2023-01-02T03:04:00+00:00", @@ -305,6 +316,7 @@ async def test_creates_volume(self, test_db, session: AsyncSession, client: Asyn "name": configuration.name, "project_name": project.name, "configuration": configuration, + "user": user.name, "external": False, "created_at": "2023-01-02T03:04:00+00:00", "status": "submitted", @@ -338,6 +350,7 @@ async def test_deletes_volumes(self, test_db, session: AsyncSession, client: Asy volume = await create_volume( session=session, project=project, + user=user, volume_provisioning_data=get_volume_provisioning_data(), ) with patch( @@ -369,6 +382,7 @@ async def test_returns_400_when_volumes_in_use( volume = await create_volume( session=session, project=project, + user=user, volume_provisioning_data=get_volume_provisioning_data(), ) instance = await create_instance( From 2fae777bba7426a8c659cb83630907f0b00d1e0b Mon Sep 17 00:00:00 2001 From: james-boydell Date: Mon, 4 Nov 2024 08:38:38 -0500 Subject: [PATCH 09/17] updating on how user is fetched --- src/dstack/_internal/server/routers/volumes.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/dstack/_internal/server/routers/volumes.py b/src/dstack/_internal/server/routers/volumes.py index c04dc12b28..a63fa8ba24 100644 --- a/src/dstack/_internal/server/routers/volumes.py +++ b/src/dstack/_internal/server/routers/volumes.py @@ -66,10 +66,9 @@ async def get_volume( async def create_volume( body: CreateVolumeRequest, session: AsyncSession = Depends(get_session), - user: UserModel = Depends(Authenticated()), user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), ) -> Volume: - _, project = user_project + user, project = user_project return await volumes_services.create_volume( session=session, project=project, From d6486894ee798b1fd8408067c2dd6ea2637f6718 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Mon, 4 Nov 2024 11:02:51 -0500 Subject: [PATCH 10/17] updating tests --- .../server/background/tasks/test_process_running_jobs.py | 1 + .../server/background/tasks/test_process_submitted_jobs.py | 1 + src/tests/_internal/server/routers/test_volumes.py | 1 + 3 files changed, 3 insertions(+) diff --git a/src/tests/_internal/server/background/tasks/test_process_running_jobs.py b/src/tests/_internal/server/background/tasks/test_process_running_jobs.py index c8c02e0c84..149ff6f98d 100644 --- a/src/tests/_internal/server/background/tasks/test_process_running_jobs.py +++ b/src/tests/_internal/server/background/tasks/test_process_running_jobs.py @@ -222,6 +222,7 @@ async def test_provisioning_shim_with_volumes( volume = await create_volume( session=session, project=project, + user=user, status=VolumeStatus.ACTIVE, configuration=get_volume_configuration( name="my-vol", backend=BackendType.AWS, region="us-east-1" diff --git a/src/tests/_internal/server/background/tasks/test_process_submitted_jobs.py b/src/tests/_internal/server/background/tasks/test_process_submitted_jobs.py index 8ffb44f9f5..995c6dd37d 100644 --- a/src/tests/_internal/server/background/tasks/test_process_submitted_jobs.py +++ b/src/tests/_internal/server/background/tasks/test_process_submitted_jobs.py @@ -392,6 +392,7 @@ async def test_assigns_job_to_instance_with_volumes(self, test_db, session: Asyn volume = await create_volume( session=session, project=project, + user=user, status=VolumeStatus.ACTIVE, volume_provisioning_data=get_volume_provisioning_data(), backend=BackendType.AWS, diff --git a/src/tests/_internal/server/routers/test_volumes.py b/src/tests/_internal/server/routers/test_volumes.py index 7ef5154420..7b08fe2e6d 100644 --- a/src/tests/_internal/server/routers/test_volumes.py +++ b/src/tests/_internal/server/routers/test_volumes.py @@ -108,6 +108,7 @@ async def test_lists_volumes_across_projects( "id": str(volume1.id), "name": volume1.name, "project_name": project1.name, + "user": user.name, "configuration": json.loads(volume1.configuration), "external": False, "created_at": "2023-01-02T03:04:00+00:00", From 88ba02bda3ac681110a51b34b48292a8ae4b1f90 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Mon, 4 Nov 2024 11:04:28 -0500 Subject: [PATCH 11/17] update data migration to update volume user_id based on project.owner_id --- ...3d08_added_user_id_and_user_to_volumes_table.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py b/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py index 5534471370..0c755af785 100644 --- a/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py +++ b/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py @@ -21,10 +21,9 @@ def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### with op.batch_alter_table("volumes", schema=None) as batch_op: batch_op.add_column( - sa.Column( - "user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False - ) + sa.Column("user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True) ) + batch_op.create_foreign_key( batch_op.f("fk_volumes_user_id_users"), "users", @@ -35,6 +34,15 @@ def upgrade() -> None: # ### end Alembic commands ### + # update any existing volumes and set the user_id equal to the project_owner.id which created the volume + op.execute( + "UPDATE volumes SET user_id = (SELECT owner_id FROM projects JOIN volumes ON projects.id = volumes.project_id) WHERE user_id IS NULL" + ) + + # set volumes.user_id to non-nullable + with op.batch_alter_table("volumes", schema=None) as batch_op: + batch_op.alter_column("user_id", nullable=False) + def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### From ea6c84797b317ba2a5dd239f98613ce1338db3bd Mon Sep 17 00:00:00 2001 From: james-boydell Date: Mon, 4 Nov 2024 15:53:25 -0500 Subject: [PATCH 12/17] update volume model to remove lazy join --- src/dstack/_internal/server/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index a57b1fc63f..0a83dfc100 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -536,7 +536,7 @@ class VolumeModel(BaseModel): name: Mapped[str] = mapped_column(String(100)) user_id: Mapped["UserModel"] = mapped_column(ForeignKey("users.id", ondelete="CASCADE")) - user: Mapped["UserModel"] = relationship(lazy="joined") + user: Mapped["UserModel"] = relationship() project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) project: Mapped["ProjectModel"] = relationship(foreign_keys=[project_id]) From 7df8749fe2249992cb24c9d074b1c5cb92476369 Mon Sep 17 00:00:00 2001 From: james-boydell Date: Mon, 4 Nov 2024 15:54:39 -0500 Subject: [PATCH 13/17] regenerate migration --- ...r_to_volumes_table.py => 82b32a135ea2_.py} | 105 +++++++++--------- 1 file changed, 52 insertions(+), 53 deletions(-) rename src/dstack/_internal/server/migrations/versions/{9be642a33d08_added_user_id_and_user_to_volumes_table.py => 82b32a135ea2_.py} (88%) diff --git a/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py b/src/dstack/_internal/server/migrations/versions/82b32a135ea2_.py similarity index 88% rename from src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py rename to src/dstack/_internal/server/migrations/versions/82b32a135ea2_.py index 0c755af785..29cd3e8e66 100644 --- a/src/dstack/_internal/server/migrations/versions/9be642a33d08_added_user_id_and_user_to_volumes_table.py +++ b/src/dstack/_internal/server/migrations/versions/82b32a135ea2_.py @@ -1,53 +1,52 @@ -"""Added user_id and user to Volumes table - -Revision ID: 9be642a33d08 -Revises: afbc600ff2b2 -Create Date: 2024-11-02 01:46:48.989331 - -""" - -import sqlalchemy as sa -import sqlalchemy_utils -from alembic import op - -# revision identifiers, used by Alembic. -revision = "9be642a33d08" -down_revision = "afbc600ff2b2" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - with op.batch_alter_table("volumes", schema=None) as batch_op: - batch_op.add_column( - sa.Column("user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True) - ) - - batch_op.create_foreign_key( - batch_op.f("fk_volumes_user_id_users"), - "users", - ["user_id"], - ["id"], - ondelete="CASCADE", - ) - - # ### end Alembic commands ### - - # update any existing volumes and set the user_id equal to the project_owner.id which created the volume - op.execute( - "UPDATE volumes SET user_id = (SELECT owner_id FROM projects JOIN volumes ON projects.id = volumes.project_id) WHERE user_id IS NULL" - ) - - # set volumes.user_id to non-nullable - with op.batch_alter_table("volumes", schema=None) as batch_op: - batch_op.alter_column("user_id", nullable=False) - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - with op.batch_alter_table("volumes", schema=None) as batch_op: - batch_op.drop_constraint(batch_op.f("fk_volumes_user_id_users"), type_="foreignkey") - batch_op.drop_column("user_id") - - # ### end Alembic commands ### +"""empty message + +Revision ID: 82b32a135ea2 +Revises: afbc600ff2b2 +Create Date: 2024-11-04 15:46:37.719531 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +# revision identifiers, used by Alembic. +revision = "82b32a135ea2" +down_revision = "afbc600ff2b2" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("volumes", schema=None) as batch_op: + batch_op.add_column( + sa.Column("user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True) + ) + batch_op.create_foreign_key( + batch_op.f("fk_volumes_user_id_users"), + "users", + ["user_id"], + ["id"], + ondelete="CASCADE", + ) + + # ### end Alembic commands ### + + # update any existing volumes and set the user_id equal to the project_owner.id which created the volume + op.execute( + "UPDATE volumes SET user_id = (SELECT owner_id FROM projects JOIN volumes ON projects.id = volumes.project_id) WHERE user_id IS NULL" + ) + + # set volumes.user_id to non-nullable + with op.batch_alter_table("volumes", schema=None) as batch_op: + batch_op.alter_column("user_id", nullable=False) + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("volumes", schema=None) as batch_op: + batch_op.drop_constraint(batch_op.f("fk_volumes_user_id_users"), type_="foreignkey") + batch_op.drop_column("user_id") + + # ### end Alembic commands ### From ff4fe25242e02360b14b1d137793d3204cf8f56c Mon Sep 17 00:00:00 2001 From: james-boydell Date: Mon, 4 Nov 2024 16:46:06 -0500 Subject: [PATCH 14/17] join in users --- src/dstack/_internal/server/background/tasks/process_volumes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/dstack/_internal/server/background/tasks/process_volumes.py b/src/dstack/_internal/server/background/tasks/process_volumes.py index 9089de5f3b..da6d145bce 100644 --- a/src/dstack/_internal/server/background/tasks/process_volumes.py +++ b/src/dstack/_internal/server/background/tasks/process_volumes.py @@ -49,6 +49,7 @@ async def _process_submitted_volume(session: AsyncSession, volume_model: VolumeM select(VolumeModel) .where(VolumeModel.id == volume_model.id) .options(joinedload(VolumeModel.project).joinedload(ProjectModel.backends)) + .options(joinedload(VolumeModel.user)) .execution_options(populate_existing=True) ) volume_model = res.unique().scalar_one() From b2ad8bf1515367ccc4ac689acac466d3fca47be7 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 5 Nov 2024 12:20:37 +0500 Subject: [PATCH 15/17] Set default Volume.user for client backward compatibility --- src/dstack/_internal/core/models/volumes.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/dstack/_internal/core/models/volumes.py b/src/dstack/_internal/core/models/volumes.py index 5ff5decb4a..54a63edb9f 100644 --- a/src/dstack/_internal/core/models/volumes.py +++ b/src/dstack/_internal/core/models/volumes.py @@ -64,7 +64,9 @@ class VolumeAttachmentData(CoreModel): class Volume(CoreModel): id: uuid.UUID name: str - user: str + # Default user to "" for client backward compatibility (old 0.18 servers). + # TODO: Remove in 0.19 + user: str = "" project_name: str configuration: VolumeConfiguration external: bool From a1c2b97c180a3afe6099e42999aca4b94d74393d Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 5 Nov 2024 12:42:43 +0500 Subject: [PATCH 16/17] Load VolumeModel.user The loading is required everywhere VolumeModel is converted to Volume. If not loaded, the server errors when processing/returning other users' volumes. E.g. when global admins lists all volumes. --- .../background/tasks/process_submitted_jobs.py | 5 ++++- src/dstack/_internal/server/services/volumes.py | 17 +++++++++++++---- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py index 5c9585489a..b2272d1511 100644 --- a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py +++ b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py @@ -290,7 +290,10 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel): # Take lock to prevent attaching volumes that are to be deleted. # If the volume was deleted before the lock, the volume will fail to attach and the job will fail. await session.execute( - select(VolumeModel).where(VolumeModel.id.in_(volumes_ids)).with_for_update() + select(VolumeModel) + .where(VolumeModel.id.in_(volumes_ids)) + .options(joinedload(VolumeModel.user)) + .with_for_update() ) async with get_locker().lock_ctx(VolumeModel.__tablename__, volumes_ids): if len(volume_models) > 0: diff --git a/src/dstack/_internal/server/services/volumes.py b/src/dstack/_internal/server/services/volumes.py index 3267e0fccc..482f71b899 100644 --- a/src/dstack/_internal/server/services/volumes.py +++ b/src/dstack/_internal/server/services/volumes.py @@ -4,7 +4,7 @@ from sqlalchemy import and_, func, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import selectinload +from sqlalchemy.orm import joinedload, selectinload from dstack._internal.core.backends import BACKENDS_WITH_VOLUMES_SUPPORT from dstack._internal.core.errors import ( @@ -102,7 +102,11 @@ async def list_projects_volume_models( if ascending: order_by = (VolumeModel.created_at.asc(), VolumeModel.id.desc()) res = await session.execute( - select(VolumeModel).where(*filters).order_by(*order_by).limit(limit) + select(VolumeModel) + .where(*filters) + .order_by(*order_by) + .limit(limit) + .options(joinedload(VolumeModel.user)) ) volume_models = list(res.scalars().all()) return volume_models @@ -130,7 +134,9 @@ async def list_project_volume_models( filters.append(VolumeModel.name.in_(names)) if not include_deleted: filters.append(VolumeModel.deleted == False) - res = await session.execute(select(VolumeModel).where(*filters)) + res = await session.execute( + select(VolumeModel).where(*filters).options(joinedload(VolumeModel.user)) + ) return list(res.scalars().all()) @@ -157,7 +163,9 @@ async def get_project_volume_model_by_name( ] if not include_deleted: filters.append(VolumeModel.deleted == False) - res = await session.execute(select(VolumeModel).where(*filters)) + res = await session.execute( + select(VolumeModel).where(*filters).options(joinedload(VolumeModel.user)) + ) return res.scalar_one_or_none() @@ -226,6 +234,7 @@ async def delete_volumes(session: AsyncSession, project: ProjectModel, names: Li VolumeModel.name.in_(names), VolumeModel.deleted == False, ) + .options(joinedload(VolumeModel.user)) .options(selectinload(VolumeModel.instances)) .execution_options(populate_existing=True) .with_for_update() From 45ce69fc1a37fbbe487f9d19a35fb659cab4ff22 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 5 Nov 2024 13:05:56 +0500 Subject: [PATCH 17/17] Fix VolumeModel.user loading on postgres --- .../_internal/server/background/tasks/process_submitted_jobs.py | 2 +- src/dstack/_internal/server/services/volumes.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py index b2272d1511..1cfb33a63b 100644 --- a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py +++ b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py @@ -292,7 +292,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel): await session.execute( select(VolumeModel) .where(VolumeModel.id.in_(volumes_ids)) - .options(joinedload(VolumeModel.user)) + .options(selectinload(VolumeModel.user)) .with_for_update() ) async with get_locker().lock_ctx(VolumeModel.__tablename__, volumes_ids): diff --git a/src/dstack/_internal/server/services/volumes.py b/src/dstack/_internal/server/services/volumes.py index 482f71b899..0063b3419f 100644 --- a/src/dstack/_internal/server/services/volumes.py +++ b/src/dstack/_internal/server/services/volumes.py @@ -234,7 +234,7 @@ async def delete_volumes(session: AsyncSession, project: ProjectModel, names: Li VolumeModel.name.in_(names), VolumeModel.deleted == False, ) - .options(joinedload(VolumeModel.user)) + .options(selectinload(VolumeModel.user)) .options(selectinload(VolumeModel.instances)) .execution_options(populate_existing=True) .with_for_update()