tweak to use regex to fix function call

Mause · Dec 12, 2023 · a719a8d · a719a8d
1 parent bca2d99
commit a719a8d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py
@@ -69,8 +69,7 @@ def get_check_constraints(
             raise NotImplementedError() from e
 
 
-PG_GET_CONSTRAINTDEF = "pg_get_constraintdef(cons.oid)"
-PG_GET_CONSTRAINTDEF_NEW = "pg_get_constraintdef(cons.oid, true)"
+PG_GET_CONSTRAINTDEF = re.compile(r"pg_get_constraintdef\(([a-z\.]+)\)", re.I)
 
 
 class ConnectionWrapper:
@@ -141,7 +140,9 @@ def execute(
         parameters: Optional[Tuple] = None,
         context: Optional[Any] = None,
     ) -> None:
-        statement = statement.replace(PG_GET_CONSTRAINTDEF, PG_GET_CONSTRAINTDEF_NEW)
+        statement = PG_GET_CONSTRAINTDEF.sub(
+            lambda match: f"pg_get_constraintdef({match[1]}, true)", statement
+        )
         try:
             if statement.lower() == "commit":  # this is largely for ipython-sql
                 self.__c.commit()