misc: Refactor scan process by splitting single function

This change mainly refactors the `scan_platforms` function, moving part
of its logic to `_identify_platform`, `_identify_firmware`, and
`_identify_rom`.

The logic is simpler this way, each smaller function returns `ScanStats`
that can be merged by the caller, and it simplifies future performance
improvements.
This commit is contained in:
Michael Manganiello
2024-07-25 10:24:32 -03:00
parent e7e37f4680
commit 2bfc3c4c7e
3 changed files with 233 additions and 168 deletions

View File

@@ -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

View File

@@ -0,0 +1 @@
class ScanStoppedException(Exception): ...

View File

@@ -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()