diff --git a/apps/epg/tasks.py b/apps/epg/tasks.py
index 248e9a5f..0e8eda42 100644
--- a/apps/epg/tasks.py
+++ b/apps/epg/tasks.py
@@ -561,9 +561,10 @@ def refresh_epg_data(source_id, force=False):
gc.collect()
return
- # Build byte-offset index for preview lookups in the background so refresh isn't blocked by it
- build_programme_index_task.delay(source.id)
+ # Build byte-offset index after programme data is committed so refresh
+ # does not compete for memory/IO during the programme swap.
parse_programs_for_source(source)
+ build_programme_index_task.delay(source.id)
elif source.source_type == 'schedules_direct':
fetch_schedules_direct(source, force=force)
@@ -590,6 +591,7 @@ def refresh_epg_data(source_id, force=False):
source = None
# Force garbage collection before releasing the lock
gc.collect()
+ connection.close()
lock_renewer.stop()
release_task_lock('refresh_epg_data', source_id)
@@ -1804,6 +1806,163 @@ def parse_programs_for_tvg_id(epg_id, force=False):
release_task_lock('parse_epg_programs', epg_id)
+_EPG_PROGRAM_STAGING_TABLE = 'epg_program_staging'
+# Parse batches bound Python memory during XML iterparse; swap batches bound each
+# DELETE/INSERT statement inside the single atomic swap transaction.
+_EPG_PARSE_BATCH_SIZE = 2500
+_EPG_SWAP_BATCH_SIZE = 5000
+
+
+def _epg_program_staging_supported():
+ return connection.vendor == 'postgresql'
+
+
+def _prepare_epg_program_staging_table():
+ """Create/truncate a session-scoped temp table for streaming EPG programme inserts."""
+ if not _epg_program_staging_supported():
+ return False
+
+ with connection.cursor() as cursor:
+ cursor.execute(
+ f"""
+ CREATE TEMP TABLE IF NOT EXISTS {_EPG_PROGRAM_STAGING_TABLE} (
+ epg_id bigint NOT NULL,
+ start_time timestamptz NOT NULL,
+ end_time timestamptz NOT NULL,
+ title varchar(255) NOT NULL,
+ sub_title text,
+ description text,
+ tvg_id varchar(255),
+ custom_properties jsonb
+ ) ON COMMIT PRESERVE ROWS
+ """
+ )
+ cursor.execute(f"TRUNCATE {_EPG_PROGRAM_STAGING_TABLE}")
+ return True
+
+
+def _clear_epg_program_staging_table():
+ if not _epg_program_staging_supported():
+ return
+ with connection.cursor() as cursor:
+ cursor.execute(f"TRUNCATE {_EPG_PROGRAM_STAGING_TABLE}")
+
+
+def _flush_epg_program_staging_batch(programs_batch):
+ """Insert a batch of unsaved ProgramData rows into the session staging table."""
+ if not programs_batch or not _epg_program_staging_supported():
+ return
+
+ values_sql = []
+ params = []
+ for program in programs_batch:
+ values_sql.append("(%s, %s, %s, %s, %s, %s, %s, %s)")
+ custom_properties = program.custom_properties
+ if custom_properties is not None and not isinstance(custom_properties, str):
+ custom_properties = json.dumps(custom_properties)
+ params.extend([
+ program.epg_id,
+ program.start_time,
+ program.end_time,
+ program.title,
+ program.sub_title,
+ program.description,
+ program.tvg_id,
+ custom_properties,
+ ])
+
+ with connection.cursor() as cursor:
+ cursor.execute(
+ f"""
+ INSERT INTO {_EPG_PROGRAM_STAGING_TABLE} (
+ epg_id, start_time, end_time, title, sub_title, description, tvg_id, custom_properties
+ ) VALUES {', '.join(values_sql)}
+ """,
+ params,
+ )
+
+
+def _swap_staged_epg_programs(mapped_epg_ids, epg_source, batch_size=_EPG_SWAP_BATCH_SIZE):
+ """
+ Atomically replace mapped programme rows with staged data.
+ Must be called inside transaction.atomic().
+
+ Staged rows are moved in batches (DELETE ... RETURNING + INSERT) so Postgres
+ does not need to materialize the entire catalogue in one statement.
+ """
+ with connection.cursor() as cursor:
+ cursor.execute("SET LOCAL statement_timeout = '10min'")
+
+ deleted_count = ProgramData.objects.filter(epg_id__in=mapped_epg_ids).delete()[0]
+ logger.debug(f"Deleted {deleted_count} existing programs")
+
+ unmapped_epg_ids = list(
+ EPGData.objects.filter(epg_source=epg_source)
+ .exclude(id__in=mapped_epg_ids)
+ .values_list('id', flat=True)
+ )
+ if unmapped_epg_ids:
+ orphaned_count = ProgramData.objects.filter(epg_id__in=unmapped_epg_ids).delete()[0]
+ if orphaned_count > 0:
+ logger.info(
+ f"Cleaned up {orphaned_count} orphaned programs for "
+ f"{len(unmapped_epg_ids)} unmapped EPG entries"
+ )
+
+ if not _epg_program_staging_supported():
+ raise RuntimeError('_swap_staged_epg_programs requires PostgreSQL staging support')
+
+ program_table = ProgramData._meta.db_table
+ total_inserted = 0
+ while True:
+ with connection.cursor() as cursor:
+ cursor.execute(
+ f"""
+ WITH moved AS (
+ DELETE FROM {_EPG_PROGRAM_STAGING_TABLE}
+ WHERE ctid IN (
+ SELECT ctid FROM {_EPG_PROGRAM_STAGING_TABLE} LIMIT %s
+ )
+ RETURNING
+ epg_id, start_time, end_time, title, sub_title,
+ description, tvg_id, custom_properties
+ )
+ INSERT INTO {program_table} (
+ epg_id, start_time, end_time, title, sub_title,
+ description, tvg_id, custom_properties
+ )
+ SELECT
+ epg_id, start_time, end_time, title, sub_title,
+ description, tvg_id, custom_properties
+ FROM moved
+ """,
+ [batch_size],
+ )
+ moved_count = cursor.rowcount
+ if moved_count == 0:
+ break
+ total_inserted += moved_count
+
+ logger.debug(f"Inserted {total_inserted} staged programs in batches of {batch_size}")
+
+ return deleted_count
+
+
+def _swap_parsed_epg_programs(mapped_epg_ids, epg_source, programs_to_create, batch_size=_EPG_SWAP_BATCH_SIZE):
+ """SQLite/dev fallback: atomic delete + bulk insert from an in-memory batch list."""
+ with transaction.atomic():
+ deleted_count = ProgramData.objects.filter(epg_id__in=mapped_epg_ids).delete()[0]
+ unmapped_epg_ids = list(
+ EPGData.objects.filter(epg_source=epg_source)
+ .exclude(id__in=mapped_epg_ids)
+ .values_list('id', flat=True)
+ )
+ if unmapped_epg_ids:
+ ProgramData.objects.filter(epg_id__in=unmapped_epg_ids).delete()
+ for i in range(0, len(programs_to_create), batch_size):
+ ProgramData.objects.bulk_create(programs_to_create[i:i + batch_size])
+ return deleted_count
+
def parse_programs_for_source(epg_source, tvg_id=None):
"""
@@ -1910,106 +2069,164 @@ def parse_programs_for_source(epg_source, tvg_id=None):
send_epg_update(epg_source.id, "parsing_programs", 100, status="error", error="No URL provided")
return False
- # SINGLE PASS PARSING: Parse the XML file once and collect all programs in memory
- # We parse FIRST, then do an atomic delete+insert to avoid race conditions
- # where clients might see empty/partial EPG data during the transition
- all_programs_to_create = []
- programs_by_channel = {tvg_id: 0 for tvg_id in mapped_tvg_ids} # Track count per channel
+ # Stream parsed rows into a session temp table, then swap in a short transaction.
+ # This bounds Python memory (batched staging inserts) and Postgres memory (no
+ # long-lived transaction spanning the entire XML parse).
+ programs_by_channel = {tvg_id: 0 for tvg_id in mapped_tvg_ids}
total_programs = 0
skipped_programs = 0
last_progress_update = 0
+ parse_batch_size = _EPG_PARSE_BATCH_SIZE
+ swap_batch_size = _EPG_SWAP_BATCH_SIZE
+ programs_batch = []
+ deleted_count = 0
+ staging_prepared = False
+ use_staging = False
+ programs_accumulator = []
+
+ send_epg_update(epg_source.id, "parsing_programs", 10, message="Parsing programs...")
try:
- logger.debug(f"Opening file for single-pass parsing: {file_path}")
+ staging_prepared = _prepare_epg_program_staging_table()
+ use_staging = staging_prepared
+
+ logger.debug(f"Opening file for streaming parse: {file_path}")
source_file = _open_xmltv_file(file_path)
+ try:
+ program_parser = etree.iterparse(
+ source_file,
+ events=('end',),
+ tag='programme',
+ remove_blank_text=True,
+ recover=True,
+ )
- # Stream parse the file using lxml's iterparse
- program_parser = etree.iterparse(source_file, events=('end',), tag='programme', remove_blank_text=True, recover=True)
+ for _, elem in program_parser:
+ channel_id = elem.get('channel')
- for _, elem in program_parser:
- channel_id = elem.get('channel')
+ if channel_id not in mapped_tvg_ids:
+ skipped_programs += 1
+ clear_element(elem)
+ continue
- # Skip programmes for unmapped channels immediately
- if channel_id not in mapped_tvg_ids:
- skipped_programs += 1
- # Clear element to free memory
- clear_element(elem)
- continue
+ try:
+ start_time = parse_xmltv_time(elem.get('start'))
+ end_time = parse_xmltv_time(elem.get('stop'))
+ title = None
+ desc = None
+ sub_title = None
- # This programme is for a mapped channel - process it
- try:
- start_time = parse_xmltv_time(elem.get('start'))
- end_time = parse_xmltv_time(elem.get('stop'))
- title = None
- desc = None
- sub_title = None
+ for child in elem:
+ if child.tag == 'title':
+ title = child.text or 'No Title'
+ elif child.tag == 'desc':
+ desc = child.text or ''
+ elif child.tag == 'sub-title':
+ sub_title = child.text or ''
- # Efficiently process child elements
- for child in elem:
- if child.tag == 'title':
- title = child.text or 'No Title'
- elif child.tag == 'desc':
- desc = child.text or ''
- elif child.tag == 'sub-title':
- sub_title = child.text or ''
+ if not title:
+ title = 'No Title'
- if not title:
- title = 'No Title'
+ custom_props = extract_custom_properties(elem)
+ custom_properties_json = custom_props if custom_props else None
- # Extract custom properties
- custom_props = extract_custom_properties(elem)
- custom_properties_json = custom_props if custom_props else None
+ if desc:
+ has_season = (custom_properties_json or {}).get('season') is not None
+ has_episode = (custom_properties_json or {}).get('episode') is not None
+ if not has_season or not has_episode:
+ d_season, d_episode, cleaned_desc = extract_season_episode_from_description(desc)
+ if d_season is not None and d_episode is not None:
+ if custom_properties_json is None:
+ custom_properties_json = {}
+ if not has_season:
+ custom_properties_json['season'] = d_season
+ if not has_episode:
+ custom_properties_json['episode'] = d_episode
+ custom_properties_json['season_episode_source'] = 'description'
+ desc = cleaned_desc
- # Fallback: extract S/E from description when episode-num
- # elements didn't provide them
- if desc:
- has_season = (custom_properties_json or {}).get('season') is not None
- has_episode = (custom_properties_json or {}).get('episode') is not None
- if not has_season or not has_episode:
- d_season, d_episode, cleaned_desc = extract_season_episode_from_description(desc)
- if d_season is not None and d_episode is not None:
- if custom_properties_json is None:
- custom_properties_json = {}
- if not has_season:
- custom_properties_json['season'] = d_season
- if not has_episode:
- custom_properties_json['episode'] = d_episode
- custom_properties_json['season_episode_source'] = 'description'
- desc = cleaned_desc
+ epg_id = tvg_id_to_epg_id[channel_id]
+ programs_batch.append(ProgramData(
+ epg_id=epg_id,
+ start_time=start_time,
+ end_time=end_time,
+ title=title[:255],
+ description=desc,
+ sub_title=sub_title,
+ tvg_id=channel_id,
+ custom_properties=custom_properties_json,
+ ))
+ total_programs += 1
+ programs_by_channel[channel_id] += 1
+ clear_element(elem)
- epg_id = tvg_id_to_epg_id[channel_id]
- all_programs_to_create.append(ProgramData(
- epg_id=epg_id,
- start_time=start_time,
- end_time=end_time,
- title=title[:255],
- description=desc,
- sub_title=sub_title,
- tvg_id=channel_id,
- custom_properties=custom_properties_json
- ))
- total_programs += 1
- programs_by_channel[channel_id] += 1
+ if len(programs_batch) >= parse_batch_size:
+ if use_staging:
+ _flush_epg_program_staging_batch(programs_batch)
+ programs_batch = []
+ else:
+ programs_accumulator.extend(programs_batch)
+ programs_batch = []
- # Clear the element to free memory
- clear_element(elem)
+ if total_programs - last_progress_update >= 5000:
+ last_progress_update = total_programs
+ progress = min(
+ 85,
+ 10 + int((total_programs / max(total_programs + 10000, 1)) * 75),
+ )
+ send_epg_update(
+ epg_source.id,
+ "parsing_programs",
+ progress,
+ processed=total_programs,
+ channels=mapped_count,
+ message=f"Staging programs... {total_programs:,}",
+ )
- # Send progress update (estimate based on programs processed)
- if total_programs - last_progress_update >= 5000:
- last_progress_update = total_programs
- # Cap at 70% during parsing phase (save 30% for DB operations)
- progress = min(70, 10 + int((total_programs / max(total_programs + 10000, 1)) * 60))
- send_epg_update(epg_source.id, "parsing_programs", progress,
- processed=total_programs, channels=mapped_count)
+ if total_programs % 5000 == 0:
+ gc.collect()
- # Periodic garbage collection during parsing
- if total_programs % 5000 == 0:
- gc.collect()
+ except Exception as e:
+ logger.error(f"Error processing program for {channel_id}: {e}", exc_info=True)
+ clear_element(elem)
+ continue
- except Exception as e:
- logger.error(f"Error processing program for {channel_id}: {e}", exc_info=True)
- clear_element(elem)
- continue
+ if programs_batch:
+ if use_staging:
+ _flush_epg_program_staging_batch(programs_batch)
+ else:
+ programs_accumulator.extend(programs_batch)
+ programs_batch = []
+ finally:
+ if source_file:
+ source_file.close()
+ source_file = None
+
+ try:
+ send_epg_update(epg_source.id, "parsing_programs", 90, message="Updating database...")
+ if use_staging:
+ with transaction.atomic():
+ deleted_count = _swap_staged_epg_programs(
+ mapped_epg_ids, epg_source, batch_size=swap_batch_size
+ )
+ else:
+ deleted_count = _swap_parsed_epg_programs(
+ mapped_epg_ids, epg_source, programs_accumulator, batch_size=swap_batch_size
+ )
+ programs_accumulator = []
+
+ logger.info(
+ f"Atomic update complete: deleted {deleted_count}, inserted {total_programs} programs"
+ )
+ except Exception as db_error:
+ logger.error(f"Database error during atomic update: {db_error}", exc_info=True)
+ epg_source.status = EPGSource.STATUS_ERROR
+ epg_source.last_message = f"Database error: {str(db_error)}"
+ epg_source.save(update_fields=['status', 'last_message'])
+ send_epg_update(
+ epg_source.id, "parsing_programs", 100, status="error", message=str(db_error)
+ )
+ return False
except etree.XMLSyntaxError as xml_error:
logger.error(f"XML syntax error parsing program data: {xml_error}")
@@ -2018,63 +2235,23 @@ def parse_programs_for_source(epg_source, tvg_id=None):
epg_source.save(update_fields=['status', 'last_message'])
send_epg_update(epg_source.id, "parsing_programs", 100, status="error", message=str(xml_error))
return False
- except Exception as e:
- logger.error(f"Error parsing XML for programs: {e}", exc_info=True)
- raise
- finally:
- if source_file:
- source_file.close()
- source_file = None
-
- # Now perform atomic delete + bulk insert
- # This ensures clients never see empty/partial EPG data
- logger.info(f"Parsed {total_programs} programs, performing atomic database update...")
- send_epg_update(epg_source.id, "parsing_programs", 75, message="Updating database...")
-
- batch_size = 1000
- try:
- with transaction.atomic():
- # Kill any individual statement that hangs longer than 10 minutes.
- # SET LOCAL automatically resets when this transaction ends (commit or rollback).
- with connection.cursor() as cursor:
- cursor.execute("SET LOCAL statement_timeout = '10min'")
- # Delete existing programs for mapped EPGs
- deleted_count = ProgramData.objects.filter(epg_id__in=mapped_epg_ids).delete()[0]
- logger.debug(f"Deleted {deleted_count} existing programs")
-
- # Clean up orphaned programs for unmapped EPG entries
- unmapped_epg_ids = list(EPGData.objects.filter(
- epg_source=epg_source
- ).exclude(id__in=mapped_epg_ids).values_list('id', flat=True))
-
- if unmapped_epg_ids:
- orphaned_count = ProgramData.objects.filter(epg_id__in=unmapped_epg_ids).delete()[0]
- if orphaned_count > 0:
- logger.info(f"Cleaned up {orphaned_count} orphaned programs for {len(unmapped_epg_ids)} unmapped EPG entries")
-
- # Bulk insert all new programs in batches within the same transaction
- for i in range(0, len(all_programs_to_create), batch_size):
- batch = all_programs_to_create[i:i + batch_size]
- ProgramData.objects.bulk_create(batch)
-
- # Update progress during insertion
- progress = 75 + int((i / len(all_programs_to_create)) * 20) if all_programs_to_create else 95
- if i % (batch_size * 5) == 0:
- send_epg_update(epg_source.id, "parsing_programs", min(95, progress),
- message=f"Inserting programs... {i}/{len(all_programs_to_create)}")
-
- logger.info(f"Atomic update complete: deleted {deleted_count}, inserted {total_programs} programs")
-
- except Exception as db_error:
- logger.error(f"Database error during atomic update: {db_error}", exc_info=True)
+ except Exception as parse_error:
+ logger.error(f"Error parsing programs from XML: {parse_error}", exc_info=True)
epg_source.status = EPGSource.STATUS_ERROR
- epg_source.last_message = f"Database error: {str(db_error)}"
+ epg_source.last_message = f"Error parsing programs: {str(parse_error)}"
epg_source.save(update_fields=['status', 'last_message'])
- send_epg_update(epg_source.id, "parsing_programs", 100, status="error", message=str(db_error))
+ send_epg_update(
+ epg_source.id, "parsing_programs", 100, status="error", message=str(parse_error)
+ )
return False
finally:
- # Clear the large list to free memory
- all_programs_to_create = None
+ programs_batch = None
+ programs_accumulator = None
+ if staging_prepared:
+ try:
+ _clear_epg_program_staging_table()
+ except Exception:
+ pass
gc.collect()
# Count channels that actually got programs
@@ -2130,7 +2307,7 @@ def parse_programs_for_source(epg_source, tvg_id=None):
source_file = None
# Explicitly release any remaining large data structures
- programs_to_create = None
+ programs_batch = None
programs_by_channel = None
mapped_epg_ids = None
mapped_tvg_ids = None
diff --git a/apps/epg/tests/test_parse_programs_for_source.py b/apps/epg/tests/test_parse_programs_for_source.py
new file mode 100644
index 00000000..f49a5d70
--- /dev/null
+++ b/apps/epg/tests/test_parse_programs_for_source.py
@@ -0,0 +1,238 @@
+import os
+import tempfile
+from datetime import timedelta
+from unittest.mock import patch
+
+from django.db import connection, transaction
+from django.test import TestCase
+from django.utils import timezone
+
+from apps.channels.models import Channel
+from apps.epg.models import EPGSource, EPGData, ProgramData
+from apps.epg.tasks import (
+ parse_programs_for_source,
+ _flush_epg_program_staging_batch,
+ _swap_staged_epg_programs,
+ _EPG_PARSE_BATCH_SIZE,
+)
+
+
+def _programme_xml(channel_id, title, start, stop):
+ return (
+ f' \n'
+ f' {title}\n'
+ f' \n'
+ )
+
+
+def _xmltv_file(programmes):
+ body = (
+ '\n'
+ '\n'
+ f'{programmes}'
+ '\n'
+ )
+ handle = tempfile.NamedTemporaryFile(
+ mode='w',
+ suffix='.xml',
+ delete=False,
+ encoding='utf-8',
+ )
+ handle.write(body)
+ handle.close()
+ return handle.name
+
+
+class ParseProgramsForSourceTests(TestCase):
+ def setUp(self):
+ self.source = EPGSource.objects.create(
+ name='XMLTV Parse Test',
+ source_type='xmltv',
+ )
+ self.mapped_epg = EPGData.objects.create(
+ epg_source=self.source,
+ tvg_id='mapped.channel',
+ name='Mapped Channel',
+ )
+ self.unmapped_epg = EPGData.objects.create(
+ epg_source=self.source,
+ tvg_id='unmapped.channel',
+ name='Unmapped Channel',
+ )
+ Channel.objects.create(
+ channel_number=1,
+ name='Mapped Channel',
+ epg_data=self.mapped_epg,
+ )
+ self.base_time = timezone.now().replace(minute=0, second=0, microsecond=0)
+ self.start = self.base_time.strftime('%Y%m%d%H%M%S +0000')
+ self.stop = (self.base_time + timedelta(hours=1)).strftime('%Y%m%d%H%M%S +0000')
+
+ def tearDown(self):
+ if getattr(self, 'xml_path', None) and os.path.exists(self.xml_path):
+ os.unlink(self.xml_path)
+
+ def _configure_source_file(self, programmes):
+ self.xml_path = _xmltv_file(programmes)
+ self.source.file_path = self.xml_path
+ self.source.save(update_fields=['file_path'])
+
+ @patch('apps.epg.tasks.log_system_event')
+ @patch('apps.epg.tasks.send_epg_update')
+ def test_replaces_programs_for_mapped_channels(self, _send_update, _log_event):
+ old_start = self.base_time - timedelta(days=1)
+ ProgramData.objects.create(
+ epg=self.mapped_epg,
+ start_time=old_start,
+ end_time=old_start + timedelta(hours=1),
+ title='Old Programme',
+ tvg_id=self.mapped_epg.tvg_id,
+ )
+ orphan_start = self.base_time - timedelta(days=1)
+ ProgramData.objects.create(
+ epg=self.unmapped_epg,
+ start_time=orphan_start,
+ end_time=orphan_start + timedelta(hours=1),
+ title='Orphan Programme',
+ tvg_id=self.unmapped_epg.tvg_id,
+ )
+
+ programmes = (
+ _programme_xml('mapped.channel', 'New Show', self.start, self.stop)
+ + _programme_xml('unmapped.channel', 'Skipped Show', self.start, self.stop)
+ )
+ self._configure_source_file(programmes)
+
+ result = parse_programs_for_source(self.source)
+
+ self.assertTrue(result)
+ mapped_programs = ProgramData.objects.filter(epg=self.mapped_epg)
+ self.assertEqual(mapped_programs.count(), 1)
+ self.assertEqual(mapped_programs.get().title, 'New Show')
+ self.assertFalse(ProgramData.objects.filter(epg=self.unmapped_epg).exists())
+
+ @patch('apps.epg.tasks.log_system_event')
+ @patch('apps.epg.tasks.send_epg_update')
+ def test_atomic_failure_rolls_back_and_preserves_existing_programs(self, _send_update, _log_event):
+ old_start = self.base_time - timedelta(days=1)
+ ProgramData.objects.create(
+ epg=self.mapped_epg,
+ start_time=old_start,
+ end_time=old_start + timedelta(hours=1),
+ title='Keep Me',
+ tvg_id=self.mapped_epg.tvg_id,
+ )
+
+ self._configure_source_file(
+ _programme_xml('mapped.channel', 'Replacement', self.start, self.stop)
+ )
+
+ swap_path = (
+ 'apps.epg.tasks._swap_staged_epg_programs'
+ if connection.vendor == 'postgresql'
+ else 'apps.epg.tasks._swap_parsed_epg_programs'
+ )
+ with patch(swap_path, side_effect=RuntimeError('simulated insert failure')):
+ result = parse_programs_for_source(self.source)
+
+ self.assertFalse(result)
+ self.assertEqual(ProgramData.objects.filter(epg=self.mapped_epg).count(), 1)
+ self.assertEqual(
+ ProgramData.objects.get(epg=self.mapped_epg).title,
+ 'Keep Me',
+ )
+
+ @patch('apps.epg.tasks.log_system_event')
+ @patch('apps.epg.tasks.send_epg_update')
+ def test_streams_batches_without_holding_full_program_list(self, _send_update, _log_event):
+ if connection.vendor != 'postgresql':
+ self.skipTest('PostgreSQL staging batches are required for this assertion')
+
+ programme_count = _EPG_PARSE_BATCH_SIZE * 2
+ programmes = ''.join(
+ _programme_xml(
+ 'mapped.channel',
+ f'Show {idx}',
+ self.start,
+ self.stop,
+ )
+ for idx in range(programme_count)
+ )
+ self._configure_source_file(programmes)
+ flush_sizes = []
+ original_flush = _flush_epg_program_staging_batch
+
+ def tracking_flush(batch):
+ flush_sizes.append(len(batch))
+ return original_flush(batch)
+
+ with patch('apps.epg.tasks._flush_epg_program_staging_batch', side_effect=tracking_flush):
+ result = parse_programs_for_source(self.source)
+
+ self.assertTrue(result)
+ self.assertEqual(ProgramData.objects.filter(epg=self.mapped_epg).count(), programme_count)
+ self.assertEqual(sum(flush_sizes), programme_count)
+ self.assertTrue(all(size <= _EPG_PARSE_BATCH_SIZE for size in flush_sizes))
+ self.assertGreater(len(flush_sizes), 1)
+
+ @patch('apps.epg.tasks.log_system_event')
+ @patch('apps.epg.tasks.send_epg_update')
+ def test_live_programs_remain_until_swap_commits(self, _send_update, _log_event):
+ if connection.vendor != 'postgresql':
+ self.skipTest('PostgreSQL staging swap is required for this assertion')
+
+ old_start = self.base_time - timedelta(days=1)
+ ProgramData.objects.create(
+ epg=self.mapped_epg,
+ start_time=old_start,
+ end_time=old_start + timedelta(hours=1),
+ title='Old Programme',
+ tvg_id=self.mapped_epg.tvg_id,
+ )
+ self._configure_source_file(
+ _programme_xml('mapped.channel', 'New Show', self.start, self.stop)
+ )
+
+ observed_titles_at_swap = []
+
+ def swap_with_visibility_check(mapped_epg_ids, epg_source, *args, **kwargs):
+ observed_titles_at_swap.append(
+ ProgramData.objects.get(epg=self.mapped_epg).title
+ )
+ return _swap_staged_epg_programs(mapped_epg_ids, epg_source, *args, **kwargs)
+
+ with patch('apps.epg.tasks._swap_staged_epg_programs', side_effect=swap_with_visibility_check):
+ result = parse_programs_for_source(self.source)
+
+ self.assertTrue(result)
+ self.assertEqual(observed_titles_at_swap, ['Old Programme'])
+ self.assertEqual(ProgramData.objects.get(epg=self.mapped_epg).title, 'New Show')
+
+ @patch('apps.epg.tasks.log_system_event')
+ @patch('apps.epg.tasks.send_epg_update')
+ def test_swap_delete_is_rolled_back_when_insert_fails(self, _send_update, _log_event):
+ if connection.vendor != 'postgresql':
+ self.skipTest('PostgreSQL staging swap is required for this assertion')
+
+ old_start = self.base_time - timedelta(days=1)
+ ProgramData.objects.create(
+ epg=self.mapped_epg,
+ start_time=old_start,
+ end_time=old_start + timedelta(hours=1),
+ title='Keep Me',
+ tvg_id=self.mapped_epg.tvg_id,
+ )
+ self._configure_source_file(
+ _programme_xml('mapped.channel', 'Replacement', self.start, self.stop)
+ )
+
+ def failing_swap(mapped_epg_ids, epg_source, *args, **kwargs):
+ with transaction.atomic():
+ ProgramData.objects.filter(epg_id__in=mapped_epg_ids).delete()
+ raise RuntimeError('simulated insert failure')
+
+ with patch('apps.epg.tasks._swap_staged_epg_programs', side_effect=failing_swap):
+ result = parse_programs_for_source(self.source)
+
+ self.assertFalse(result)
+ self.assertEqual(ProgramData.objects.get(epg=self.mapped_epg).title, 'Keep Me')
diff --git a/dispatcharr/celery.py b/dispatcharr/celery.py
index 7ce0ab29..bdb5e48d 100644
--- a/dispatcharr/celery.py
+++ b/dispatcharr/celery.py
@@ -94,6 +94,7 @@ def cleanup_task_memory(**kwargs):
'apps.epg.tasks.refresh_all_epg_data',
'apps.epg.tasks.parse_programs_for_source',
'apps.epg.tasks.parse_programs_for_tvg_id',
+ 'apps.epg.tasks.build_programme_index_task',
'apps.channels.tasks.match_epg_channels',
'apps.channels.tasks.match_selected_channels_epg',
'apps.channels.tasks.match_single_channel_epg',