diff --git a/awx/main/models/unified_jobs.py b/awx/main/models/unified_jobs.py index 2a651ffa1c..7b5b74d632 100644 --- a/awx/main/models/unified_jobs.py +++ b/awx/main/models/unified_jobs.py @@ -20,6 +20,9 @@ from dispatcherd.factories import get_control_from_settings # Django from django.conf import settings from django.db import models, connection, transaction + +# psycopg +from psycopg import sql from django.db.models.constraints import UniqueConstraint from django.core.exceptions import NON_FIELD_ERRORS from django.utils.translation import gettext_lazy as _ @@ -1179,17 +1182,23 @@ class UnifiedJob( raise StdoutMaxBytesExceeded(total, max_supported) tbl = self._meta.db_table + 'event' - created_by_cond = '' + where_parts = [ + sql.SQL('{} = {}').format(sql.Identifier(self.event_parent_key), sql.Literal(self.id)), + sql.SQL("stdout != ''"), + ] if self.has_unpartitioned_events: - tbl = f'_unpartitioned_{tbl}' + tbl = '_unpartitioned_' + tbl else: - created_by_cond = f"job_created='{self.created.isoformat()}' AND " + where_parts.insert(0, sql.SQL('job_created = {}').format(sql.Literal(self.created))) - sql = f"copy (select stdout from {tbl} where {created_by_cond}{self.event_parent_key}={self.id} and stdout != '' order by start_line) to stdout" # nosql + copy_sql = sql.SQL('COPY (SELECT stdout FROM {} WHERE {} ORDER BY start_line) TO STDOUT').format( + sql.Identifier(tbl), + sql.SQL(' AND ').join(where_parts), + ) # psycopg3's copy writes bytes, but callers of this # function assume a str-based fd will be returned; decode # .write() calls on the fly to maintain this interface - with cursor.copy(sql) as copy: + with cursor.copy(copy_sql) as copy: while data := copy.read(): fd.write(smart_str(bytes(data))) diff --git a/awx/main/tests/functional/conftest.py b/awx/main/tests/functional/conftest.py index 8f21861120..bd3e9a0178 100644 --- a/awx/main/tests/functional/conftest.py +++ b/awx/main/tests/functional/conftest.py @@ -830,14 +830,13 @@ class MockCopy: events = [] index = -1 - def __init__(self, sql): + def __init__(self): self.events = [] - parts = sql.split(' ') - tablename = parts[parts.index('from') + 1] for cls in (JobEvent, AdHocCommandEvent, ProjectUpdateEvent, InventoryUpdateEvent, SystemJobEvent): - if cls._meta.db_table == tablename: - for event in cls.objects.order_by('start_line').all(): - self.events.append(event.stdout) + events = list(cls.objects.order_by('start_line').values_list('stdout', flat=True)) + if events: + self.events = events + break def read(self): self.index = self.index + 1 @@ -858,9 +857,8 @@ def sqlite_copy(request, mocker): # copy is postgres-specific, and SQLite doesn't support it; mock its # behavior to test that it writes a file that contains stdout from events - def write_stdout(self, sql): - mock_copy = MockCopy(sql) - return mock_copy + def write_stdout(self, sql, params=None): + return MockCopy() mocker.patch.object(SQLiteCursorWrapper, 'copy', write_stdout, create=True)