Ability to specify Aurora DB Name for storing embeddings

cli/magic-config.ts

@@ -164,6 +164,7 @@ const embeddingModels = [
       ) {
         options.ragsToEnable.pop("kendra");
       }
+      options.auroraDefaultDBName = config.rag.engines.aurora.defaultDatabaseName
       options.embeddings = config.rag.embeddingsModels.map((m: any) => m.name);
       options.defaultEmbedding = (config.rag.embeddingsModels ?? []).filter(
         (m: any) => m.default
@@ -568,6 +569,20 @@ async function processCreateOptions(options: any): Promise<void> {
       },
       initial: options.ragsToEnable || [],
     },
+    {
+      type:"input",
+      name:"auroraDefaultDBName",
+      message: "Enter a default database name for Aurora",
+      initial: options.auroraDefaultDBName || undefined,
+      validate(v: string) {
+        return (v.length === 0 || (this as any).skipped || 
+          RegExp(/^[a-zA-Z][a-zA-Z0-9_]{0,62}$/).test(v)) ? 
+            true : "You need to enter a valid database name";
+      },
+      skip(): boolean {
+        return !(this as any).state.answers.ragsToEnable.includes("aurora");
+      },
+    },
     {
       type: "confirm",
       name: "kendraEnterprise",
@@ -744,6 +759,8 @@ async function processCreateOptions(options: any): Promise<void> {
       engines: {
         aurora: {
           enabled: answers.ragsToEnable.includes("aurora"),
+          defaultDatabaseName: answers.ragsToEnable.includes("aurora") && answers.auroraDefaultDBName.length > 0 ? 
+            answers.auroraDefaultDBName : undefined
         },
         opensearch: {
           enabled: answers.ragsToEnable.includes("opensearch"),

lib/rag-engines/aurora-pgvector/functions/pgvector-setup/index.py

@@ -15,25 +15,44 @@ def lambda_handler(event, context: LambdaContext):
     request_type = event["RequestType"]
     resource_properties = event["ResourceProperties"]
     AURORA_DB_SECRET_ID = resource_properties["AURORA_DB_SECRET_ID"]
+    AURORA_DB_NAME= resource_properties["AURORA_DB_NAME"]
 
     secret_response = secretsmanager_client.get_secret_value(
         SecretId=AURORA_DB_SECRET_ID
     )
     database_secrets = json.loads(secret_response["SecretString"])
-    dbhost = database_secrets["host"]
-    dbport = database_secrets["port"]
-    dbuser = database_secrets["username"]
-    dbpass = database_secrets["password"]
 
     if request_type == "Create" or request_type == "Update":
-        dbconn = psycopg2.connect(
-            host=dbhost, user=dbuser, password=dbpass, port=dbport, connect_timeout=10
-        )
-
+        db_conn_args = {
+            "host": database_secrets["host"],
+            "user": database_secrets["username"],
+            "password": database_secrets["password"],
+            "port": database_secrets["port"],
+            "connect_timeout":10,
+        }
+        if "dbname" in database_secrets:
+            db_conn_args["dbname"] = database_secrets["dbname"]
+        dbconn = psycopg2.connect(**db_conn_args)
         dbconn.set_session(autocommit=True)
-
         cur = dbconn.cursor()
 
+        # Check if database exists, if not create it.
+        cur.execute(f"SELECT * FROM pg_database WHERE LOWER(datname) = LOWER('{AURORA_DB_NAME}')")
+        rows = cur.fetchall()
+        if len(rows) == 0:
+            cur.execute(f"CREATE DATABASE {AURORA_DB_NAME}")
+            logger.info(f"Created database {AURORA_DB_NAME}")
+            # close existing connection and re-open connect to the new database
+            cur.close()
+            db_conn_args["dbname"] = AURORA_DB_NAME
+            dbconn = psycopg2.connect(**db_conn_args)
+            dbconn.set_session(autocommit=True)
+            cur = dbconn.cursor()
+        else:
+            logger.info(f"Database {AURORA_DB_NAME} already exists")
+        cur.close()
+
+        # Create vector extension if not exists. This will create the vector type.
         cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
         register_vector(dbconn)
 

lib/rag-engines/aurora-pgvector/index.ts

@@ -77,6 +77,7 @@ export class AuroraPgVector extends Construct {
         serviceToken: databaseSetupProvider.serviceToken,
         properties: {
           AURORA_DB_SECRET_ID: dbCluster.secret?.secretArn as string,
+          AURORA_DB_NAME: props.config.rag.engines.aurora.defaultDatabaseName,
         },
       }
     );

lib/shared/layers/python-sdk/python/genai_core/aurora/connection.py

@@ -21,15 +21,20 @@ def __init__(self, autocommit=True):
         self.dbport = database_secrets["port"]
         self.dbuser = database_secrets["username"]
         self.dbpass = database_secrets["password"]
+        if "dbname" in database_secrets:
+            self.dbname = database_secrets["dbname"]
 
     def __enter__(self):
-        connection = psycopg2.connect(
-            host=self.dbhost,
-            user=self.dbuser,
-            password=self.dbpass,
-            port=self.dbport,
-            connect_timeout=10,
-        )
+        db_conn_args = {
+            "host":self.dbhost,
+            "user":self.dbuser,
+            "password":self.dbpass,
+            "port":self.dbport,
+            "connect_timeout":10,
+        }
+        if hasattr(self, "dbname"):
+            db_conn_args["dbname"] = self.dbname
+        connection = psycopg2.connect(**db_conn_args)
 
         connection.set_session(autocommit=self.autocommit)
         psycopg2.extras.register_uuid()

lib/shared/types.ts

@@ -106,6 +106,7 @@ export interface SystemConfig {
     engines: {
       aurora: {
         enabled: boolean;
+        defaultDatabaseName?: string;
       };
       opensearch: {
         enabled: boolean;