File size: 12,990 Bytes
b05ef13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b8cd397
 
 
 
 
 
 
 
 
 
 
 
 
818d24c
 
 
 
b8cd397
 
818d24c
 
 
 
b8cd397
 
818d24c
 
b8cd397
818d24c
 
 
 
 
 
 
 
 
 
 
 
b8cd397
 
 
 
 
 
 
818d24c
 
 
 
 
b8cd397
818d24c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b05ef13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
818d24c
b05ef13
 
 
 
 
 
818d24c
b05ef13
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
import streamlit as st
from azure.cosmos import CosmosClient, exceptions
import os
import pandas as pd
import json
from datetime import datetime
import shutil
from github import Github
from git import Repo
import base64

# Cosmos DB configuration
ENDPOINT = "https://acae-afd.documents.azure.com:443/"
SUBSCRIPTION_ID = "003fba60-5b3f-48f4-ab36-3ed11bc40816"
DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
Key = os.environ.get("Key")

# GitHub configuration
def download_github_repo(url, local_path):
    if os.path.exists(local_path):
        shutil.rmtree(local_path)
    Repo.clone_from(url, local_path)

def create_zip_file(source_dir, output_filename):
    shutil.make_archive(output_filename, 'zip', source_dir)

def create_repo(g, repo_name):
    user = g.get_user()
    return user.create_repo(repo_name)

def push_to_github(local_path, repo, github_token):
    repo_url = f"https://{github_token}@github.com/{repo.full_name}.git"
    local_repo = Repo(local_path)
    
    if 'origin' in [remote.name for remote in local_repo.remotes]:
        origin = local_repo.remote('origin')
        origin.set_url(repo_url)
    else:
        origin = local_repo.create_remote('origin', repo_url)
    
    if not local_repo.heads:
        local_repo.git.checkout('-b', 'main')
        current_branch = 'main'
    else:
        current_branch = local_repo.active_branch.name
    
    local_repo.git.add(A=True)
    
    if local_repo.is_dirty():
        local_repo.git.commit('-m', 'Initial commit')
    
    origin.push(refspec=f'{current_branch}:{current_branch}')

def get_base64_download_link(file_path, file_name):
    with open(file_path, "rb") as file:
        contents = file.read()
    base64_encoded = base64.b64encode(contents).decode()
    return f'<a href="data:application/zip;base64,{base64_encoded}" download="{file_name}">Download {file_name}</a>'

# Function to fetch all databases and containers
def fetch_db_structure(client):
    structure = {}
    for database in client.list_databases():
        db_name = database['id']
        structure[db_name] = []
        db_client = client.get_database_client(db_name)
        for container in db_client.list_containers():
            structure[db_name].append(container['id'])
    return structure

# Function to fetch items from a container
def fetch_items(client, db_name, container_name):
    container = client.get_database_client(db_name).get_container_client(container_name)
    items = list(container.read_all_items())
    return items

# Function to render sidebar
def render_sidebar(client, structure):
    st.sidebar.title("πŸ“ Cosmos DB Manager")
    
    # Database and Container selection
    selected_db = st.sidebar.selectbox("πŸ—„οΈ Select Database", list(structure.keys()))
    selected_container = st.sidebar.selectbox("πŸ“¦ Select Container", structure[selected_db])
    
    # Fetch items for the selected container
    items = fetch_items(client, selected_db, selected_container)
    
    # Item selection
    item_ids = [item['id'] for item in items]
    selected_item = st.sidebar.selectbox("πŸ“„ Select Item", ["Create New Item"] + item_ids)
    
    # CRUD buttons
    if selected_item != "Create New Item":
        if st.sidebar.button("πŸ–ŠοΈ Edit Item"):
            st.session_state.action = "edit"
            st.session_state.selected_item = selected_item
        if st.sidebar.button("πŸ—‘οΈ Delete Item"):
            st.session_state.action = "delete"
            st.session_state.selected_item = selected_item
    else:
        if st.sidebar.button("βž• Create New Item"):
            st.session_state.action = "create"
    
    # GitHub section
    st.sidebar.subheader("πŸ™ GitHub Operations")
    github_token = os.environ.get("GITHUB")
    source_repo = st.sidebar.text_input("πŸ”— Source Repo URL", value="https://github.com/AaronCWacker/AIExamples-8-24-Streamlit")
    new_repo_name = st.sidebar.text_input("πŸ“ New Repo Name", value=f"AIExample-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    
    if st.sidebar.button("πŸ“₯ Clone Repository"):
        clone_github_repo(github_token, source_repo, new_repo_name)
    
    if st.sidebar.button("πŸ“€ Push to New Repository"):
        push_to_new_repo(github_token, source_repo, new_repo_name)
    
    # Archive data button
    if st.sidebar.button("πŸ“¦ Archive All Data"):
        download_link = archive_all_data(client)
        st.sidebar.markdown(download_link, unsafe_allow_html=True)
    
    # Logout button
    if st.sidebar.button("πŸšͺ Logout"):
        st.session_state.logged_in = False
        st.session_state.action = None
        st.session_state.selected_item = None
        st.rerun()
    
    return selected_db, selected_container, selected_item

# GitHub operations
def clone_github_repo(github_token, source_repo, new_repo_name):
    if github_token and source_repo:
        try:
            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            download_github_repo(source_repo, local_path)
            zip_filename = f"{new_repo_name}.zip"
            create_zip_file(local_path, zip_filename[:-4])
            st.sidebar.markdown(get_base64_download_link(zip_filename, zip_filename), unsafe_allow_html=True)
            st.sidebar.success("Repository cloned successfully!")
        except Exception as e:
            st.sidebar.error(f"An error occurred: {str(e)}")
        finally:
            if os.path.exists(local_path):
                shutil.rmtree(local_path)
            if os.path.exists(zip_filename):
                os.remove(zip_filename)
    else:
        st.sidebar.error("Please ensure GitHub token is set in environment variables and source repository URL is provided.")

def push_to_new_repo(github_token, source_repo, new_repo_name):
    if github_token and source_repo:
        try:
            g = Github(github_token)
            new_repo = create_repo(g, new_repo_name)
            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            download_github_repo(source_repo, local_path)
            push_to_github(local_path, new_repo, github_token)
            st.sidebar.success(f"Repository pushed successfully to {new_repo.html_url}")
        except Exception as e:
            st.sidebar.error(f"An error occurred: {str(e)}")
        finally:
            if os.path.exists(local_path):
                shutil.rmtree(local_path)
    else:
        st.sidebar.error("Please ensure GitHub token is set in environment variables and source repository URL is provided.")

# Function to archive all data
def archive_all_data(client):
    try:
        base_dir = "./cosmos_archive"
        if os.path.exists(base_dir):
            shutil.rmtree(base_dir)
        os.makedirs(base_dir)

        for database in client.list_databases():
            db_name = database['id']
            db_dir = os.path.join(base_dir, db_name)
            os.makedirs(db_dir)

            db_client = client.get_database_client(db_name)
            for container in db_client.list_containers():
                container_name = container['id']
                container_dir = os.path.join(db_dir, container_name)
                os.makedirs(container_dir)

                container_client = db_client.get_container_client(container_name)
                items = list(container_client.read_all_items())

                with open(os.path.join(container_dir, f"{container_name}.json"), 'w') as f:
                    json.dump(items, f, indent=2)

        archive_name = f"cosmos_archive_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        shutil.make_archive(archive_name, 'zip', base_dir)
        
        return get_base64_download_link(f"{archive_name}.zip", f"{archive_name}.zip")
    except Exception as e:
        return f"An error occurred while archiving data: {str(e)}"

# New function to convert text to JSON
def text_to_json(text):
    lines = text.split('\n')
    json_data = {}
    for line in lines:
        if ':' in line:
            key, value = line.split(':', 1)
            json_data[key.strip()] = value.strip()
        else:
            # If there's no colon, use the whole line as a key with an empty string as value
            json_data[line.strip()] = ""
    return json_data

# Cosmos DB operations
def create_item(client, db_name, container_name):
    st.subheader("Create New Item")
    item_id = st.text_input("Item ID")
    item_data = st.text_area("Item Data (Text format)", 
                             value="🌟 Cosmos DB and GitHub Integration\nCreate New Item\nItem ID\nItem Data (Text format)\nπŸ’Ύ Save New Item")
    
    if st.button("πŸ’Ύ Save New Item"):
        try:
            container = client.get_database_client(db_name).get_container_client(container_name)
            json_data = text_to_json(item_data)
            new_item = {"id": item_id, "content": json_data}
            container.create_item(body=new_item)
            st.success("New item created successfully!")
            st.json(new_item)  # Display the created item in JSON format
            st.rerun()
        except Exception as e:
            st.error(f"Error creating item: {str(e)}")

def edit_item(client, db_name, container_name, item_id):
    container = client.get_database_client(db_name).get_container_client(container_name)
    item = container.read_item(item=item_id, partition_key=item_id)
    
    st.subheader(f"Editing Item: {item_id}")
    edited_item = {}
    for key, value in item.items():
        if key not in ['_rid', '_self', '_etag', '_attachments', '_ts']:
            if key == 'content':
                # Convert the JSON content back to text for editing
                text_content = "\n".join([f"{k}: {v}" for k, v in value.items()])
                edited_text = st.text_area(key, text_content)
                edited_item[key] = text_to_json(edited_text)
            else:
                edited_item[key] = st.text_input(key, value)
    
    if st.button("πŸ’Ύ Save Changes"):
        try:
            container.upsert_item(body=edited_item)
            st.success("Item updated successfully!")
            st.json(edited_item)  # Display the updated item in JSON format
            st.rerun()
        except Exception as e:
            st.error(f"Error updating item: {str(e)}")

def delete_item(client, db_name, container_name, item_id):
    st.subheader(f"Delete Item: {item_id}")
    st.write("Are you sure you want to delete this item?")
    if st.button("πŸ—‘οΈ Confirm Delete"):
        try:
            container = client.get_database_client(db_name).get_container_client(container_name)
            container.delete_item(item=item_id, partition_key=item_id)
            st.success(f"Item {item_id} deleted successfully!")
            st.rerun()
        except Exception as e:
            st.error(f"Error deleting item: {str(e)}")

# Main Streamlit app
def main():
    st.set_page_config(layout="wide")
    st.title("🌟 Cosmos DB and GitHub Integration")

    # Initialize session state
    if 'logged_in' not in st.session_state:
        st.session_state.logged_in = False
    if 'action' not in st.session_state:
        st.session_state.action = None
    if 'selected_item' not in st.session_state:
        st.session_state.selected_item = None

    # Login section
    if not st.session_state.logged_in:
        st.subheader("πŸ” Login")
        input_key = Key  # Get the key from the environment variable
        if st.button("πŸš€ Login"):
            if input_key:
                st.session_state.primary_key = input_key
                st.session_state.logged_in = True
                st.rerun()
            else:
                st.error("Invalid key. Please check your environment variables.")
    else:
        # Initialize Cosmos DB client
        try:
            client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
            
            # Fetch database structure
            structure = fetch_db_structure(client)
            
            # Render sidebar and get selections
            selected_db, selected_container, selected_item = render_sidebar(client, structure)
            
            # Main area
            if st.session_state.action == "create":
                create_item(client, selected_db, selected_container)
            elif st.session_state.action == "edit":
                edit_item(client, selected_db, selected_container, st.session_state.selected_item)
            elif st.session_state.action == "delete":
                delete_item(client, selected_db, selected_container, st.session_state.selected_item)
            else:
                st.write("Select an action from the sidebar to get started.")

        except exceptions.CosmosHttpResponseError as e:
            st.error(f"Failed to connect to Cosmos DB. HTTP error: {str(e)}. Status code: {e.status_code}")
        except Exception as e:
            st.error(f"An unexpected error occurred: {str(e)}")

if __name__ == "__main__":
    main()