Update app.py

app.py

@@ -1,64 +1,128 @@
 import os
 import streamlit as st
 from azure.cosmos import CosmosClient, PartitionKey
+from azure.storage.blob import BlobServiceClient
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+import requests
 import glob
 
 # Environment Variables
 COSMOS_CONNECTION_STRING = os.getenv('COSMOS_CONNECTION_STRING')
+BLOB_STORAGE_CONNECTION_STRING = os.getenv('BLOB_STORAGE_CONNECTION_STRING')
 
-# Initialize Azure Cosmos DB Client
+# Initialize Azure Clients
 cosmos_client = CosmosClient.from_connection_string(COSMOS_CONNECTION_STRING)
+blob_service = BlobServiceClient.from_connection_string(BLOB_STORAGE_CONNECTION_STRING)
 
-# Function to Retrieve, Display, and Manage Cosmos DB Images
-def manage_cosmos_db_images():
-    st.subheader('Manage Cosmos DB Images')
+@st.cache(allow_output_mutation=True)
+def get_cosmos_client():
+    return cosmos_client
+
+def get_database_client(db_name):
+    cosmos_client = get_cosmos_client()
+    return cosmos_client.get_database_client(db_name)
+
+def get_container_client(db_name, container_name):
+    database_client = get_database_client(db_name)
+    return database_client.get_container_client(container_name)
+
+def list_databases():
+    cosmos_client = get_cosmos_client()
+    return list(cosmos_client.list_databases())
+
+def list_containers(db_name):
+    database_client = get_database_client(db_name)
+    return list(database_client.list_containers())
+
+def list_items(db_name, container_name):
+    container_client = get_container_client(db_name, container_name)
+    return list(container_client.read_all_items())
+
+def create_item(db_name, container_name, item_data):
+    container_client = get_container_client(db_name, container_name)
+    container_client.create_item(body=item_data)
+
+def update_item(db_name, container_name, item_id, new_data):
+    container_client = get_container_client(db_name, container_name)
+    try:
+        item = container_client.read_item(item_id, partition_key=item_id)
+        for key in new_data:
+            item[key] = new_data[key]
+        container_client.replace_item(item_id, item)
+    except CosmosResourceNotFoundError:
+        st.error("Item not found for update")
+
+def delete_item(db_name, container_name, item_id):
+    container_client = get_container_client(db_name, container_name)
+    try:
+        container_client.delete_item(item_id, partition_key=item_id)
+    except CosmosResourceNotFoundError:
+        st.error("Item not found for deletion")
+
+def display_cosmos_db_structure():
+    # Existing function logic here
+
+def manage_png_images():
+    # Get database and container names
     db_properties = next(cosmos_client.list_databases(), None)
-    if db_properties:
-        db_name = db_properties['id']
-        database_client = cosmos_client.get_database_client(db_name)
-        container_properties = next(database_client.list_containers(), None)
-        if container_properties:
-            container_name = container_properties['id']
-            container_client = database_client.get_container_client(container_name)
-            items = list(container_client.read_all_items())
-            
-            # Display existing items with delete option
-            for item in items:
-                item_desc = f"Item: `{item['id']}`"
-                if 'file_name' in item and item['file_name'].endswith('.png'):
-                    st.markdown(item_desc)
-                    st.image(item['file_name'])
-                    if st.button(f"🗑️ Delete {item['id']}", key=f"delete_{item['id']}"):
-                        container_client.delete_item(item=item, partition_key=item['id'])
-                        st.success(f"Deleted Item: {item['id']}")
-
-            # Add or update PNG files from directory
-            png_files = glob.glob('*.png')
-            for file_name in png_files:
-                item_id = os.path.splitext(file_name)[0]
-                item_data = {"id": item_id, "file_name": file_name}
-                if item_id not in {item['id'] for item in items}:
-                    container_client.create_item(body=item_data)
-                    st.write(f"Added Item: {item_id}")
-
-# UI to Manage PNG Images
-st.title('Manage PNG Images in Cosmos DB')
-if st.button('Manage PNG Images'):
-    manage_cosmos_db_images()
+    if not db_properties:
+        st.error("No database found.")
+        return
+
+    db_name = db_properties['id']
+    database_client = get_database_client(db_name)
+    container_properties = next(database_client.list_containers(), None)
+    if not container_properties:
+        st.error("No container found.")
+        return
+
+    container_name = container_properties['id']
+    container_client = get_container_client(db_name, container_name)
+
+    # List existing items
+    existing_items = list_items(db_name, container_name)
+    existing_ids = {item['id']: item for item in existing_items}
+
+    # Display existing items
+    st.subheader("Existing PNG Images")
+    for item_id, item in existing_ids.items():
+        st.markdown(f"**{item_id}**: {item.get('file_name', 'No file name')}")
+        if st.button(f"🗑️ Delete {item_id}", key=f"delete_{item_id}"):
+            delete_item(db_name, container_name, item_id)
+            st.success(f"Deleted Item: {item_id}")
 
-# Azure Blob Storage - Upload/Download
-st.subheader('Azure Blob Storage - Upload/Download')
-blob_container = st.text_input('Blob Container')
-blob_file = st.file_uploader('Upload file to Blob')
+    # Add or Update PNG files
+    st.subheader("Add or Update PNG Images")
+    png_files = glob.glob('*.png')
+    for file_name in png_files:
+        item_id = os.path.splitext(file_name)[0]
+        item_data = {"id": item_id, "file_name": file_name}
+        if item_id not in existing_ids:
+            if st.button(f"➕ Add {item_id}", key=f"add_{item_id}"):
+                create_item(db_name, container_name, item_data)
+                st.success(f"Added Item: {item_id}")
+        else:
+            if st.button(f"✏️ Update {item_id}", key=f"update_{item_id}"):
+                update_item(db_name, container_name, item_id, {"file_name": file_name})
+                st.success(f"Updated Item: {item_id}")
 
-if blob_file is not None and st.button('Upload to Blob'):
-    blob_client = blob_service.get_blob_client(container=blob_container, blob=blob_file.name)
-    blob_client.upload_blob(blob_file.getvalue())
 
-# Azure Functions - Trigger
-st.subheader('Azure Functions - Trigger')
-function_url = st.text_input('Function URL')
+def list_blobs_in_container(container_name):
+    container_client = blob_service.get_container_client(container_name)
+    return list(container_client.list_blobs())
+
+# UI Elements
+if st.button('Show Cosmos DB Structure'):
+    display_cosmos_db_structure()
+
+if st.button('Manage PNG Images'):
+    manage_png_images()
+
+# Listing Blobs in a Container
+blob_container = st.text_input('Enter Blob Container Name')
+if st.button('List Blobs in Container'):
+    blobs = list_blobs_in_container(blob_container)
+    for blob in blobs:
+        st.write(blob.name)
 
-if st.button('Call Azure Function'):
-    response = requests.get(function_url)
-    st.write('Function Response:', response.text)
+# Remaining UI for Azure Blob Storage and Azure Functions as in the original code