From 82fac253798a4bdccae2ebcca9aa14e3a7d0ac4d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alja=C5=BE=20Mur=20Er=C5=BEen?=
 <aljazerzen@users.noreply.github.com>
Date: Thu, 5 Dec 2024 10:46:17 +0100
Subject: [PATCH] Show a better message for SET/RESET/SHOW over SQL native
 protocol (#8056)

---
 edb/pgsql/resolver/dispatch.py  |  2 +-
 edb/server/compiler/compiler.py |  2 +-
 edb/server/compiler/sql.py      | 12 ++++++++++++
 edb/server/defines.py           |  3 +++
 edb/server/protocol/pg_ext.pyx  |  7 ++++++-
 tests/test_sql_query.py         | 29 +++++++++++++++++++++++++++++
 6 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/edb/pgsql/resolver/dispatch.py b/edb/pgsql/resolver/dispatch.py
index d3ec583fcf5..f20ff6c0a80 100644
--- a/edb/pgsql/resolver/dispatch.py
+++ b/edb/pgsql/resolver/dispatch.py
@@ -108,7 +108,7 @@ def _raise_unsupported(expr: pgast.Base) -> typing.Never:
     # title case to spaces
     pretty_name = re.sub(r'(?<!^)(?=[A-Z])', ' ', pretty_name).upper()
 
-    raise errors.QueryError(
+    raise errors.UnsupportedFeatureError(
         f'not supported: {pretty_name}',
         span=expr.span,
         pgext_code=pgerror.ERROR_FEATURE_NOT_SUPPORTED,
diff --git a/edb/server/compiler/compiler.py b/edb/server/compiler/compiler.py
index 01c67aec6db..fd9e1abadd8 100644
--- a/edb/server/compiler/compiler.py
+++ b/edb/server/compiler/compiler.py
@@ -568,7 +568,7 @@ def compile_sql(
             apply_access_policies_sql=apply_access_policies_sql,
             disambiguate_column_names=False,
             backend_runtime_params=self.state.backend_runtime_params,
-            protocol_version=(-3, 0),  # emulated PG binary protocol version
+            protocol_version=defines.POSTGRES_PROTOCOL,
         )
 
     def compile_serialized_request(
diff --git a/edb/server/compiler/sql.py b/edb/server/compiler/sql.py
index 1f0ce8a52e3..a0637ab9071 100644
--- a/edb/server/compiler/sql.py
+++ b/edb/server/compiler/sql.py
@@ -100,6 +100,10 @@ def compile_sql(
         )
 
         if isinstance(stmt, (pgast.VariableSetStmt, pgast.VariableResetStmt)):
+            if protocol_version != defines.POSTGRES_PROTOCOL:
+                from edb.pgsql import resolver as pg_resolver
+                pg_resolver.dispatch._raise_unsupported(stmt)
+
             value: Optional[dbstate.SQLSetting]
             if isinstance(stmt, pgast.VariableSetStmt):
                 value = pg_arg_list_to_python(stmt.args)
@@ -137,6 +141,10 @@ def compile_sql(
                 unit.capabilities |= enums.Capability.SESSION_CONFIG
 
         elif isinstance(stmt, pgast.VariableShowStmt):
+            if protocol_version != defines.POSTGRES_PROTOCOL:
+                from edb.pgsql import resolver as pg_resolver
+                pg_resolver.dispatch._raise_unsupported(stmt)
+
             unit.get_var = stmt.name
             unit.frontend_only = (
                 stmt.name in FE_SETTINGS_MUTABLE
@@ -146,6 +154,10 @@ def compile_sql(
                 unit.command_complete_tag = dbstate.TagPlain(tag=b"SHOW")
 
         elif isinstance(stmt, pgast.SetTransactionStmt):
+            if protocol_version != defines.POSTGRES_PROTOCOL:
+                from edb.pgsql import resolver as pg_resolver
+                pg_resolver.dispatch._raise_unsupported(stmt)
+
             if stmt.scope == pgast.OptionsScope.SESSION:
                 unit.set_vars = {
                     f"default_{name}": (
diff --git a/edb/server/defines.py b/edb/server/defines.py
index 5199f8d481b..980283a0dab 100644
--- a/edb/server/defines.py
+++ b/edb/server/defines.py
@@ -86,6 +86,9 @@
 MIN_PROTOCOL: ProtocolVersion = (1, 0)
 CURRENT_PROTOCOL: ProtocolVersion = (3, 0)
 
+# Emulated PG binary protocol
+POSTGRES_PROTOCOL: ProtocolVersion = (-3, 0)
+
 MIN_SUGGESTED_CLIENT_POOL_SIZE = 10
 MAX_SUGGESTED_CLIENT_POOL_SIZE = 100
 
diff --git a/edb/server/protocol/pg_ext.pyx b/edb/server/protocol/pg_ext.pyx
index 2d136b8c8b7..33297da7a19 100644
--- a/edb/server/protocol/pg_ext.pyx
+++ b/edb/server/protocol/pg_ext.pyx
@@ -428,7 +428,12 @@ cdef class PgConnection(frontend.FrontendConnection):
                 str(exc), detail=exc.details, severity="FATAL"
             )
         elif isinstance(exc, errors.UnsupportedFeatureError):
-            exc = pgerror.FeatureNotSupported(str(exc))
+            exc = pgerror.new(
+                pgerror.ERROR_FEATURE_NOT_SUPPORTED,
+                str(exc),
+                L=str(exc.line) if exc.line >= 0 else None,
+                P=str(exc.position + 1) if exc.position >= 0 else None,
+            )
         elif isinstance(exc, errors.EdgeDBError):
             args = dict(hint=exc.hint, detail=exc.details)
             if exc.line >= 0:
diff --git a/tests/test_sql_query.py b/tests/test_sql_query.py
index 2eff2b60bd9..2501ff1a483 100644
--- a/tests/test_sql_query.py
+++ b/tests/test_sql_query.py
@@ -2720,3 +2720,32 @@ async def test_native_sql_query_15(self):
             'SELECT title FROM "Content" ORDER BY title',
             [{'title': 'Forrest Gump'}]
         )
+
+    async def test_native_sql_query_16(self):
+        with self.assertRaisesRegex(
+            edgedb.UnsupportedFeatureError,
+            "not supported: VARIABLE SET",
+            _position=14,  # this point to `1`, but hey, better than nothing
+        ):
+            await self.assert_sql_query_result(
+                'SET my_var TO 1',
+                []
+            )
+
+        with self.assertRaisesRegex(
+            edgedb.UnsupportedFeatureError,
+            "not supported: VARIABLE RESET",
+        ):
+            await self.assert_sql_query_result(
+                'RESET my_var',
+                []
+            )
+
+        with self.assertRaisesRegex(
+            edgedb.UnsupportedFeatureError,
+            "not supported: VARIABLE SHOW",
+        ):
+            await self.assert_sql_query_result(
+                'SHOW my_var',
+                []
+            )