Replace volatility_map with a classmethod on FunctionVolatility

Author: enobayram

src/sqlalchemy_declarative_extensions/dialects/postgresql/function.py

@@ -21,6 +21,17 @@ class FunctionVolatility(enum.Enum):
     STABLE = "STABLE"
     IMMUTABLE = "IMMUTABLE"
 
+    @classmethod
+    def from_provolatile(cls, provolatile: str) -> FunctionVolatility:
+        """Convert a `pg_proc.provolatile` value to a `FunctionVolatility` enum."""
+        if provolatile == "v":
+            return cls.VOLATILE
+        if provolatile == "s":
+            return cls.STABLE
+        if provolatile == "i":
+            return cls.IMMUTABLE
+        raise ValueError(f"Invalid volatility: {provolatile}")
+
 
 def normalize_arg(arg: str) -> str:
     parts = arg.strip().split(maxsplit=1)

src/sqlalchemy_declarative_extensions/dialects/postgresql/query.py

@@ -217,11 +217,6 @@ def get_procedures_postgresql(connection: Connection) -> Sequence[BaseProcedure]
 
 def get_functions_postgresql(connection: Connection) -> Sequence[BaseFunction]:
     functions = []
-    volatility_map = {
-        "v": FunctionVolatility.VOLATILE,
-        "s": FunctionVolatility.STABLE,
-        "i": FunctionVolatility.IMMUTABLE,
-    }
 
     for f in connection.execute(functions_query).fetchall():
         name = f.name
@@ -233,7 +228,7 @@ def get_functions_postgresql(connection: Connection) -> Sequence[BaseFunction]:
             parameters=(
                 [p.strip() for p in f.parameters.split(",")] if f.parameters else None
             ),
-            volatility=volatility_map[f.volatility],
+            volatility=FunctionVolatility.from_provolatile(f.volatility),
             name=name,
             definition=definition,
             language=language,