fix(dynamo_db.py): don't auto-create tables, allow database_type == 'dynamodb'

This commit is contained in:
Krrish Dholakia 2024-01-12 11:33:40 +05:30
parent aeff32ac62
commit 352f943dcf
2 changed files with 33 additions and 29 deletions

View file

@ -51,38 +51,40 @@ class DynamoDBWrapper(CustomDB):
## User ## User
table = client.table(DBTableNames.user.value) table = client.table(DBTableNames.user.value)
if not await table.exists(): if not await table.exists():
try: sample_code_snippet = f"""
table = client.table({DBTableNames.user.value})
await table.create( await table.create(
self.throughput_type, self.throughput_type,
KeySchema(hash_key=KeySpec("user_id", KeyType.string)), KeySchema(hash_key=KeySpec("user_id", KeyType.string)),
) )
except: """
raise Exception( raise Exception(
f"Failed to create table - {DBTableNames.user.value}.\nPlease create a new table called {DBTableNames.user.value}\nAND set `hash_key` as 'user_id'" f"Failed to create table - {DBTableNames.user.value}.\nPlease create a new table called {DBTableNames.user.value}\nAND set `hash_key` as 'user_id'\n\nEg.: {sample_code_snippet}"
) )
## Token ## Token
table = client.table(DBTableNames.key.value)
if not await table.exists(): if not await table.exists():
try: sample_code_snippet = f"""
table = client.table({DBTableNames.key.value})
await table.create( await table.create(
self.throughput_type, self.throughput_type,
KeySchema(hash_key=KeySpec("token", KeyType.string)), KeySchema(hash_key=KeySpec("token", KeyType.string)),
) )
except: """
raise Exception( raise Exception(
f"Failed to create table - {DBTableNames.key.value}.\nPlease create a new table called {DBTableNames.key.value}\nAND set `hash_key` as 'token'" f"Failed to create table - {DBTableNames.key.value}.\nPlease create a new table called {DBTableNames.key.value}\nAND set `hash_key` as 'token'\n\nE.g.: {sample_code_snippet}"
) )
## Config ## Config
table = client.table(DBTableNames.config.value) table = client.table(DBTableNames.config.value)
if not await table.exists(): if not await table.exists():
try: sample_code_snippet = f"""
table = client.table({DBTableNames.config.value})
await table.create( await table.create(
self.throughput_type, self.throughput_type,
KeySchema(hash_key=KeySpec("param_name", KeyType.string)), KeySchema(hash_key=KeySpec("param_name", KeyType.string)),
) )
except: """
raise Exception( raise Exception(
f"Failed to create table - {DBTableNames.config.value}.\nPlease create a new table called {DBTableNames.config.value}\nAND set `hash_key` as 'token'" f"Failed to create table - {DBTableNames.config.value}.\nPlease create a new table called {DBTableNames.config.value}\nAND set `hash_key` as 'param_name'\n\nE.g.: {sample_code_snippet}"
) )
async def insert_data( async def insert_data(

View file

@ -853,7 +853,9 @@ class ProxyConfig:
) )
## dynamodb ## dynamodb
database_type = general_settings.get("database_type", None) database_type = general_settings.get("database_type", None)
if database_type is not None and database_type == "dynamo_db": if database_type is not None and (
database_type == "dynamo_db" or database_type == "dynamodb"
):
database_args = general_settings.get("database_args", None) database_args = general_settings.get("database_args", None)
custom_db_client = DBClient( custom_db_client = DBClient(
custom_db_args=database_args, custom_db_type=database_type custom_db_args=database_args, custom_db_type=database_type