Upstream 16522 - Refactor formatted raw SQL in unified_jobs result_stdout_raw_handle (#587)

Author: cigamit

awx/main/models/unified_jobs.py

@@ -23,6 +23,9 @@
 from django.utils.encoding import smart_str
 from django.contrib.contenttypes.models import ContentType
 
+# psycopg
+from psycopg import sql
+
 # REST Framework
 from rest_framework.exceptions import ParseError
 
@@ -1141,17 +1144,23 @@ def result_stdout_raw_handle(self, enforce_max_bytes=True):
                         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)))
 

awx/main/tests/functional/conftest.py

@@ -790,14 +790,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
@@ -818,9 +817,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)