Code Editor for app.py

import streamlit as st
import duckdb
import pandas as pd

# Initialize DuckDB connection
conn = duckdb.connect()

def get_tables():
    # Get all tables from BLUE schema
    query = """
    SELECT table_name 
    FROM information_schema.tables 
    WHERE table_schema = 'BLUE'
    """
    try:
        tables = conn.execute(query).fetchall()
        return [table[0] for table in tables]
    except:
        # If schema doesn't exist, create it with a sample table
        conn.execute("CREATE SCHEMA IF NOT EXISTS BLUE")
        conn.execute("""
        CREATE TABLE IF NOT EXISTS BLUE.sample_table (
            id INTEGER,
            name VARCHAR,
            value DOUBLE
        )
        """)
        conn.execute("INSERT INTO BLUE.sample_table VALUES (1, 'test', 1.0), (2, 'test2', 2.0)")
        return ['sample_table']

def get_columns(table_name):
    query = f"""
    SELECT column_name 
    FROM information_schema.columns 
    WHERE table_schema = 'BLUE' 
    AND table_name = '{table_name}'
    """
    columns = conn.execute(query).fetchall()
    return [col[0] for col in columns]

# Set up the Streamlit app
st.title("DuckDB Table Viewer")

# Sidebar with table selection
with st.sidebar:
    tables = get_tables()
    selected_table = st.selectbox("Select a table", tables)

if selected_table:
    # Get columns for the selected table
    columns = get_columns(selected_table)
    
    # Create column selection
    st.subheader("Select columns to display")
    
    # Initialize session state for column selection if not exists
    if 'selected_columns' not in st.session_state:
        st.session_state.selected_columns = set()

    # Add "All" and "None" checkboxes
    col1, col2 = st.columns(2)
    
    with col1:
        if st.checkbox("Select All"):
            st.session_state.selected_columns = set(columns)
    
    with col2:
        if st.checkbox("Select None"):
            st.session_state.selected_columns = set()

    # Individual column checkboxes
    for col in columns:
        if st.checkbox(col, value=col in st.session_state.selected_columns):
            st.session_state.selected_columns.add(col)
        else:
            st.session_state.selected_columns.discard(col)

    # Display the table
    if st.session_state.selected_columns:
        selected_cols = list(st.session_state.selected_columns)
        query = f"""
        SELECT {', '.join(selected_cols)}
        FROM BLUE.{selected_table}
        LIMIT 10
        """
        df = conn.execute(query).df()
        st.subheader(f"First 10 records of {selected_table}")
        st.dataframe(df)
    else:
        st.warning("Please select at least one column to display.")