diff --git a/backend/endpoints/sockets/scan.py b/backend/endpoints/sockets/scan.py index f1686975b..1deb5567c 100644 --- a/backend/endpoints/sockets/scan.py +++ b/backend/endpoints/sockets/scan.py @@ -1,4 +1,7 @@ -from typing import Final +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Final import emoji import socketio # type: ignore @@ -10,6 +13,7 @@ from exceptions.fs_exceptions import ( FolderStructureNotMatchException, RomsNotFoundException, ) +from exceptions.socket_exceptions import ScanStoppedException from handler.database import db_firmware_handler, db_platform_handler, db_rom_handler from handler.filesystem import ( fs_firmware_handler, @@ -23,6 +27,7 @@ from handler.redis_handler import high_prio_queue, redis_client, redis_url from handler.scan_handler import ScanType, scan_firmware, scan_platform, scan_rom from handler.socket_handler import socket_handler from logger.logger import log +from models.platform import Platform from models.rom import Rom from rq import Worker from rq.job import Job @@ -32,16 +37,30 @@ from utils.context import initialize_context STOP_SCAN_FLAG: Final = "scan:stop" +@dataclass class ScanStats: - def __init__(self): - self.scanned_platforms = 0 - self.added_platforms = 0 - self.metadata_platforms = 0 - self.scanned_roms = 0 - self.added_roms = 0 - self.metadata_roms = 0 - self.scanned_firmware = 0 - self.added_firmware = 0 + scanned_platforms: int = 0 + added_platforms: int = 0 + metadata_platforms: int = 0 + scanned_roms: int = 0 + added_roms: int = 0 + metadata_roms: int = 0 + scanned_firmware: int = 0 + added_firmware: int = 0 + + def __add__(self, other: Any) -> ScanStats: + if not isinstance(other, ScanStats): + return NotImplemented + return ScanStats( + scanned_platforms=self.scanned_platforms + other.scanned_platforms, + added_platforms=self.added_platforms + other.added_platforms, + metadata_platforms=self.metadata_platforms + other.metadata_platforms, + scanned_roms=self.scanned_roms + other.scanned_roms, + added_roms=self.added_roms + other.added_roms, + metadata_roms=self.metadata_roms + other.metadata_roms, + scanned_firmware=self.scanned_firmware + other.scanned_firmware, + added_firmware=self.added_firmware + other.added_firmware, + ) def _get_socket_manager(): @@ -135,164 +154,14 @@ async def scan_platforms( log.info(f"Found {len(platform_list)} platforms in file system ") for platform_slug in platform_list: - # Stop the scan if the flag is set - if redis_client.get(STOP_SCAN_FLAG): - await stop_scan() - break - - platform = db_platform_handler.get_platform_by_fs_slug(platform_slug) - if platform and scan_type == ScanType.NEW_PLATFORMS: - continue - - scanned_platform = scan_platform( - platform_slug, fs_platforms, metadata_sources=metadata_sources + scan_stats += await _identify_platform( + platform_slug=platform_slug, + scan_type=scan_type, + fs_platforms=fs_platforms, + selected_roms=selected_roms, + metadata_sources=metadata_sources, + socket_manager=sm, ) - if platform: - scanned_platform.id = platform.id - # Keep the existing ids if they exist on the platform - scanned_platform.igdb_id = scanned_platform.igdb_id or platform.igdb_id - scanned_platform.moby_id = scanned_platform.moby_id or platform.moby_id - - scan_stats.scanned_platforms += 1 - scan_stats.added_platforms += 1 if not platform else 0 - scan_stats.metadata_platforms += ( - 1 if scanned_platform.igdb_id or scanned_platform.moby_id else 0 - ) - - platform = db_platform_handler.add_platform(scanned_platform) - - await sm.emit( - "scan:scanning_platform", - PlatformSchema.model_validate(platform).model_dump( - include={"id", "name", "slug"} - ), - ) - await sm.emit("", None) - - # Scanning firmware - try: - fs_firmware = fs_firmware_handler.get_firmware(platform) - except FirmwareNotFoundException: - fs_firmware = [] - - if len(fs_firmware) == 0: - log.warning( - " ⚠️ No firmware found, skipping firmware scan for this platform" - ) - else: - log.info(f" {len(fs_firmware)} firmware files found") - - for fs_fw in fs_firmware: - # Break early if the flag is set - if redis_client.get(STOP_SCAN_FLAG): - break - - firmware = db_firmware_handler.get_firmware_by_filename( - platform.id, fs_fw - ) - - scanned_firmware = scan_firmware( - platform=platform, - file_name=fs_fw, - firmware=firmware, - ) - - scan_stats.scanned_firmware += 1 - scan_stats.added_firmware += 1 if not firmware else 0 - - _added_firmware = db_firmware_handler.add_firmware(scanned_firmware) - firmware = db_firmware_handler.get_firmware(_added_firmware.id) - - # Scanning roms - try: - fs_roms = fs_rom_handler.get_roms(platform) - except RomsNotFoundException as e: - log.error(e) - continue - - if len(fs_roms) == 0: - log.warning( - " ⚠️ No roms found, verify that the folder structure is correct" - ) - else: - log.info(f" {len(fs_roms)} roms found") - - for fs_rom in fs_roms: - # Break early if the flag is set - if redis_client.get(STOP_SCAN_FLAG): - break - - rom = db_rom_handler.get_rom_by_filename( - platform.id, fs_rom["file_name"] - ) - - if _should_scan_rom( - scan_type=scan_type, rom=rom, selected_roms=selected_roms - ): - scanned_rom = await scan_rom( - platform=platform, - rom_attrs=fs_rom, - scan_type=scan_type, - rom=rom, - metadata_sources=metadata_sources, - ) - - scan_stats.scanned_roms += 1 - scan_stats.added_roms += 1 if not rom else 0 - scan_stats.metadata_roms += ( - 1 if scanned_rom.igdb_id or scanned_rom.moby_id else 0 - ) - - _added_rom = db_rom_handler.add_rom(scanned_rom) - - path_cover_s, path_cover_l = await fs_resource_handler.get_cover( - overwrite=True, - entity=_added_rom, - url_cover=_added_rom.url_cover, - ) - - path_screenshots = await fs_resource_handler.get_rom_screenshots( - rom=_added_rom, - url_screenshots=_added_rom.url_screenshots, - ) - - _added_rom.path_cover_s = path_cover_s - _added_rom.path_cover_l = path_cover_l - _added_rom.path_screenshots = path_screenshots - # Update the scanned rom with the cover and screenshots paths and update database - db_rom_handler.update_rom( - _added_rom.id, - { - c: getattr(_added_rom, c) - for c in inspect(_added_rom).mapper.column_attrs.keys() - }, - ) - - await sm.emit( - "scan:scanning_rom", - { - "platform_name": platform.name, - "platform_slug": platform.slug, - **RomSchema.model_validate(_added_rom).model_dump( - exclude={"created_at", "updated_at", "rom_user"} - ), - }, - ) - await sm.emit("", None) - - # Only purge entries if there are some file remaining in the library - # This protects against accidental deletion of entries when - # the folder structure is not correct or the drive is not mounted - if len(fs_roms) > 0: - db_rom_handler.purge_roms( - platform.id, [rom["file_name"] for rom in fs_roms] - ) - - # Same protection for firmware - if len(fs_firmware) > 0: - db_firmware_handler.purge_firmware( - platform.id, [fw for fw in fs_firmware] - ) # Same protection for platforms if len(fs_platforms) > 0: @@ -300,6 +169,9 @@ async def scan_platforms( log.info(emoji.emojize(":check_mark: Scan completed ")) await sm.emit("scan:done", scan_stats.__dict__) + except ScanStoppedException: + await stop_scan() + return except Exception as e: log.error(e) # Catch all exceptions and emit error to the client @@ -307,6 +179,198 @@ async def scan_platforms( return +async def _identify_platform( + platform_slug: str, + scan_type: ScanType, + fs_platforms: list[str], + selected_roms: list[str], + metadata_sources: list[str], + socket_manager: socketio.AsyncRedisManager, +) -> ScanStats: + # Stop the scan if the flag is set + if redis_client.get(STOP_SCAN_FLAG): + raise ScanStoppedException() + + scan_stats = ScanStats() + + platform = db_platform_handler.get_platform_by_fs_slug(platform_slug) + if platform and scan_type == ScanType.NEW_PLATFORMS: + return scan_stats + + scanned_platform = scan_platform( + platform_slug, fs_platforms, metadata_sources=metadata_sources + ) + if platform: + scanned_platform.id = platform.id + # Keep the existing ids if they exist on the platform + scanned_platform.igdb_id = scanned_platform.igdb_id or platform.igdb_id + scanned_platform.moby_id = scanned_platform.moby_id or platform.moby_id + + scan_stats.scanned_platforms += 1 + scan_stats.added_platforms += 1 if not platform else 0 + scan_stats.metadata_platforms += ( + 1 if scanned_platform.igdb_id or scanned_platform.moby_id else 0 + ) + + platform = db_platform_handler.add_platform(scanned_platform) + + await socket_manager.emit( + "scan:scanning_platform", + PlatformSchema.model_validate(platform).model_dump( + include={"id", "name", "slug"} + ), + ) + await socket_manager.emit("", None) + + # Scanning firmware + try: + fs_firmware = fs_firmware_handler.get_firmware(platform) + except FirmwareNotFoundException: + fs_firmware = [] + + if len(fs_firmware) == 0: + log.warning(" ⚠️ No firmware found, skipping firmware scan for this platform") + else: + log.info(f" {len(fs_firmware)} firmware files found") + + for fs_fw in fs_firmware: + scan_stats += await _identify_firmware( + platform=platform, + fs_fw=fs_fw, + ) + + # Scanning roms + try: + fs_roms = fs_rom_handler.get_roms(platform) + except RomsNotFoundException as e: + log.error(e) + return scan_stats + + if len(fs_roms) == 0: + log.warning(" ⚠️ No roms found, verify that the folder structure is correct") + else: + log.info(f" {len(fs_roms)} roms found") + + for fs_rom in fs_roms: + scan_stats += await _identify_rom( + platform=platform, + fs_rom=fs_rom, + scan_type=scan_type, + selected_roms=selected_roms, + metadata_sources=metadata_sources, + socket_manager=socket_manager, + ) + + # Only purge entries if there are some file remaining in the library + # This protects against accidental deletion of entries when + # the folder structure is not correct or the drive is not mounted + + if len(fs_roms) > 0: + db_rom_handler.purge_roms(platform.id, [rom["file_name"] for rom in fs_roms]) + + # Same protection for firmware + if len(fs_firmware) > 0: + db_firmware_handler.purge_firmware(platform.id, [fw for fw in fs_firmware]) + + return scan_stats + + +async def _identify_firmware( + platform: Platform, + fs_fw: str, +) -> ScanStats: + scan_stats = ScanStats() + + # Break early if the flag is set + if redis_client.get(STOP_SCAN_FLAG): + return scan_stats + + firmware = db_firmware_handler.get_firmware_by_filename(platform.id, fs_fw) + + scanned_firmware = scan_firmware( + platform=platform, + file_name=fs_fw, + firmware=firmware, + ) + + scan_stats.scanned_firmware += 1 + scan_stats.added_firmware += 1 if not firmware else 0 + + db_firmware_handler.add_firmware(scanned_firmware) + return scan_stats + + +async def _identify_rom( + platform: Platform, + fs_rom: dict, + scan_type: ScanType, + selected_roms: list[str], + metadata_sources: list[str], + socket_manager: socketio.AsyncRedisManager, +) -> ScanStats: + scan_stats = ScanStats() + + # Break early if the flag is set + if redis_client.get(STOP_SCAN_FLAG): + return scan_stats + + rom = db_rom_handler.get_rom_by_filename(platform.id, fs_rom["file_name"]) + + if not _should_scan_rom(scan_type=scan_type, rom=rom, selected_roms=selected_roms): + return scan_stats + + scanned_rom = await scan_rom( + platform=platform, + rom_attrs=fs_rom, + scan_type=scan_type, + rom=rom, + metadata_sources=metadata_sources, + ) + + scan_stats.scanned_roms += 1 + scan_stats.added_roms += 1 if not rom else 0 + scan_stats.metadata_roms += 1 if scanned_rom.igdb_id or scanned_rom.moby_id else 0 + + _added_rom = db_rom_handler.add_rom(scanned_rom) + + path_cover_s, path_cover_l = await fs_resource_handler.get_cover( + overwrite=True, + entity=_added_rom, + url_cover=_added_rom.url_cover, + ) + + path_screenshots = await fs_resource_handler.get_rom_screenshots( + rom=_added_rom, + url_screenshots=_added_rom.url_screenshots, + ) + + _added_rom.path_cover_s = path_cover_s + _added_rom.path_cover_l = path_cover_l + _added_rom.path_screenshots = path_screenshots + # Update the scanned rom with the cover and screenshots paths and update database + db_rom_handler.update_rom( + _added_rom.id, + { + c: getattr(_added_rom, c) + for c in inspect(_added_rom).mapper.column_attrs.keys() + }, + ) + + await socket_manager.emit( + "scan:scanning_rom", + { + "platform_name": platform.name, + "platform_slug": platform.slug, + **RomSchema.model_validate(_added_rom).model_dump( + exclude={"created_at", "updated_at", "rom_user"} + ), + }, + ) + await socket_manager.emit("", None) + + return scan_stats + + @socket_handler.socket_server.on("scan") async def scan_handler(_sid: str, options: dict): """Scan socket endpoint diff --git a/backend/exceptions/socket_exceptions.py b/backend/exceptions/socket_exceptions.py new file mode 100644 index 000000000..dfbc6bf56 --- /dev/null +++ b/backend/exceptions/socket_exceptions.py @@ -0,0 +1 @@ +class ScanStoppedException(Exception): ... diff --git a/backend/handler/database/platforms_handler.py b/backend/handler/database/platforms_handler.py index 82a9708de..f2c0aee5f 100644 --- a/backend/handler/database/platforms_handler.py +++ b/backend/handler/database/platforms_handler.py @@ -29,7 +29,7 @@ class DBPlatformsHandler(DBBaseHandler): @with_roms def add_platform( self, platform: Platform, query: Query = None, session: Session = None - ) -> Platform | None: + ) -> Platform: platform = session.merge(platform) session.flush()