diff --git a/cli/magic-config.ts b/cli/magic-config.ts index 2b272c3f0..b361fb5f2 100644 --- a/cli/magic-config.ts +++ b/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 { }, 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 { 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"), diff --git a/lib/rag-engines/aurora-pgvector/functions/pgvector-setup/index.py b/lib/rag-engines/aurora-pgvector/functions/pgvector-setup/index.py index 146435216..aeea4e25e 100644 --- a/lib/rag-engines/aurora-pgvector/functions/pgvector-setup/index.py +++ b/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) diff --git a/lib/rag-engines/aurora-pgvector/index.ts b/lib/rag-engines/aurora-pgvector/index.ts index 043c4d963..7e89c0699 100644 --- a/lib/rag-engines/aurora-pgvector/index.ts +++ b/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, }, } ); diff --git a/lib/shared/layers/python-sdk/python/genai_core/aurora/connection.py b/lib/shared/layers/python-sdk/python/genai_core/aurora/connection.py index b303fad7c..f81b167b7 100644 --- a/lib/shared/layers/python-sdk/python/genai_core/aurora/connection.py +++ b/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() diff --git a/lib/shared/types.ts b/lib/shared/types.ts index 423247c48..1076667a6 100644 --- a/lib/shared/types.ts +++ b/lib/shared/types.ts @@ -106,6 +106,7 @@ export interface SystemConfig { engines: { aurora: { enabled: boolean; + defaultDatabaseName?: string; }; opensearch: { enabled: boolean;