[sql] Fix non-normalized queries over native protocol (#8315)

Queries where we had to fall back to the non-normalized version were
broken under the native protocol, because we were still feeding the
arguments to postgres.

Also, make `pg_get_serial_sequence` force non-normalization, since it
ignores its arguments. It still fails if actual params are used,
though.

edb/edgeql/tokenizer.py

@@ -97,6 +97,9 @@ def from_string(text: str) -> Source:
     def __repr__(self):
         return f'<edgeql.Source text={self._text!r}>'
 
+    def denormalized(self) -> Source:
+        return self
+
 
 class NormalizedSource(Source):
     def __init__(
@@ -140,6 +143,9 @@ def from_string(text: str) -> NormalizedSource:
         normalized = _normalize(text)
         return NormalizedSource(normalized, text, normalized.pack())
 
+    def denormalized(self) -> Source:
+        return Source.from_string(self._text)
+
 
 def inflate_span(
     source: str, span: Tuple[int, Optional[int]]

edb/pgsql/parser/parser.pyx

@@ -271,6 +271,9 @@ cdef class Source:
     def from_string(cls, text: str) -> Source:
         return Source(text)
 
+    def denormalized(self) -> Source:
+        return self
+
 
 cdef class NormalizedSource(Source):
     def __init__(
@@ -383,6 +386,8 @@ cdef class NormalizedSource(Source):
             serialized,
         )
 
+    def denormalized(self) -> Source:
+        return Source.from_string(self._orig_text)
 
 def deserialize(serialized: bytes) -> Source:
     if serialized[0] == 0:

edb/pgsql/resolver/static.py

@@ -344,6 +344,7 @@ def eval_FuncCall(
         )
 
     if fn_name == "pg_get_serial_sequence":
+        eval_list(expr.args, ctx=ctx)
         # we do not expose sequences, so any calls to this function returns NULL
         return pgast.NullConstant()
 

edb/server/compiler/compiler.py

@@ -586,7 +586,7 @@ def compile_sql(
             disambiguate_column_names=False,
             backend_runtime_params=self.state.backend_runtime_params,
             protocol_version=defines.POSTGRES_PROTOCOL,
-        )
+        )[0]
 
     def compile_serialized_request(
         self,
@@ -2682,7 +2682,7 @@ def compile_sql_as_unit_group(
         ],
     )
 
-    sql_units = sql.compile_sql(
+    sql_units, force_non_normalized = sql.compile_sql(
         source,
         schema=schema,
         tx_state=sql_tx_state,
@@ -2702,6 +2702,7 @@ def compile_sql_as_unit_group(
     qug = dbstate.QueryUnitGroup(
         cardinality=sql_units[-1].cardinality,
         cacheable=True,
+        force_non_normalized=force_non_normalized,
     )
 
     for sql_unit in sql_units:

edb/server/compiler/dbstate.py

@@ -423,6 +423,8 @@ class QueryUnitGroup:
     cache_state: int = 0
     tx_seq_id: int = 0
 
+    force_non_normalized: bool = False
+
     @property
     def units(self) -> List[QueryUnit]:
         if self._unpacked_units is None:

edb/server/compiler/sql.py

@@ -81,7 +81,7 @@ def compile_sql(
     backend_runtime_params: pg_params.BackendRuntimeParams,
     protocol_version: defines.ProtocolVersion,
     implicit_limit: Optional[int] = None,
-) -> List[dbstate.SQLQueryUnit]:
+) -> tuple[list[dbstate.SQLQueryUnit], bool]:
     def _try(
         q: str, normalized_params: List[int]
     ) -> List[dbstate.SQLQueryUnit]:
@@ -108,7 +108,7 @@ def _try(
     normalized_params = list(source.extra_type_oids())
     try:
         try:
-            return _try(source.text(), normalized_params)
+            return _try(source.text(), normalized_params), False
         except DisableNormalization:
             # compiler requested non-normalized query (it needs it for static
             # evaluation)
@@ -120,7 +120,7 @@ def _try(
                     # TODO: Can we tell the server to cache using non-extracted?
                     for unit in units:
                         unit.cacheable = False
-                    return units
+                    return units, True
             except DisableNormalization:
                 pass
 

edb/server/dbview/dbview.pyx

@@ -1436,6 +1436,8 @@ cdef class DatabaseConnectionView:
             )
 
         source = query_req.source
+        if query_unit_group.force_non_normalized:
+            source = source.denormalized()
         return CompiledQuery(
             query_unit_group=query_unit_group,
             first_extra=source.first_extra(),
@@ -1464,6 +1466,9 @@ cdef class DatabaseConnectionView:
 
         desc_map = {}
         source = query_req.source
+        if qug.force_non_normalized:
+            source = source.denormalized()
+
         first_extra = source.first_extra()
         num_injected_params = 0
         if qug.globals is not None:

tests/test_sql_query.py

@@ -2957,6 +2957,20 @@ async def test_sql_native_query_27(self):
                 select (), asdf
             ''')
 
+    async def test_sql_native_query_28(self):
+        await self.assert_sql_query_result(
+            "SELECT current_setting('search_path') as path",
+            [{"path": 'public'}],
+        )
+
+        await self.assert_sql_query_result(
+            '''
+            SELECT pg_get_serial_sequence('"public"."Book"', 1)
+                ::regclass::text as seq;
+            ''',
+            [{"seq": None}],
+        )
+
 
 class TestSQLQueryNonTransactional(tb.SQLQueryTestCase):