Update app.py

app.py

@@ -36,11 +36,9 @@ def display_cosmos_db_structure():
 if st.button('Show Cosmos DB Structure'):
     display_cosmos_db_structure()
 
-
-
 # Function to Add or Update an Item
-def add_or_update_item(database_name, container_name, item_id, item_data):
-    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
+def add_or_update_item(database_client, container_name, item_id, item_data):
+    container = database_client.get_container_client(container_name)
     try:
         existing_item = container.read_item(item_id, partition_key=PartitionKey(item_id))
         existing_item.update(item_data)
@@ -49,29 +47,39 @@ def add_or_update_item(database_name, container_name, item_id, item_data):
         container.create_item(item_data)
 
 # Test Function to Insert PNG Images
-def test_insert_png_images(database_name, container_name):
-    png_files = glob.glob('*.png')
-    for file_name in png_files:
-        item_id = os.path.splitext(file_name)[0]  # Use file name without extension as ID
-        item_data = {"id": item_id, "file_name": file_name}
-        add_or_update_item(database_name, container_name, item_id, item_data)
-        st.write(f"Inserted: {file_name}")
-
-    # Displaying Images
-    st.subheader("Displaying Images in Container")
-    items = list(cosmos_client.get_database_client(database_name).get_container_client(container_name).read_all_items())
-    for item in items:
-        if 'file_name' in item and item['file_name'].endswith('.png'):
-            st.image(item['file_name'], caption=item['file_name'])
+def test_insert_png_images():
+    # Get the first database and container
+    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']
+
+            # Insert PNG files
+            png_files = glob.glob('*.png')
+            for file_name in png_files:
+                item_id = os.path.splitext(file_name)[0]  # Use file name without extension as ID
+                item_data = {"id": item_id, "file_name": file_name}
+                add_or_update_item(database_client, container_name, item_id, item_data)
+                st.write(f"Inserted: {file_name}")
+
+            # Displaying Images
+            st.subheader("Displaying Images in Container")
+            items = list(database_client.get_container_client(container_name).read_all_items())
+            for item in items:
+                if 'file_name' in item and item['file_name'].endswith('.png'):
+                    st.image(item['file_name'], caption=item['file_name'])
+        else:
+            st.error("No container found in the database.")
+    else:
+        st.error("No database found.")
 
 # UI to Test Image Insertion Function
 if st.button('Test Insert PNG Images'):
-    test_database_name = st.text_input('Enter Test Database Name')
-    test_container_name = st.text_input('Enter Test Container Name')
-    if test_database_name and test_container_name:
-        test_insert_png_images(test_database_name, test_container_name)
-    else:
-        st.error('Please enter the test database and container names.')
+    test_insert_png_images()
+
 
 
 # Function to Add an Item