Update app.py

app.py

@@ -155,7 +155,7 @@ def update_record(container, updated_record):
 def delete_record(container, record):
     try:
         doc_id = record["id"]
-        partition_key_value = record.get("pk", doc_id)  # Fallback to id if pk missing
+        partition_key_value = record.get("pk", doc_id)
         st.write(f"Deleting {doc_id} with partition key {partition_key_value}")  # Debug log
         container.delete_item(item=doc_id, partition_key=partition_key_value)
         return True, f"Record {doc_id} deleted. 🗑️"
@@ -353,6 +353,11 @@ def edit_all_documents(container):
     if not documents:
         st.info("No documents in this container.")
         return
+    
+    # Initialize saved_docs if not present
+    if 'saved_docs' not in st.session_state:
+        st.session_state.saved_docs = {}
+    
     for doc in documents:
         ts = doc.get("_ts", 0)
         dt = datetime.fromtimestamp(ts) if ts else datetime.now()
@@ -360,26 +365,24 @@ def edit_all_documents(container):
         header = f"{doc.get('name', 'Unnamed')} - {formatted_ts}"
         with st.expander(header):
             doc_key = f"editor_{doc['id']}"
-            if doc_key not in st.session_state:
-                st.session_state[doc_key] = json.dumps(doc, indent=2)
-            edited_content = st.text_area("Edit JSON", value=st.session_state[doc_key], height=300, key=doc_key)
+            # Use saved_docs for initial value if available, else original doc
+            initial_value = st.session_state.saved_docs.get(doc['id'], json.dumps(doc, indent=2))
+            edited_content = st.text_area("Edit JSON", value=initial_value, height=300, key=doc_key)
             col_save, col_delete = st.columns(2)
             with col_save:
                 if st.button("💾 Save", key=f"save_{doc['id']}"):
                     try:
-                        # Parse edited content
                         cleaned_content = sanitize_json_text(edited_content)
                         updated_doc = json.loads(cleaned_content)
-                        # Preserve identity fields
                         updated_doc['id'] = doc['id']
                         updated_doc['pk'] = doc.get('pk', doc['id'])
-                        # Remove system fields
                         for field in ['_ts', '_rid', '_self', '_etag', '_attachments']:
                             updated_doc.pop(field, None)
                         success, message = update_record(container, updated_doc)
                         if success:
                             st.success(f"Saved {doc['id']}")
-                            st.session_state[doc_key] = json.dumps(updated_doc, indent=2)
+                            st.session_state.saved_docs[doc['id']] = json.dumps(updated_doc, indent=2)
+                            st.rerun()  # Refresh UI to reflect saved state
                         else:
                             st.error(message)
                     except json.JSONDecodeError as e:
@@ -391,8 +394,8 @@ def edit_all_documents(container):
                     success, message = delete_record(container, doc)
                     if success:
                         st.success(message)
-                        if doc_key in st.session_state:
-                            del st.session_state[doc_key]  # Clean up session state
+                        if doc['id'] in st.session_state.saved_docs:
+                            del st.session_state.saved_docs[doc['id']]
                         st.rerun()
                     else:
                         st.error(message)