From 2da9e3152b478d462c7f734b446f6e4be1bad214 Mon Sep 17 00:00:00 2001 From: Jason Rasmussen Date: Wed, 5 Mar 2025 08:38:23 -0500 Subject: [PATCH] refactor: download service (#16600) --- .../src/repositories/download.repository.ts | 36 ++++++ server/src/repositories/index.ts | 2 + server/src/services/base.service.ts | 2 + server/src/services/download.service.spec.ts | 118 ++++++++++-------- server/src/services/download.service.ts | 108 ++++++++-------- server/src/utils/pagination.ts | 1 + .../repositories/download.repository.mock.ts | 12 ++ server/test/utils.ts | 5 + 8 files changed, 172 insertions(+), 112 deletions(-) create mode 100644 server/src/repositories/download.repository.ts create mode 100644 server/test/repositories/download.repository.mock.ts diff --git a/server/src/repositories/download.repository.ts b/server/src/repositories/download.repository.ts new file mode 100644 index 0000000000..c9c62c90ce --- /dev/null +++ b/server/src/repositories/download.repository.ts @@ -0,0 +1,36 @@ +import { Injectable } from '@nestjs/common'; +import { Kysely } from 'kysely'; +import { InjectKysely } from 'nestjs-kysely'; +import { DB } from 'src/db'; +import { anyUuid } from 'src/utils/database'; + +const builder = (db: Kysely) => + db + .selectFrom('assets') + .innerJoin('exif', 'assetId', 'id') + .select(['assets.id', 'assets.livePhotoVideoId', 'exif.fileSizeInByte as size']) + .where('assets.deletedAt', 'is', null); + +@Injectable() +export class DownloadRepository { + constructor(@InjectKysely() private db: Kysely) {} + + downloadAssetIds(ids: string[]) { + return builder(this.db).where('assets.id', '=', anyUuid(ids)).stream(); + } + + downloadMotionAssetIds(ids: string[]) { + return builder(this.db).select(['assets.originalPath']).where('assets.id', '=', anyUuid(ids)).stream(); + } + + downloadAlbumId(albumId: string) { + return builder(this.db) + .innerJoin('albums_assets_assets', 'assets.id', 'albums_assets_assets.assetsId') + .where('albums_assets_assets.albumsId', '=', albumId) + .stream(); + } + + downloadUserId(userId: string) { + return builder(this.db).where('assets.ownerId', '=', userId).where('assets.isVisible', '=', true).stream(); + } +} diff --git a/server/src/repositories/index.ts b/server/src/repositories/index.ts index 180d8ccd4f..8c262edcde 100644 --- a/server/src/repositories/index.ts +++ b/server/src/repositories/index.ts @@ -9,6 +9,7 @@ import { ConfigRepository } from 'src/repositories/config.repository'; import { CronRepository } from 'src/repositories/cron.repository'; import { CryptoRepository } from 'src/repositories/crypto.repository'; import { DatabaseRepository } from 'src/repositories/database.repository'; +import { DownloadRepository } from 'src/repositories/download.repository'; import { EventRepository } from 'src/repositories/event.repository'; import { JobRepository } from 'src/repositories/job.repository'; import { LibraryRepository } from 'src/repositories/library.repository'; @@ -51,6 +52,7 @@ export const repositories = [ CronRepository, CryptoRepository, DatabaseRepository, + DownloadRepository, EventRepository, JobRepository, LibraryRepository, diff --git a/server/src/services/base.service.ts b/server/src/services/base.service.ts index a1720aeb31..f8c995c007 100644 --- a/server/src/services/base.service.ts +++ b/server/src/services/base.service.ts @@ -17,6 +17,7 @@ import { ConfigRepository } from 'src/repositories/config.repository'; import { CronRepository } from 'src/repositories/cron.repository'; import { CryptoRepository } from 'src/repositories/crypto.repository'; import { DatabaseRepository } from 'src/repositories/database.repository'; +import { DownloadRepository } from 'src/repositories/download.repository'; import { EventRepository } from 'src/repositories/event.repository'; import { JobRepository } from 'src/repositories/job.repository'; import { LibraryRepository } from 'src/repositories/library.repository'; @@ -66,6 +67,7 @@ export class BaseService { protected cronRepository: CronRepository, protected cryptoRepository: CryptoRepository, protected databaseRepository: DatabaseRepository, + protected downloadRepository: DownloadRepository, protected eventRepository: EventRepository, protected jobRepository: JobRepository, protected libraryRepository: LibraryRepository, diff --git a/server/src/services/download.service.spec.ts b/server/src/services/download.service.spec.ts index d9e60dfdb4..5139fbd58f 100644 --- a/server/src/services/download.service.spec.ts +++ b/server/src/services/download.service.spec.ts @@ -1,10 +1,9 @@ import { BadRequestException } from '@nestjs/common'; import { DownloadResponseDto } from 'src/dtos/download.dto'; -import { AssetEntity } from 'src/entities/asset.entity'; import { DownloadService } from 'src/services/download.service'; import { assetStub } from 'test/fixtures/asset.stub'; import { authStub } from 'test/fixtures/auth.stub'; -import { newTestService, ServiceMocks } from 'test/utils'; +import { makeStream, newTestService, ServiceMocks } from 'test/utils'; import { Readable } from 'typeorm/platform/PlatformTools.js'; import { vitest } from 'vitest'; @@ -12,7 +11,7 @@ const downloadResponse: DownloadResponseDto = { totalSize: 105_000, archives: [ { - assetIds: ['asset-id', 'asset-id'], + assetIds: ['asset-1', 'asset-2'], size: 105_000, }, ], @@ -172,53 +171,60 @@ describe(DownloadService.name, () => { }); it('should return a list of archives (assetIds)', async () => { - mocks.access.asset.checkOwnerAccess.mockResolvedValue(new Set(['asset-1', 'asset-2'])); - mocks.asset.getByIds.mockResolvedValue([assetStub.image, assetStub.video]); - const assetIds = ['asset-1', 'asset-2']; + + mocks.access.asset.checkOwnerAccess.mockResolvedValue(new Set(assetIds)); + mocks.downloadRepository.downloadAssetIds.mockReturnValue( + makeStream([ + { id: 'asset-1', livePhotoVideoId: null, size: 100_000 }, + { id: 'asset-2', livePhotoVideoId: null, size: 5000 }, + ]), + ); + await expect(sut.getDownloadInfo(authStub.admin, { assetIds })).resolves.toEqual(downloadResponse); - expect(mocks.asset.getByIds).toHaveBeenCalledWith(['asset-1', 'asset-2'], { exifInfo: true }); + expect(mocks.downloadRepository.downloadAssetIds).toHaveBeenCalledWith(['asset-1', 'asset-2']); }); it('should return a list of archives (albumId)', async () => { mocks.access.album.checkOwnerAccess.mockResolvedValue(new Set(['album-1'])); - mocks.asset.getByAlbumId.mockResolvedValue({ - items: [assetStub.image, assetStub.video], - hasNextPage: false, - }); + mocks.downloadRepository.downloadAlbumId.mockReturnValue( + makeStream([ + { id: 'asset-1', livePhotoVideoId: null, size: 100_000 }, + { id: 'asset-2', livePhotoVideoId: null, size: 5000 }, + ]), + ); await expect(sut.getDownloadInfo(authStub.admin, { albumId: 'album-1' })).resolves.toEqual(downloadResponse); expect(mocks.access.album.checkOwnerAccess).toHaveBeenCalledWith(authStub.admin.user.id, new Set(['album-1'])); - expect(mocks.asset.getByAlbumId).toHaveBeenCalledWith({ take: 2500, skip: 0 }, 'album-1'); + expect(mocks.downloadRepository.downloadAlbumId).toHaveBeenCalledWith('album-1'); }); it('should return a list of archives (userId)', async () => { - mocks.asset.getByUserId.mockResolvedValue({ - items: [assetStub.image, assetStub.video], - hasNextPage: false, - }); + mocks.downloadRepository.downloadUserId.mockReturnValue( + makeStream([ + { id: 'asset-1', livePhotoVideoId: null, size: 100_000 }, + { id: 'asset-2', livePhotoVideoId: null, size: 5000 }, + ]), + ); await expect(sut.getDownloadInfo(authStub.admin, { userId: authStub.admin.user.id })).resolves.toEqual( downloadResponse, ); - expect(mocks.asset.getByUserId).toHaveBeenCalledWith({ take: 2500, skip: 0 }, authStub.admin.user.id, { - isVisible: true, - }); + expect(mocks.downloadRepository.downloadUserId).toHaveBeenCalledWith(authStub.admin.user.id); }); it('should split archives by size', async () => { - mocks.asset.getByUserId.mockResolvedValue({ - items: [ - { ...assetStub.image, id: 'asset-1' }, - { ...assetStub.video, id: 'asset-2' }, - { ...assetStub.withLocation, id: 'asset-3' }, - { ...assetStub.noWebpPath, id: 'asset-4' }, - ], - hasNextPage: false, - }); + mocks.downloadRepository.downloadUserId.mockReturnValue( + makeStream([ + { id: 'asset-1', livePhotoVideoId: null, size: 5000 }, + { id: 'asset-2', livePhotoVideoId: null, size: 100_000 }, + { id: 'asset-3', livePhotoVideoId: null, size: 23_456 }, + { id: 'asset-4', livePhotoVideoId: null, size: 123_000 }, + ]), + ); await expect( sut.getDownloadInfo(authStub.admin, { @@ -235,49 +241,53 @@ describe(DownloadService.name, () => { }); it('should include the video portion of a live photo', async () => { - const assetIds = [assetStub.livePhotoStillAsset.id]; - const assets = [assetStub.livePhotoStillAsset, assetStub.livePhotoMotionAsset]; + const assetIds = ['asset-1', 'asset-2']; mocks.access.asset.checkOwnerAccess.mockResolvedValue(new Set(assetIds)); - mocks.asset.getByIds.mockImplementation( - (ids) => - Promise.resolve( - ids.map((id) => assets.find((asset) => asset.id === id)).filter((asset) => !!asset), - ) as Promise, + mocks.downloadRepository.downloadAssetIds.mockReturnValue( + makeStream([ + { id: 'asset-1', livePhotoVideoId: 'asset-3', size: 5000 }, + { id: 'asset-2', livePhotoVideoId: 'asset-4', size: 100_000 }, + ]), ); - await expect(sut.getDownloadInfo(authStub.admin, { assetIds })).resolves.toEqual({ - totalSize: 125_000, + mocks.downloadRepository.downloadMotionAssetIds.mockReturnValue( + makeStream([ + { id: 'asset-3', livePhotoVideoId: null, size: 23_456, originalPath: '/path/to/file.mp4' }, + { id: 'asset-4', livePhotoVideoId: null, size: 123_000, originalPath: '/path/to/file.mp4' }, + ]), + ); + + await expect(sut.getDownloadInfo(authStub.admin, { assetIds, archiveSize: 30_000 })).resolves.toEqual({ + totalSize: 251_456, archives: [ - { - assetIds: [assetStub.livePhotoStillAsset.id, assetStub.livePhotoMotionAsset.id], - size: 125_000, - }, + { assetIds: ['asset-1', 'asset-2'], size: 105_000 }, + { assetIds: ['asset-3', 'asset-4'], size: 146_456 }, ], }); }); it('should skip the video portion of an android live photo by default', async () => { - const assetIds = [assetStub.livePhotoStillAsset.id]; - const assets = [ - assetStub.livePhotoStillAsset, - { ...assetStub.livePhotoMotionAsset, originalPath: 'upload/encoded-video/uuid-MP.mp4' }, - ]; + const assetIds = ['asset-1']; mocks.access.asset.checkOwnerAccess.mockResolvedValue(new Set(assetIds)); - mocks.asset.getByIds.mockImplementation( - (ids) => - Promise.resolve( - ids.map((id) => assets.find((asset) => asset.id === id)).filter((asset) => !!asset), - ) as Promise, + + mocks.downloadRepository.downloadAssetIds.mockReturnValue( + makeStream([{ id: 'asset-1', livePhotoVideoId: 'asset-3', size: 5000 }]), + ); + + mocks.downloadRepository.downloadMotionAssetIds.mockReturnValue( + makeStream([ + { id: 'asset-2', livePhotoVideoId: null, size: 23_456, originalPath: 'upload/encoded-video/uuid-MP.mp4' }, + ]), ); await expect(sut.getDownloadInfo(authStub.admin, { assetIds })).resolves.toEqual({ - totalSize: 25_000, + totalSize: 5000, archives: [ { - assetIds: [assetStub.livePhotoStillAsset.id], - size: 25_000, + assetIds: ['asset-1'], + size: 5000, }, ], }); diff --git a/server/src/services/download.service.ts b/server/src/services/download.service.ts index 8b18bd0a07..cb664aea32 100644 --- a/server/src/services/download.service.ts +++ b/server/src/services/download.service.ts @@ -4,53 +4,72 @@ import { StorageCore } from 'src/cores/storage.core'; import { AssetIdsDto } from 'src/dtos/asset.dto'; import { AuthDto } from 'src/dtos/auth.dto'; import { DownloadArchiveInfo, DownloadInfoDto, DownloadResponseDto } from 'src/dtos/download.dto'; -import { AssetEntity } from 'src/entities/asset.entity'; import { Permission } from 'src/enum'; import { ImmichReadStream } from 'src/repositories/storage.repository'; import { BaseService } from 'src/services/base.service'; import { HumanReadableSize } from 'src/utils/bytes'; -import { usePagination } from 'src/utils/pagination'; import { getPreferences } from 'src/utils/preferences'; @Injectable() export class DownloadService extends BaseService { async getDownloadInfo(auth: AuthDto, dto: DownloadInfoDto): Promise { + let assets; + + if (dto.assetIds) { + const assetIds = dto.assetIds; + await this.requireAccess({ auth, permission: Permission.ASSET_DOWNLOAD, ids: assetIds }); + assets = this.downloadRepository.downloadAssetIds(assetIds); + } else if (dto.albumId) { + const albumId = dto.albumId; + await this.requireAccess({ auth, permission: Permission.ALBUM_DOWNLOAD, ids: [albumId] }); + assets = this.downloadRepository.downloadAlbumId(albumId); + } else if (dto.userId) { + const userId = dto.userId; + await this.requireAccess({ auth, permission: Permission.TIMELINE_DOWNLOAD, ids: [userId] }); + assets = this.downloadRepository.downloadUserId(userId); + } else { + throw new BadRequestException('assetIds, albumId, or userId is required'); + } + const targetSize = dto.archiveSize || HumanReadableSize.GiB * 4; + const metadata = await this.userRepository.getMetadata(auth.user.id); + const preferences = getPreferences(auth.user.email, metadata); + const motionIds = new Set(); const archives: DownloadArchiveInfo[] = []; let archive: DownloadArchiveInfo = { size: 0, assetIds: [] }; - const metadata = await this.userRepository.getMetadata(auth.user.id); - const preferences = getPreferences(auth.user.email, metadata); + const addToArchive = ({ id, size }: { id: string; size: number | null }) => { + archive.assetIds.push(id); + archive.size += Number(size || 0); - const assetPagination = await this.getDownloadAssets(auth, dto); - for await (const assets of assetPagination) { - // motion part of live photos - const motionIds = assets.map((asset) => asset.livePhotoVideoId).filter((id): id is string => !!id); - if (motionIds.length > 0) { - const motionAssets = await this.assetRepository.getByIds(motionIds, { exifInfo: true }); - for (const motionAsset of motionAssets) { - if ( - !StorageCore.isAndroidMotionPath(motionAsset.originalPath) || - preferences.download.includeEmbeddedVideos - ) { - assets.push(motionAsset); - } - } - } - - for (const asset of assets) { - archive.size += Number(asset.exifInfo?.fileSizeInByte || 0); - archive.assetIds.push(asset.id); - - if (archive.size > targetSize) { - archives.push(archive); - archive = { size: 0, assetIds: [] }; - } - } - - if (archive.assetIds.length > 0) { + if (archive.size > targetSize) { archives.push(archive); + archive = { size: 0, assetIds: [] }; } + }; + + for await (const asset of assets) { + // motion part of live photos + if (asset.livePhotoVideoId) { + motionIds.add(asset.livePhotoVideoId); + } + + addToArchive(asset); + } + + if (motionIds.size > 0) { + const motionAssets = this.downloadRepository.downloadMotionAssetIds([...motionIds]); + for await (const motionAsset of motionAssets) { + if (StorageCore.isAndroidMotionPath(motionAsset.originalPath) && !preferences.download.includeEmbeddedVideos) { + continue; + } + + addToArchive(motionAsset); + } + } + + if (archive.assetIds.length > 0) { + archives.push(archive); } let totalSize = 0; @@ -99,31 +118,4 @@ export class DownloadService extends BaseService { return { stream: zip.stream }; } - - private async getDownloadAssets(auth: AuthDto, dto: DownloadInfoDto): Promise> { - const PAGINATION_SIZE = 2500; - - if (dto.assetIds) { - const assetIds = dto.assetIds; - await this.requireAccess({ auth, permission: Permission.ASSET_DOWNLOAD, ids: assetIds }); - const assets = await this.assetRepository.getByIds(assetIds, { exifInfo: true }); - return usePagination(PAGINATION_SIZE, () => ({ hasNextPage: false, items: assets })); - } - - if (dto.albumId) { - const albumId = dto.albumId; - await this.requireAccess({ auth, permission: Permission.ALBUM_DOWNLOAD, ids: [albumId] }); - return usePagination(PAGINATION_SIZE, (pagination) => this.assetRepository.getByAlbumId(pagination, albumId)); - } - - if (dto.userId) { - const userId = dto.userId; - await this.requireAccess({ auth, permission: Permission.TIMELINE_DOWNLOAD, ids: [userId] }); - return usePagination(PAGINATION_SIZE, (pagination) => - this.assetRepository.getByUserId(pagination, userId, { isVisible: true }), - ); - } - - throw new BadRequestException('assetIds, albumId, or userId is required'); - } } diff --git a/server/src/utils/pagination.ts b/server/src/utils/pagination.ts index 7cb31d1e04..eb4106c86a 100644 --- a/server/src/utils/pagination.ts +++ b/server/src/utils/pagination.ts @@ -10,6 +10,7 @@ export interface PaginationResult { export type Paginated = Promise>; +/** @deprecated use `this.db. ... .stream()` instead */ export async function* usePagination( pageSize: number, getNextPage: (pagination: PaginationOptions) => PaginationResult | Paginated, diff --git a/server/test/repositories/download.repository.mock.ts b/server/test/repositories/download.repository.mock.ts new file mode 100644 index 0000000000..50824c4f3d --- /dev/null +++ b/server/test/repositories/download.repository.mock.ts @@ -0,0 +1,12 @@ +import { DownloadRepository } from 'src/repositories/download.repository'; +import { RepositoryInterface } from 'src/types'; +import { Mocked, vitest } from 'vitest'; + +export const newDownloadRepositoryMock = (): Mocked> => { + return { + downloadAssetIds: vitest.fn(), + downloadMotionAssetIds: vitest.fn(), + downloadAlbumId: vitest.fn(), + downloadUserId: vitest.fn(), + }; +}; diff --git a/server/test/utils.ts b/server/test/utils.ts index b15888e4a3..a4a06f91d3 100644 --- a/server/test/utils.ts +++ b/server/test/utils.ts @@ -17,6 +17,7 @@ import { ConfigRepository } from 'src/repositories/config.repository'; import { CronRepository } from 'src/repositories/cron.repository'; import { CryptoRepository } from 'src/repositories/crypto.repository'; import { DatabaseRepository } from 'src/repositories/database.repository'; +import { DownloadRepository } from 'src/repositories/download.repository'; import { EventRepository } from 'src/repositories/event.repository'; import { JobRepository } from 'src/repositories/job.repository'; import { LibraryRepository } from 'src/repositories/library.repository'; @@ -59,6 +60,7 @@ import { newConfigRepositoryMock } from 'test/repositories/config.repository.moc import { newCronRepositoryMock } from 'test/repositories/cron.repository.mock'; import { newCryptoRepositoryMock } from 'test/repositories/crypto.repository.mock'; import { newDatabaseRepositoryMock } from 'test/repositories/database.repository.mock'; +import { newDownloadRepositoryMock } from 'test/repositories/download.repository.mock'; import { newEventRepositoryMock } from 'test/repositories/event.repository.mock'; import { newJobRepositoryMock } from 'test/repositories/job.repository.mock'; import { newLibraryRepositoryMock } from 'test/repositories/library.repository.mock'; @@ -103,6 +105,7 @@ export type ServiceOverrides = { cron: CronRepository; crypto: CryptoRepository; database: DatabaseRepository; + downloadRepository: DownloadRepository; event: EventRepository; job: JobRepository; library: LibraryRepository; @@ -162,6 +165,7 @@ export const newTestService = ( asset: newAssetRepositoryMock(), config: newConfigRepositoryMock(), database: newDatabaseRepositoryMock(), + downloadRepository: newDownloadRepositoryMock(), event: newEventRepositoryMock(), job: newJobRepositoryMock(), apiKey: newKeyRepositoryMock(), @@ -206,6 +210,7 @@ export const newTestService = ( overrides.cron || (mocks.cron as As), overrides.crypto || (mocks.crypto as As), overrides.database || (mocks.database as As), + overrides.downloadRepository || (mocks.downloadRepository as As), overrides.event || (mocks.event as As), overrides.job || (mocks.job as As), overrides.library || (mocks.library as As),