Enable indexing and json->scalar calls to be inlined by postgres (#8804)

Functions that are marked immutable but make calls to stable functions
can't be inlined, and we have a lot of functions like that. (Because
error raising is stable in postgres but we want to allow calling some
functions that can error in indexes/constraints.)

Fix this by marking the functions as stable but making sure that their
*trampoline* wrappers are immutable, since those are what actually go
into indexes/constraints.

Author: msullivan

edb/pgsql/metaschema.py

@@ -1772,7 +1772,8 @@ def __init__(self) -> None:
                 ('detail', ('text',), "''"),
             ],
             returns=('text',),
-            volatility='immutable',
+            volatility='stable',
+            wrapper_volatility='immutable',
             text=self.text,
         )
 
@@ -2006,7 +2007,8 @@ def __init__(self) -> None:
             returns=('anyelement',),
             # Min volatility of exception helpers and pg_typeof is 'stable',
             # but for all practical purposes, we can assume 'immutable'
-            volatility='immutable',
+            volatility='stable',
+            wrapper_volatility='immutable',
             text=self.text,
         )
 
@@ -2035,7 +2037,8 @@ def __init__(self) -> None:
                 ("stop", ("bigint",)),
             ],
             returns=("anyarray",),
-            volatility="immutable",
+            volatility="stable",
+            wrapper_volatility='immutable',
             text=self.text,
         )
 
@@ -2073,7 +2076,8 @@ def __init__(self) -> None:
             returns=('text',),
             # Min volatility of exception helpers and pg_typeof is 'stable',
             # but for all practical purposes, we can assume 'immutable'
-            volatility='immutable',
+            volatility='stable',
+            wrapper_volatility='immutable',
             text=self.text,
         )
 
@@ -2110,7 +2114,8 @@ def __init__(self) -> None:
             returns=('bytea',),
             # Min volatility of exception helpers and pg_typeof is 'stable',
             # but for all practical purposes, we can assume 'immutable'
-            volatility='immutable',
+            volatility='stable',
+            wrapper_volatility='immutable',
             text=self.text,
         )
 
@@ -2222,7 +2227,7 @@ def __init__(self) -> None:
                 ('stop', ('bigint',)),
             ],
             returns=('text',),
-            volatility='immutable',
+            volatility='stable',
             text=self.text)
 
 
@@ -2241,7 +2246,7 @@ def __init__(self) -> None:
                 ('stop', ('bigint',)),
             ],
             returns=('bytea',),
-            volatility='immutable',
+            volatility='stable',
             text=self.text)
 
 
@@ -2299,7 +2304,8 @@ def __init__(self) -> None:
             returns=('jsonb',),
             # Min volatility of exception helpers 'stable',
             # but for all practical purposes, we can assume 'immutable'
-            volatility='immutable',
+            volatility='stable',
+            wrapper_volatility='immutable',
             strict=True,
             text=self.text,
         )
@@ -2365,7 +2371,8 @@ def __init__(self) -> None:
             returns=('jsonb',),
             # Min volatility of exception helpers and pg_typeof is 'stable',
             # but for all practical purposes, we can assume 'immutable'
-            volatility='immutable',
+            volatility='stable',
+            wrapper_volatility='immutable',
             strict=True,
             text=self.text,
         )
@@ -2417,7 +2424,8 @@ def __init__(self) -> None:
             returns=("jsonb",),
             # Min volatility of to_jsonb is 'stable',
             # but for all practical purposes, we can assume 'immutable'
-            volatility="immutable",
+            volatility="stable",
+            wrapper_volatility='immutable',
             text=self.text,
         )
 

edb/pgsql/patches.py

@@ -303,4 +303,13 @@ def get_version_key(num_patches: int) -> str:
     USING (to_bytes(to_str(j)));
 };
 '''),
+    ('metaschema-sql', 'ArrayIndexWithBoundsFunction'),
+    ('metaschema-sql', 'ArraySliceFunction'),
+    ('metaschema-sql', 'StringIndexWithBoundsFunction'),
+    ('metaschema-sql', 'BytesIndexWithBoundsFunction'),
+    ('metaschema-sql', 'StringSliceFunction'),
+    ('metaschema-sql', 'BytesSliceFunction'),
+    ('metaschema-sql', 'JSONIndexByTextFunction'),
+    ('metaschema-sql', 'JSONIndexByIntFunction'),
+    ('metaschema-sql', 'JSONSliceFunction'),
 ]

edb/pgsql/trampoline.py

@@ -36,6 +36,8 @@
 from __future__ import annotations
 from typing import (
     TYPE_CHECKING,
+    Optional,
+    Sequence,
 )
 
 import abc
@@ -62,13 +64,44 @@ def fixup_query(query: str) -> str:
 
 
 class VersionedFunction(dbops.Function):
-    if not TYPE_CHECKING:
-        def __init__(self, *args, **kwargs):
+    if TYPE_CHECKING:
+        # What the volatility of the trampoline wrapper should be.
+        # This is sometimes immutable even when the underlying is
+        # stable, for functions that must be immutable so they can go
+        # into indexes/constraints but might do something technically
+        # stable (like raise an error).
+        #
+        # This allows the real function to be inlined while allowing
+        # the wrapper to get used in indexes/constraints.
+        wrapper_volatility: Optional[str]
+
+        def __init__(
+            self,
+            name: tuple[str, ...],
+            *,
+            args: Optional[Sequence[dbops.FunctionArg]] = None,
+            returns: str | tuple[str, ...],
+            text: str,
+            volatility: str = "volatile",
+            language: str = "sql",
+            has_variadic: Optional[bool] = None,
+            strict: bool = False,
+            parallel_safe: bool = False,
+            set_returning: bool = False,
+
+            wrapper_volatility: Optional[str] = None,
+        ):
+            pass
+
+    else:
+        def __init__(self, *args, wrapper_volatility=None, **kwargs):
             super().__init__(*args, **kwargs)
             self.name = (
                 common.maybe_versioned_schema(self.name[0]), *self.name[1:])
             self.text = fixup_query(self.text)
 
+            self.wrapper_volatility = wrapper_volatility
+
             if self.args:
                 nargs = []
                 for arg in self.args:
@@ -156,6 +189,8 @@ def make_trampoline(func: dbops.Function) -> TrampolineFunction:
     new_func.text = f'select {q(*func.name)}({", ".join(args)})'
     new_func.language = 'sql'
     new_func.strict = False
+    if isinstance(func, VersionedFunction) and func.wrapper_volatility:
+        new_func.volatility = func.wrapper_volatility
     return TrampolineFunction(new_func.name, new_func)