Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
+import pandas as pd
 import streamlit as st
-from azure.cosmos import CosmosClient, PartitionKey
+from azure.cosmos import CosmosClient
 from azure.cosmos.exceptions import CosmosResourceNotFoundError
 
 # Initialize Azure Cosmos DB Client
@@ -8,7 +9,7 @@ COSMOS_CONNECTION_STRING = os.getenv('COSMOS_CONNECTION_STRING')
 cosmos_client = CosmosClient.from_connection_string(COSMOS_CONNECTION_STRING)
 
 # Function to Delete an Item in Cosmos DB
-def delete_cosmos_item(db_name, container_name, item_id, partition_key):
+def delete_cosmos_item(db_name, container_name, item_id):
     try:
         database_client = cosmos_client.get_database_client(db_name)
         container_client = database_client.get_container_client(container_name)
@@ -17,9 +18,9 @@ def delete_cosmos_item(db_name, container_name, item_id, partition_key):
     except CosmosResourceNotFoundError:
         st.error(f"Item not found: {item_id}")
 
-# Display and Manage Cosmos DB Structure
-def display_and_manage_cosmos_db():
-    st.subheader('Azure Cosmos DB Structure')
+# Load Cosmos DB Data into a DataFrame
+def load_cosmos_data_into_dataframe():
+    data = []
     for db_properties in cosmos_client.list_databases():
         db_name = db_properties['id']
         database_client = cosmos_client.get_database_client(db_name)
@@ -27,18 +28,25 @@ def display_and_manage_cosmos_db():
             container_name = container_properties['id']
             container_client = database_client.get_container_client(container_name)
             for item in container_client.read_all_items():
-                st.json(item)  # Display the full data of the item
-                if 'file_name' in item and os.path.isfile(item['file_name']):
-                    st.image(item['file_name'])  # Show image if available
-                # Assume the partition key is the item ID
-                partition_key = item.get('id')
-                delete_button_key = f"delete_{item['id']}"
-                if st.button(f"🗑️ Delete {item['id']}", key=delete_button_key):
-                    delete_cosmos_item(db_name, container_name, item['id'], partition_key)
+                data.append(item)
+    return pd.DataFrame(data)
+
+# Display and Manage Cosmos DB Data
+def display_and_manage_cosmos_data(df):
+    st.subheader('Azure Cosmos DB Data')
+    for index, row in df.iterrows():
+        st.json(row.to_json())  # Display the full data of the item
+        if 'file_name' in row and os.path.isfile(row['file_name']):
+            st.image(row['file_name'])  # Show image if available
+        delete_button_key = f"delete_{row['id']}"
+        if st.button(f"🗑️ Delete {row['id']}", key=delete_button_key):
+            delete_cosmos_item(row['_database_id'], row['_container_id'], row['id'])
+            st.experimental_rerun()
 
 # Main UI Function
 def main():
-    display_and_manage_cosmos_db()
+    df = load_cosmos_data_into_dataframe()
+    display_and_manage_cosmos_data(df)
 
 # Run the main function
 if __name__ == "__main__":