From d7a02b5735d469fa710f5e6c3f7bbd847e4a16fc Mon Sep 17 00:00:00 2001 From: Roland Date: Fri, 5 Dec 2025 19:49:39 -0700 Subject: [PATCH] Speeding up channels config import. --- modules/archive/lib.py | 29 +++++--- modules/videos/__init__.py | 29 ++++++++ modules/videos/lib.py | 90 ++++++++++++++++++++--- modules/videos/models.py | 63 ++++++++++++++++ wrolpi/collections/models.py | 139 +++++++++++++++++++++++++++++++++++ 5 files changed, 331 insertions(+), 19 deletions(-) diff --git a/modules/archive/lib.py b/modules/archive/lib.py index b37f91f1..7fe31195 100644 --- a/modules/archive/lib.py +++ b/modules/archive/lib.py @@ -87,7 +87,7 @@ class DomainsConfig(ConfigFile): return self._config.get('collections', []) def import_config(self, file: pathlib.Path = None, send_events=False): - """Import domain collections from config file into database.""" + """Import domain collections from config file into database using batch operations.""" ConfigFile.import_config(self, file, send_events) file_str = str(self.get_relative_file()) @@ -102,10 +102,8 @@ class DomainsConfig(ConfigFile): try: with get_db_session(commit=True) as session: - # Track imported domain names - imported_domains = set() - - # Import each domain collection, forcing kind='domain' + # Validate and prepare data first + valid_data_list = [] for idx, collection_data in enumerate(collections_data): try: name = collection_data.get('name') @@ -113,6 +111,14 @@ class DomainsConfig(ConfigFile): logger.error(f'Domain collection at index {idx} has no name, skipping') continue + # Validate domain name format before batching + if not Collection.is_valid_domain_name(name): + logger.error( + f"Domain collection at index {idx} has invalid name '{name}' " + f"(must contain at least one '.' and not start/end with '.'), skipping" + ) + continue + # Ensure kind is 'domain' collection_data = collection_data.copy() collection_data['kind'] = 'domain' @@ -127,14 +133,17 @@ class DomainsConfig(ConfigFile): ) collection_data.pop('tag_name', None) - # Use Collection.from_config to create/update - collection = Collection.from_config(collection_data, session) - imported_domains.add(collection.name) - + valid_data_list.append(collection_data) except Exception as e: - logger.error(f'Failed to import domain collection at index {idx}', exc_info=e) + logger.error(f'Failed to validate domain collection at index {idx}', exc_info=e) continue + # Batch create/update collections + imported_domains = set() + if valid_data_list: + collections = Collection.batch_from_config(valid_data_list, session=session) + imported_domains = {c.name for c in collections} + # Delete domain collections that are no longer in config all_domain_collections = session.query(Collection).filter_by(kind='domain').all() for collection in all_domain_collections: diff --git a/modules/videos/__init__.py b/modules/videos/__init__.py index e1c6de42..92fb4d09 100644 --- a/modules/videos/__init__.py +++ b/modules/videos/__init__.py @@ -91,3 +91,32 @@ def video_cleanup(): WHERE v.channel_id IS NULL AND fg.id = v.file_group_id ''') + + +def claim_videos_for_channels(channel_ids: List[int]): + """ + Assign Videos to specific Channels based on directory paths. + + This is a targeted version of video_cleanup() that only processes specific channel IDs. + More efficient than video_cleanup() for channel import because it limits the scope + of the SQL update to only the channels that were just imported/updated. + + Args: + channel_ids: List of channel IDs to process + """ + if not channel_ids: + return + + logger.info(f'Claiming Videos for {len(channel_ids)} channel(s)') + with get_db_curs(commit=True) as curs: + # Note: %% is escaped to % in psycopg2 when parameters are used + curs.execute(''' + UPDATE video v + SET channel_id = c.id + FROM channel c + INNER JOIN collection col ON col.id = c.collection_id + LEFT JOIN file_group fg ON fg.primary_path LIKE col.directory || '/%%'::VARCHAR + WHERE v.channel_id IS NULL + AND fg.id = v.file_group_id + AND c.id = ANY(%s) + ''', (list(channel_ids),)) diff --git a/modules/videos/lib.py b/modules/videos/lib.py index c6362dd0..4d543786 100644 --- a/modules/videos/lib.py +++ b/modules/videos/lib.py @@ -252,29 +252,97 @@ class ChannelsConfig(ConfigFile): self.update({'channels': value}) def import_config(self, file: pathlib.Path = None, send_events=False): + """Import channels from config file into database using batch operations.""" super().import_config() try: channels = self.channels + if not channels: + channel_import_logger.info('No channels to import') + self.successful_import = True + return + channel_directories = [i['directory'] for i in channels] if len(channel_directories) != len(set(channel_directories)): raise RuntimeError('Refusing to import channels config because it contains duplicate directories!') + + # Validate all channels have required options + for data in channels: + directory = get_media_directory() / data['directory'] + for option in (i for i in REQUIRED_OPTIONS if i not in data): + raise ConfigError(f'Channel "{directory}" is required to have "{option}"') + updated_channel_ids = set() with get_db_session(commit=True) as session: + # Prepare data with kind='channel' for all entries + channel_data_list = [] for data in channels: - # A Channel's directory is saved (in config) relative to the media directory. - directory = get_media_directory() / data['directory'] - for option in (i for i in REQUIRED_OPTIONS if i not in data): - raise ConfigError(f'Channel "{directory}" is required to have "{option}"') + data = data.copy() + data['kind'] = 'channel' + channel_data_list.append(data) - # Create or update Channel from config (handles Collection creation) - channel = Channel.from_config(data, session=session) + # Batch create/update Collections + from wrolpi.collections import Collection + collections = Collection.batch_from_config(channel_data_list, session=session) + + # Pre-fetch existing Channels by collection_id in ONE query + collection_ids = [c.id for c in collections] + existing_channels_list = session.query(Channel).filter( + Channel.collection_id.in_(collection_ids) + ).all() + existing_channels = {ch.collection_id: ch for ch in existing_channels_list} + + # Pre-fetch all tags for channel update in ONE query + tag_names = {d.get('tag_name') for d in channels if d.get('tag_name')} + tags_by_name = {} + if tag_names: + from wrolpi.tags import Tag + tags = session.query(Tag).filter(Tag.name.in_(tag_names)).all() + tags_by_name = {t.name: t for t in tags} + + # Pre-fetch existing downloads for update in ONE query + download_urls = set() + for data in channels: + for dl in data.get('downloads', []): + if isinstance(dl, dict): + download_urls.add(dl['url']) + existing_downloads = {} + if download_urls: + from wrolpi.downloader import Download + downloads = session.query(Download).filter(Download.url.in_(download_urls)).all() + existing_downloads = {d.url: d for d in downloads} + + # Process channels + new_channels = [] + for collection, data in zip(collections, channel_data_list): + channel = existing_channels.get(collection.id) + + if not channel: + # Create new Channel + channel = Channel( + collection_id=collection.id, + url=data.get('url'), + source_id=data.get('source_id'), + ) + new_channels.append(channel) + existing_channels[collection.id] = channel + + # Batch add new channels + if new_channels: + session.add_all(new_channels) + session.flush() + + # Now update all channels (both new and existing) + for collection, data in zip(collections, channel_data_list): + channel = existing_channels[collection.id] + # Call batch_update with pre-fetched data + channel.batch_update(data, tags_by_name, existing_downloads) updated_channel_ids.add(channel.id) + # Background task for missing source_id if not channel.source_id and channel.url and flags.have_internet.is_set(): - # If we can download from a channel, we must have its source_id. if download_manager.can_download and channel.download_missing_data: - logger.info(f'Fetching channel source id for {channel}') + channel_import_logger.info(f'Fetching channel source id for {channel}') background_task(fetch_channel_source_id(channel.id)) channel_import_logger.debug(f'Updated {repr(channel.name)}' @@ -290,7 +358,11 @@ class ChannelsConfig(ConfigFile): logger.warning(f'Deleting {channel} because it is not in the config.') channel.delete_with_videos() - # Create any missing Downloads. Associated Downloads with any necessary Channels. + # Assign Videos to Channels using bulk SQL + from modules.videos import claim_videos_for_channels + claim_videos_for_channels(list(updated_channel_ids)) + + # Create any missing Downloads. Associate Downloads with any necessary Channels. with get_db_session() as session: link_channel_and_downloads(session) diff --git a/modules/videos/models.py b/modules/videos/models.py index b4a25b5c..e714db3f 100644 --- a/modules/videos/models.py +++ b/modules/videos/models.py @@ -700,6 +700,69 @@ class Channel(ModelHelper, Base): download.collection_id = self.collection_id session.add(download) + def batch_update(self, data: dict, tags_by_name: dict = None, existing_downloads: dict = None): + """ + Update the attributes of this Channel using pre-fetched lookup dictionaries. + More efficient than update() when processing many channels because it avoids + individual database queries for tags and downloads. + + Args: + data: Config data dict with channel fields + tags_by_name: Pre-fetched {tag_name: Tag} lookup dictionary + existing_downloads: Pre-fetched {url: Download} lookup dictionary + """ + data = data.copy() + tags_by_name = tags_by_name or {} + existing_downloads = existing_downloads or {} + + # URL should not be empty string. + url = data.pop('url', None) + self.url = url or None + + session: Session = Session.object_session(self) + downloads = data.pop('downloads', []) + + # Use pre-fetched tag lookup instead of querying + if tag_name := data.pop('tag_name', None): + tag = tags_by_name.get(tag_name) + if tag: + self.tag = tag + else: + logger.warning(f'Tag {repr(tag_name)} not found for channel') + + for key, value in data.items(): + setattr(self, key, value) + + # We need an absolute directory. + if isinstance(self.directory, pathlib.Path) and not self.directory.is_absolute(): + self.directory = get_media_directory() / self.directory + elif isinstance(self.directory, str) and not pathlib.Path(self.directory).is_absolute(): + self.directory = get_media_directory() / self.directory + + for download in downloads: + url = None + frequency = None + if isinstance(download, dict): + url = download['url'] + # Use pre-fetched download lookup instead of querying + if existing_download := existing_downloads.get(url): + frequency = existing_download.frequency + else: + frequency = download.get('frequency') + elif isinstance(download, Download): + url = download.url + frequency = download.frequency + + if not url: + raise RuntimeError(f'Unknown download type: {download}') + if not frequency: + logger.error(f'Refusing to create Download for Channel without frequency: {url}') + continue + + download_obj = self.get_or_create_download(url, frequency, session=session, reset_attempts=True) + download_obj.collection_id = self.collection_id + session.add(download_obj) + def config_view(self) -> dict: """ Retrieve the data about this Channel that should be stored in a config file. diff --git a/wrolpi/collections/models.py b/wrolpi/collections/models.py index 8d50a88b..e8fd1cd6 100644 --- a/wrolpi/collections/models.py +++ b/wrolpi/collections/models.py @@ -313,6 +313,145 @@ class Collection(ModelHelper, Base): return collection + @staticmethod + @optional_session + def batch_from_config(data_list: List[dict], session: Session = None) -> List['Collection']: + """ + Create or update multiple Collections from config data in batch. + More efficient than calling from_config() in a loop because it pre-fetches + existing collections and tags in bulk queries. + + NOTE: This method does NOT call populate_from_directory(). For channel collections, + use claim_videos_for_channels() after batch_from_config() to assign videos to channels. + + Args: + data_list: List of config dicts containing collection metadata + session: Database session + + Returns: + List of created or updated Collections + """ + if not data_list: + return [] + + # Step 1: Validate and prepare directories + directories = [] + for data in data_list: + directory = data.get('directory') + if directory: + directories.append(str(validate_collection_directory(directory))) + + # Step 2: Pre-fetch existing collections by directory in ONE query + existing_by_directory = {} + if directories: + existing = session.query(Collection).filter( + Collection.directory.in_(directories) + ).all() + existing_by_directory = {str(c.directory): c for c in existing} + + # Step 3: Pre-fetch existing collections by (name, kind) for directory-less collections + name_kind_pairs = [(d['name'], d.get('kind', 'channel')) for d in data_list if not d.get('directory')] + existing_by_name_kind = {} + if name_kind_pairs: + from sqlalchemy import and_, or_ + conditions = [ + and_(Collection.name == name, Collection.kind == kind) + for name, kind in name_kind_pairs + ] + existing = session.query(Collection).filter(or_(*conditions)).all() + existing_by_name_kind = {(c.name, c.kind): c for c in existing} + + # Step 4: Pre-fetch all tags in ONE query + tag_names = {d['tag_name'] for d in data_list if d.get('tag_name')} + tags_by_name = {} + if tag_names: + tags = session.query(Tag).filter(Tag.name.in_(tag_names)).all() + tags_by_name = {t.name: t for t in tags} + + # Step 5: Process collections + collections = [] + new_collections = [] + + for data in data_list: + name = data.get('name') + if not name: + raise ValueError('Collection config must have a name') + + description = data.get('description') + directory = data.get('directory') + tag_name = data.get('tag_name') + kind = (data.get('kind') or 'channel').strip() + + # Validate known kinds + if kind not in {'channel', 'domain'}: + logger.warning(f"Unknown collection kind '{kind}', defaulting to 'channel'") + kind = 'channel' + + # Validate collection name using the type registry + if not collection_type_registry.validate(kind, name): + description_msg = collection_type_registry.get_description(kind) or "Invalid name format" + raise ValueError(f'Invalid {kind} name for collection: {repr(name)}. {description_msg}') + + # Convert directory to absolute path if provided + if directory: + directory = validate_collection_directory(directory) + directory_str = str(directory) + else: + directory_str = None + + # Find existing collection from pre-fetched data + collection = None + if directory_str: + collection = existing_by_directory.get(directory_str) + else: + collection = existing_by_name_kind.get((name, kind)) + + if collection: + # Update existing collection + logger.debug(f'Updating collection from config: {name}') + collection.description = description + collection.directory = directory + collection.kind = kind + + if tag_name: + tag = tags_by_name.get(tag_name) + if tag: + collection.tag = tag + else: + logger.warning(f'Tag {repr(tag_name)} not found for collection {repr(name)}') + collection.tag = None + else: + collection.tag = None + else: + # Create new collection + logger.info(f'Creating new collection from config: {name}') + collection = Collection( + name=name, + description=description, + directory=directory, + kind=kind, + ) + + if tag_name: + tag = tags_by_name.get(tag_name) + if tag: + collection.tag = tag + else: + logger.warning(f'Tag {repr(tag_name)} not found for collection {repr(name)}') + + new_collections.append(collection) + + collections.append(collection) + + # Step 6: Batch add new collections + if new_collections: + session.add_all(new_collections) + + # Single flush for all collections + session.flush() + + return collections + def populate_from_directory(self, session: Session = None): """ Populate this collection with all FileGroups in the collection's directory.