app.py

from flask import Flask, render_template, request, jsonify
import sqlite3
import os
import csv
import re
from werkzeug.utils import secure_filename

app = Flask(__name__)
DB_FILE = 'sustainable_energy.db'
SOURCE_DATA_DIR = 'source_data'
MAX_FILE_SIZE = 50 * 1024 * 1024  # 50MB limit

def get_db_connection():
    """Get database connection with proper configuration"""
    conn = sqlite3.connect(DB_FILE)
    conn.row_factory = sqlite3.Row
    # Set busy timeout and WAL mode to reduce locking issues
    conn.execute('PRAGMA busy_timeout = 30000')
    conn.execute('PRAGMA journal_mode = WAL')
    return conn

def validate_table_name(table_name):
    """Validate table name to prevent SQL injection"""
    if not table_name:
        return False
    # Only allow alphanumeric and underscores
    return bool(re.match(r'^[a-zA-Z0-9_]+$', table_name))

def sanitize_column_name(header):
    """Sanitize CSV header to valid SQL column name"""
    # Clean headers - remove newline characters
    header = header.replace('\n', ' ').strip()
    # Replace spaces and special characters with underscores
    sql_col = header.replace(' ', '_').replace('(', '').replace(')', '').replace('%', 'pct').replace('$', 'USD').replace('/', '_per_')
    sql_col = sql_col.replace(',', '').replace('-', '_')
    # Remove any remaining special characters
    sql_col = ''.join(c if c.isalnum() or c == '_' else '_' for c in sql_col)
    return sql_col

def detect_column_type(values):
    """Detect SQL column type from sample values"""
    if not values:
        return 'TEXT'
    
    # Check if all non-empty values are integers
    all_int = True
    has_int = False
    for val in values:
        if val and val.strip():
            has_int = True
            try:
                int(val)
            except (ValueError, TypeError):
                all_int = False
                break
    
    if all_int and has_int:
        return 'INTEGER'
    
    # Check if all non-empty values are numeric
    all_numeric = True
    has_numeric = False
    for val in values:
        if val and val.strip():
            has_numeric = True
            try:
                float(val)
            except (ValueError, TypeError):
                all_numeric = False
                break
    
    if all_numeric and has_numeric:
        return 'REAL'
    
    return 'TEXT'

def execute_query(query):
    """Execute SQL query and return results"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        
        # Only allow SELECT queries for safety
        query_upper = query.strip().upper()
        if not query_upper.startswith('SELECT'):
            conn.close()
            return None, None, "Only SELECT queries are allowed for safety."
        
        cursor.execute(query)
        rows = cursor.fetchall()
        
        # Get column names
        columns = [description[0] for description in cursor.description] if cursor.description else []
        
        # Convert rows to list of dictionaries
        results = [dict(row) for row in rows]
        
        conn.close()
        return results, columns, None
    except Exception as e:
        return None, None, str(e)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/tables', methods=['GET'])
def get_tables():
    """Get list of all tables in the database"""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name")
        tables = [row[0] for row in cursor.fetchall()]
        conn.close()
        return jsonify({'tables': tables})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/upload-csv', methods=['POST'])
def upload_csv():
    """Upload and store CSV file"""
    try:
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        # Check file extension
        if not file.filename.lower().endswith('.csv'):
            return jsonify({'error': 'Only CSV files are allowed'}), 400
        
        # Check file size
        file.seek(0, os.SEEK_END)
        file_size = file.tell()
        file.seek(0)
        if file_size > MAX_FILE_SIZE:
            return jsonify({'error': f'File size exceeds {MAX_FILE_SIZE / 1024 / 1024}MB limit'}), 400
        
        # Ensure source_data directory exists
        os.makedirs(SOURCE_DATA_DIR, exist_ok=True)
        
        # Check for duplicate
        filename = secure_filename(file.filename)
        filepath = os.path.join(SOURCE_DATA_DIR, filename)
        
        if os.path.exists(filepath):
            return jsonify({'error': 'File with this name already exists'}), 400
        
        # Save file
        file.save(filepath)
        
        return jsonify({
            'success': True,
            'filename': filename,
            'message': 'File uploaded successfully'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/preview-csv', methods=['POST'])
def preview_csv():
    """Preview CSV file before import"""
    try:
        data = request.get_json()
        filename = data.get('filename', '')
        
        if not filename:
            return jsonify({'error': 'No filename provided'}), 400
        
        filepath = os.path.join(SOURCE_DATA_DIR, secure_filename(filename))
        
        if not os.path.exists(filepath):
            return jsonify({'error': 'File not found'}), 404
        
        headers = []
        sample_rows = []
        column_types = []
        
        with open(filepath, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            headers = next(reader)
            headers = [h.replace('\n', ' ').strip() for h in headers]
            
            # Read first 20 rows for preview
            for i, row in enumerate(reader):
                if i >= 20:
                    break
                sample_rows.append(row)
            
            # Detect column types from sample data
            if sample_rows:
                for col_idx in range(len(headers)):
                    column_values = [row[col_idx] if col_idx < len(row) else '' for row in sample_rows]
                    col_type = detect_column_type(column_values)
                    column_types.append(col_type)
            else:
                column_types = ['TEXT'] * len(headers)
        
        return jsonify({
            'headers': headers,
            'sample_rows': sample_rows[:20],
            'column_types': column_types,
            'row_count': len(sample_rows)
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/create-table', methods=['POST'])
def create_table():
    """Create table from CSV file"""
    try:
        data = request.get_json()
        filename = data.get('filename', '')
        
        if not filename:
            return jsonify({'error': 'No filename provided'}), 400
        
        # Derive table name from filename
        table_name = os.path.splitext(secure_filename(filename))[0]
        table_name = re.sub(r'[^a-zA-Z0-9_]', '_', table_name)
        
        if not validate_table_name(table_name):
            return jsonify({'error': 'Invalid table name generated from filename'}), 400
        
        filepath = os.path.join(SOURCE_DATA_DIR, secure_filename(filename))
        
        if not os.path.exists(filepath):
            return jsonify({'error': 'File not found'}), 404
        
        # Check if table already exists
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
        if cursor.fetchone():
            conn.close()
            return jsonify({'error': f'Table "{table_name}" already exists'}), 400
        
        # Read CSV and create table
        with open(filepath, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            headers = next(reader)
            headers = [h.replace('\n', ' ').strip() for h in headers]
            
            # Sanitize column names
            sql_columns = [sanitize_column_name(h) for h in headers]
            
            # Read all rows to detect types properly
            all_rows = list(reader)
            
            # Detect column types
            column_types = []
            for col_idx in range(len(headers)):
                column_values = [row[col_idx] if col_idx < len(row) else '' for row in all_rows]
                col_type = detect_column_type(column_values)
                column_types.append(col_type)
            
            # Create table statement
            create_table_sql = f"CREATE TABLE {table_name} (id INTEGER PRIMARY KEY AUTOINCREMENT"
            for col, col_type in zip(sql_columns, column_types):
                create_table_sql += f",\n    {col} {col_type}"
            create_table_sql += "\n)"
            
            cursor.execute(create_table_sql)
            
            # Prepare insert statement
            placeholders = ','.join(['?' for _ in headers])
            insert_sql = f"INSERT INTO {table_name} ({','.join(sql_columns)}) VALUES ({placeholders})"
            
            # Insert data
            rows_inserted = 0
            for row in all_rows:
                processed_row = []
                for i, value in enumerate(row):
                    if i >= len(column_types):
                        break
                    if value == '' or value is None:
                        processed_row.append(None)
                    elif column_types[i] == 'INTEGER':
                        try:
                            processed_row.append(int(value))
                        except (ValueError, TypeError):
                            processed_row.append(None)
                    elif column_types[i] == 'REAL':
                        try:
                            processed_row.append(float(value))
                        except (ValueError, TypeError):
                            processed_row.append(None)
                    else:
                        processed_row.append(value)
                
                # Pad row if needed
                while len(processed_row) < len(sql_columns):
                    processed_row.append(None)
                
                cursor.execute(insert_sql, processed_row[:len(sql_columns)])
                rows_inserted += 1
        
        conn.commit()
        conn.close()
        
        return jsonify({
            'success': True,
            'table_name': table_name,
            'rows_inserted': rows_inserted,
            'message': f'Table "{table_name}" created successfully with {rows_inserted} rows'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/table/<table_name>', methods=['DELETE'])
def delete_table(table_name):
    """Delete table from database"""
    try:
        if not validate_table_name(table_name):
            return jsonify({'error': 'Invalid table name'}), 400
        
        # Prevent deletion of system tables
        if table_name.startswith('sqlite_'):
            return jsonify({'error': 'Cannot delete system tables'}), 400
        
        conn = get_db_connection()
        cursor = conn.cursor()
        
        # Check if table exists
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
        if not cursor.fetchone():
            conn.close()
            return jsonify({'error': f'Table "{table_name}" does not exist'}), 404
        
        # Drop table
        cursor.execute(f"DROP TABLE {table_name}")
        conn.commit()
        conn.close()
        
        return jsonify({
            'success': True,
            'message': f'Table "{table_name}" deleted successfully'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/query', methods=['POST'])
def query():
    data = request.get_json()
    sql_query = data.get('query', '')
    
    if not sql_query:
        return jsonify({'error': 'No query provided'}), 400
    
    results, columns, error = execute_query(sql_query)
    
    if error:
        return jsonify({'error': error}), 400
    
    return jsonify({
        'columns': columns,
        'results': results,
        'row_count': len(results) if results else 0
    })

if __name__ == '__main__':
    # Check if database exists
    if not os.path.exists(DB_FILE):
        print(f"Error: Database file '{DB_FILE}' not found!")
        print("Please run 'python3 create_database.py' first.")
    else:
        print(f"Starting server...")
        print(f"Open http://localhost:8080 in your browser")
        app.run(debug=True, host='127.0.0.1', port=8080)