Fix @FUNCTION_NAME@ not being handled in sql attributes (#443)

* Fix @FUNCTION_NAME@ not being handled in sql attributes Signed-off-by: Ana Hobden <[email protected]> * Fixup test Signed-off-by: Ana Hobden <[email protected]> * Fix test name Signed-off-by: Ana Hobden <[email protected]>
pgcentralfoundation · Feb 16, 2022 · 08828b7 · 08828b7
1 parent b527960
commit 08828b7
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git a/pgx-tests/src/tests/mod.rs b/pgx-tests/src/tests/mod.rs
@@ -21,7 +21,7 @@ mod log_tests;
 mod memcxt_tests;
 mod name_tests;
 mod numeric_tests;
-mod pg_extern_args_tests;
+mod pg_extern_tests;
 mod pg_try_tests;
 mod pgbox_tests;
 mod postgres_type_tests;

diff --git a/pgx-tests/src/tests/pg_extern_args_tests.rs → pgx-tests/src/tests/pg_extern_tests.rs b/pgx-tests/src/tests/pg_extern_args_tests.rs → pgx-tests/src/tests/pg_extern_tests.rs
@@ -30,4 +30,25 @@ mod tests {
         .expect("failed to get SPI result");
         assert!(result)
     }
+
+
+    // Ensures `@FUNCTION_NAME@` is handled.
+    #[pg_extern(sql = r#"
+        CREATE OR REPLACE FUNCTION tests."overridden_sql_with_fn_name"() RETURNS void
+        STRICT
+        LANGUAGE c /* Rust */
+        AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
+    "#)]
+    fn overridden_sql_with_fn_name() -> bool {
+        true
+    }
+
+    #[pg_test]
+    fn test_overridden_sql_with_fn_name() {
+        let result = Spi::get_one::<bool>(
+            r#"SELECT tests."overridden_sql_with_fn_name"()"#,
+        )
+        .expect("failed to get SPI result");
+        assert!(result)
+    }
 }
diff --git a/pgx-utils/src/sql_entity_graph/pg_extern/mod.rs b/pgx-utils/src/sql_entity_graph/pg_extern/mod.rs
@@ -205,6 +205,18 @@ impl PgExtern {
         }
 
         let func = syn::parse2::<syn::ItemFn>(item)?;
+
+        if let Some(ref mut to_sql_config) = to_sql_config {
+            if let Some(ref mut content) = to_sql_config.content {
+                let value = content.value();
+                let updated_value = value.replace(
+                    "@FUNCTION_NAME@",
+                    &*(func.sig.ident.to_string() + "_wrapper"),
+                ) + "\n";
+                *content = syn::LitStr::new(&updated_value, Span::call_site());
+            }
+        }
+
         Ok(Self {
             attrs,
             func,