refactor cosmos db script

Author: ks6088ts

apps/3_call_azure_cosmos_db/README.md

@@ -32,29 +32,38 @@ $ python apps/3_call_azure_cosmos_db/main.py Hello
 ### Example
 
 ```shell
-$ python apps/3_call_azure_cosmos_db/main.py Hello
-Created item:
+$ python apps/3_call_azure_cosmos_db/main.py --command create
+Create item
 {'_attachments': 'attachments/',
- '_etag': '"8100471f-0000-2300-0000-66b32f2b0000"',
- '_rid': 'Ny9FAOMmTIMKAAAAAAAAAA==',
- '_self': 'dbs/Ny9FAA==/colls/Ny9FAOMmTIM=/docs/Ny9FAOMmTIMKAAAAAAAAAA==/',
- '_ts': 1723019051,
+ '_etag': '"0000ef9c-0000-2300-0000-66dfd5140000"',
+ '_rid': 'ipcwAJFKMxICAAAAAAAAAA==',
+ '_self': 'dbs/ipcwAA==/colls/ipcwAJFKMxI=/docs/ipcwAJFKMxICAAAAAAAAAA==/',
+ '_ts': 1725945108,
  'content': 'Hello, world!',
- 'id': 'Hello',
+ 'id': 'test',
  'role': 'assistant'}
+
+$ python apps/3_call_azure_cosmos_db/main.py --command read
 Read item:
 {'_attachments': 'attachments/',
- '_etag': '"8100471f-0000-2300-0000-66b32f2b0000"',
- '_rid': 'Ny9FAOMmTIMKAAAAAAAAAA==',
- '_self': 'dbs/Ny9FAA==/colls/Ny9FAOMmTIM=/docs/Ny9FAOMmTIMKAAAAAAAAAA==/',
- '_ts': 1723019051,
+ '_etag': '"0000ef9c-0000-2300-0000-66dfd5140000"',
+ '_rid': 'ipcwAJFKMxICAAAAAAAAAA==',
+ '_self': 'dbs/ipcwAA==/colls/ipcwAJFKMxI=/docs/ipcwAJFKMxICAAAAAAAAAA==/',
+ '_ts': 1725945108,
  'content': 'Hello, world!',
- 'id': 'Hello',
+ 'id': 'test',
  'role': 'assistant'}
+
+$ python apps/3_call_azure_cosmos_db/main.py --command delete
+Delete item:
+
+$ python apps/3_call_azure_cosmos_db/main.py --command read
+Read item:
+Failed to read item: (NotFound) Entity with the specified id does not exist in the system. More info: https://aka.ms/cosmosdb-tsg-not-found,
 ```
 
 ## References
 
-- [Get started with Azure Cosmos DB for NoSQL using Python](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-python-get-started?tabs=env-virtual%2Cazure-cli%2Clinux)
-- [Examples for Azure Cosmos DB for NoSQL SDK for Python](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/samples-python)
+- [Get started with Azure Cosmos DB for NoSQL using Python](https://learn.microsoft.com/azure/cosmos-db/nosql/how-to-python-get-started?tabs=env-virtual%2Cazure-cli%2Clinux)
+- [Examples for Azure Cosmos DB for NoSQL SDK for Python](https://learn.microsoft.com/azure/cosmos-db/nosql/samples-python)
 - [Azure Cosmos DB SQL API client library for Python Samples](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples)

apps/3_call_azure_cosmos_db/main.py

@@ -1,26 +1,43 @@
+import argparse
+import logging
 import sys
 from os import getenv
 from pprint import pprint
 
 from azure.cosmos import CosmosClient, PartitionKey
 from dotenv import load_dotenv
 
+
+def init_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="cosmosdb",
+        description="",
+    )
+    parser.add_argument("-i", "--item-id", default="test", help="Item ID")
+    parser.add_argument("-p", "--partition-key-path", default="/id", help="Partition key")
+    parser.add_argument("-v", "--verbose", action="store_true")
+    parser.add_argument("-c", "--command", default="create", help="Command to execute: create, read, delete")
+    return parser.parse_args()
+
+
 if __name__ == "__main__":
+    args = init_args()
+
+    # Set verbose mode
+    if args.verbose:
+        logging.basicConfig(level=logging.DEBUG)
+
     # Parse .env file and get environment variables
     load_dotenv()
 
-    # Get item_id from command line arguments
-    if len(sys.argv) < 2:
-        print("Usage: python main.py <item_id>")
-        sys.exit(1)
-    item_id = sys.argv[1]
+    item_id = args.item_id
 
     # Connect to Azure Cosmos DB
     try:
         client = CosmosClient.from_connection_string(getenv("AZURE_COSMOS_DB_CONNECTION_STRING"))
         database = client.create_database_if_not_exists(id=getenv("AZURE_COSMOS_DB_DATABASE_NAME"))
         partition_key_path = PartitionKey(
-            path="/id",
+            path=args.partition_key_path,
             kind="Hash",
         )
         container = database.create_container_if_not_exists(
@@ -31,24 +48,45 @@
         print(f"Failed to connect to Azure Cosmos DB: {e}")
         sys.exit(1)
 
-    # Create and read an item
-    try:
-        response = container.create_item(
-            body={
-                "id": item_id,
-                "role": "assistant",
-                "content": "Hello, world!",
-            }
-        )
-        print("Created item:")
-        pprint(response)
-        response = container.read_item(
-            item=item_id,
-            partition_key=item_id,
-        )
+    if args.command == "create":
+        # Create and read an item
+        print("Create item")
+        try:
+            response = container.create_item(
+                body={
+                    "id": item_id,
+                    "role": "assistant",
+                    "content": "Hello, world!",
+                }
+            )
+            pprint(response)
+        except Exception as e:
+            print(f"Failed to create item: {e}")
+            sys.exit(1)
+    elif args.command == "read":
+        # Read an item
         print("Read item:")
-        pprint(response)
+        try:
+            response = container.read_item(
+                item=item_id,
+                partition_key=item_id,
+            )
+            pprint(response)
 
-    except Exception as e:
-        print(f"Failed to create item: {e}")
+        except Exception as e:
+            print(f"Failed to read item: {e}")
+            sys.exit(1)
+    elif args.command == "delete":
+        # Delete an item
+        print("Delete item:")
+        try:
+            container.delete_item(
+                item=item_id,
+                partition_key=item_id,
+            )
+        except Exception as e:
+            print(f"Failed to delete item: {e}")
+            sys.exit(1)
+    else:
+        print(f"Invalid command {args.command}")
         sys.exit(1)