Skip to content

Commit 6ebc296

Browse files
fix: fix SQL Injection issue (#821)
1 parent 8fa29e6 commit 6ebc296

File tree

1 file changed

+14
-12
lines changed

1 file changed

+14
-12
lines changed

lambda/amazon-aurora-pgvector-custom-resources/custom_resources/amazon_aurora_pgvector.py

+14-12
Original file line numberDiff line numberDiff line change
@@ -86,23 +86,25 @@ def execute_sql_commands(
8686
try:
8787
with conn.cursor() as cur:
8888

89-
sql_commands = [
90-
"CREATE EXTENSION IF NOT EXISTS vector;",
91-
f"CREATE SCHEMA {schema_name};",
92-
f"CREATE ROLE bedrock_user WITH PASSWORD '{password}' LOGIN;",
93-
f"GRANT ALL ON SCHEMA {schema_name} to bedrock_user;",
89+
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
90+
cur.execute("CREATE SCHEMA IF NOT EXISTS %s;", (schema_name,))
91+
92+
cur.execute("CREATE ROLE bedrock_user WITH PASSWORD %s LOGIN;", (password,))
93+
cur.execute("GRANT ALL ON SCHEMA %s TO bedrock_user;", (schema_name,))
94+
95+
cur.execute(
9496
f"CREATE TABLE {schema_name}.{table_name} ("
9597
f"{pk_field} uuid PRIMARY KEY, "
96-
f"{vector_field} vector({vector_dimensions}), "
98+
f"{vector_field} vector(%s), "
9799
f"{text_field} text, "
98100
f"{metadata_field} json);",
99-
f"CREATE INDEX on {schema_name}.{table_name} "
100-
f"USING hnsw ({vector_field} vector_cosine_ops);"
101-
]
101+
(vector_dimensions,),
102+
)
102103

103-
for command in sql_commands:
104-
logger.info(f"Executing SQL command: {command}")
105-
cur.execute(command)
104+
cur.execute(
105+
f"CREATE INDEX ON {schema_name}.{table_name} "
106+
f"USING hnsw ({vector_field} vector_cosine_ops);"
107+
)
106108
conn.commit()
107109
except pg8000.ProgrammingError as e:
108110
error_message = f"Error executing SQL commands: {e}"

0 commit comments

Comments
 (0)