Escape unprintable ASCII characters in received SQL queries

Author: dolfinus

data_rentgen/consumer/extractors/generic/operation.py

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import codecs
+import re
 from abc import ABC, abstractmethod
 from textwrap import dedent
 
@@ -17,6 +19,13 @@
     OpenLineageRunEventType,
 )
 
+# https://www.ascii-code.com/, but left \n intact
+ASCII_UNPRINTABLE = re.compile(r"[\x00-\x09\x0b-\x1f\x7f]", re.UNICODE)
+
+
+def encode_char(char: re.Match[str]) -> str:
+    return codecs.encode(char.group(0), "unicode-escape").decode("utf-8")
+
 
 class OperationExtractorMixin(ABC):
     @abstractmethod
@@ -74,5 +83,7 @@ def _enrich_operation_status(self, operation: OperationDTO, event: OpenLineageRu
     def _extract_sql_query(self, event: OpenLineageRunEvent) -> SQLQueryDTO | None:
         if event.job.facets.sql:
             query = dedent(event.job.facets.sql.query).strip()
+            # https://stackoverflow.com/questions/56237415/removing-encoding-utf8-0x00-chars-from-pandas-dataframe-for-psycopg2-cursor
+            query = ASCII_UNPRINTABLE.sub(encode_char, query)
             return SQLQueryDTO(query=query)
         return None

tests/test_consumer/test_extractors/test_extractors_operation_spark.py

@@ -310,7 +310,34 @@ def test_extractors_extract_operation_spark_job_finished(
     )
 
 
-def test_extractors_extract_operation_spark_job_sql_query():
+@pytest.mark.parametrize(
+    ["sql_query", "expected_query"],
+    [
+        pytest.param(
+            "select id, name\nfrom schema.table\nwhere id = 1 \t\n\r",
+            "select id, name\nfrom schema.table\nwhere id = 1",
+            id="spaces and newlines",
+        ),
+        pytest.param(
+            """
+                select id, name
+                from schema.table
+                where id = 1
+            """,
+            "select id, name\nfrom schema.table\nwhere id = 1",
+            id="indentation",
+        ),
+        pytest.param(
+            "[\x00-\x1f\x20abc\x7e\x7f\x80\xff]",
+            "[\\x00-\\x1f\x20abc\x7e\\x7f\x80\xff]",
+            id="ascii encoding",
+        ),
+    ],
+)
+def test_extractors_extract_operation_spark_job_sql_query(
+    sql_query: str,
+    expected_query: str,
+):
     now = datetime(2024, 7, 5, 9, 6, 29, 462000, tzinfo=timezone.utc)
     run_id = UUID("01908224-8410-79a2-8de6-a769ad6944c9")
     operation_id = UUID("01908225-1fd7-746b-910c-70d24f2898b1")
@@ -328,11 +355,7 @@ def test_extractors_extract_operation_spark_job_sql_query():
                     jobType="SQL_JOB",
                 ),
                 sql=OpenLineageSqlJobFacet(
-                    query="""
-                        select id, name
-                        from schema.table
-                        where id = 1
-                    """,
+                    query=sql_query,
                 ),
             ),
         ),
@@ -369,7 +392,7 @@ def test_extractors_extract_operation_spark_job_sql_query():
         position=None,
         description=None,
         status=OperationStatusDTO.STARTED,
-        sql_query=SQLQueryDTO(query="select id, name\nfrom schema.table\nwhere id = 1"),
+        sql_query=SQLQueryDTO(query=expected_query),
         started_at=now,
         ended_at=None,
     )