From 352f943dcff2d7ce9f360ad7bd27b38a9a897c04 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 12 Jan 2024 11:33:40 +0530
Subject: [PATCH] fix(dynamo_db.py): don't auto-create tables, allow
 database_type == 'dynamodb'

---
 litellm/proxy/db/dynamo_db.py | 58 ++++++++++++++++++-----------------
 litellm/proxy/proxy_server.py |  4 ++-
 2 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py
index 17e5bd3b8..149e73066 100644
--- a/litellm/proxy/db/dynamo_db.py
+++ b/litellm/proxy/db/dynamo_db.py
@@ -51,39 +51,41 @@ class DynamoDBWrapper(CustomDB):
             ## User
             table = client.table(DBTableNames.user.value)
             if not await table.exists():
-                try:
-                    await table.create(
-                        self.throughput_type,
-                        KeySchema(hash_key=KeySpec("user_id", KeyType.string)),
-                    )
-                except:
-                    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'"
-                    )
+                sample_code_snippet = f"""
+                table = client.table({DBTableNames.user.value})
+                await table.create(
+                    self.throughput_type,
+                    KeySchema(hash_key=KeySpec("user_id", KeyType.string)),
+                )
+                """
+                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'\n\nEg.: {sample_code_snippet}"
+                )
             ## Token
-            table = client.table(DBTableNames.key.value)
             if not await table.exists():
-                try:
-                    await table.create(
-                        self.throughput_type,
-                        KeySchema(hash_key=KeySpec("token", KeyType.string)),
-                    )
-                except:
-                    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'"
-                    )
+                sample_code_snippet = f""" 
+                table = client.table({DBTableNames.key.value})
+                await table.create(
+                    self.throughput_type,
+                    KeySchema(hash_key=KeySpec("token", KeyType.string)),
+                )
+                """
+                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'\n\nE.g.: {sample_code_snippet}"
+                )
             ## Config
             table = client.table(DBTableNames.config.value)
             if not await table.exists():
-                try:
-                    await table.create(
-                        self.throughput_type,
-                        KeySchema(hash_key=KeySpec("param_name", KeyType.string)),
-                    )
-                except:
-                    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'"
-                    )
+                sample_code_snippet = f"""
+                table = client.table({DBTableNames.config.value})
+                await table.create(
+                    self.throughput_type,
+                    KeySchema(hash_key=KeySpec("param_name", KeyType.string)),
+                )
+                """
+                raise Exception(
+                    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(
         self, value: Any, table_name: Literal["user", "key", "config"]
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 25dde7cf0..09070ef0a 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -853,7 +853,9 @@ class ProxyConfig:
                 )
             ## dynamodb
             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)
                 custom_db_client = DBClient(
                     custom_db_args=database_args, custom_db_type=database_type