From 8105f737fd8593c5a8f00493ad591bbc2e3b9159 Mon Sep 17 00:00:00 2001
From: Adam Guo <adamguo@amazon.com>
Date: Fri, 17 Nov 2023 16:50:49 -0500
Subject: [PATCH] Allow `@` character in extension names (#244)

---
 src/tleextension.c                  |  9 ++++--
 test/expected/pg_tle_management.out | 45 +++++++++++++++++++++++++++++
 test/sql/pg_tle_management.sql      | 22 ++++++++++++++
 3 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/src/tleextension.c b/src/tleextension.c
index 766d858..6bc5cc5 100644
--- a/src/tleextension.c
+++ b/src/tleextension.c
@@ -486,13 +486,16 @@ check_valid_extension_name(const char *extensionname)
 	 * Check for alphanumeric character in extension name for now. Although
 	 * this does prevent some naming schemes, it's a more straight forward
 	 * prevention for preventing certain injection attacks due to the way the
-	 * way we rely on functions currently. Allow the '_' or '-' character to
-	 * provide a nice separator if desired.
+	 * way we rely on functions currently. Allow the '_', '-', or '@'
+	 * character to provide a nice separator if desired.
 	 */
 
 	while (extensionname[idx] != '\0')
 	{
-		if (!isalnum(extensionname[idx]) && extensionname[idx] != '_' && extensionname[idx] != '-')
+		if (!isalnum(extensionname[idx]) &&
+			extensionname[idx] != '_' &&
+			extensionname[idx] != '-' &&
+			extensionname[idx] != '@')
 			ereport(ERROR,
 					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 					 errmsg("invalid extension name: \"%s\"", extensionname),
diff --git a/test/expected/pg_tle_management.out b/test/expected/pg_tle_management.out
index ce0b7e1..e9b844f 100644
--- a/test/expected/pg_tle_management.out
+++ b/test/expected/pg_tle_management.out
@@ -815,6 +815,51 @@ SELECT pgtle.available_extension_versions();
 ------------------------------
 (0 rows)
 
+-- install extension with '@' symbol in name
+SELECT pgtle.install_extension
+(
+ 'foo@bar',
+ '1.0',
+ 'Test TLE Functions',
+$_pgtle_$
+  CREATE OR REPLACE FUNCTION at_func()
+  RETURNS INT AS $$
+  (
+    SELECT 42
+  )$$ LANGUAGE sql;
+$_pgtle_$
+);
+ install_extension 
+-------------------
+ t
+(1 row)
+
+SELECT pgtle.available_extensions();
+        available_extensions        
+------------------------------------
+ (foo@bar,1.0,"Test TLE Functions")
+(1 row)
+
+CREATE EXTENSION "foo@bar";
+SELECT extname, extversion from pg_extension where extname='foo@bar';
+ extname | extversion 
+---------+------------
+ foo@bar | 1.0
+(1 row)
+
+SELECT at_func();
+ at_func 
+---------
+      42
+(1 row)
+
+DROP EXTENSION "foo@bar";
+SELECT pgtle.uninstall_extension('foo@bar');
+ uninstall_extension 
+---------------------
+ t
+(1 row)
+
 -- Skip TransactionStmts
 BEGIN;
 SELECT pgtle.available_extension_versions();
diff --git a/test/sql/pg_tle_management.sql b/test/sql/pg_tle_management.sql
index 1949b20..ace0901 100644
--- a/test/sql/pg_tle_management.sql
+++ b/test/sql/pg_tle_management.sql
@@ -531,6 +531,28 @@ SELECT pgtle.uninstall_extension('test42');
 
 SELECT pgtle.available_extension_versions();
 
+-- install extension with '@' symbol in name
+SELECT pgtle.install_extension
+(
+ 'foo@bar',
+ '1.0',
+ 'Test TLE Functions',
+$_pgtle_$
+  CREATE OR REPLACE FUNCTION at_func()
+  RETURNS INT AS $$
+  (
+    SELECT 42
+  )$$ LANGUAGE sql;
+$_pgtle_$
+);
+
+SELECT pgtle.available_extensions();
+CREATE EXTENSION "foo@bar";
+SELECT extname, extversion from pg_extension where extname='foo@bar';
+SELECT at_func();
+DROP EXTENSION "foo@bar";
+SELECT pgtle.uninstall_extension('foo@bar');
+
 -- Skip TransactionStmts
 BEGIN;
 SELECT pgtle.available_extension_versions();