Main.py

from flask import Flask, render_template, request, jsonify, session, send_file
import os
from werkzeug.utils import secure_filename
import secrets
import json
import atexit
from datetime import timedelta
from database_handler import get_database_handler

app = Flask(__name__)
app.secret_key = secrets.token_hex(16)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 50MB max file size
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=2)  # Session expires after 2 hours
ALLOWED_EXTENSIONS = {'db', 'sqlite', 'sqlite3'}

# Create uploads folder if it doesn't exist
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# Track uploaded files and database connections per session
uploaded_files = {}
active_connections = {}

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def cleanup_session_file(session_id):
    """Delete the uploaded file associated with a session"""
    # Close any active database connection
    if session_id in active_connections:
        try:
            active_connections[session_id].disconnect()
        except Exception as e:
            print(f"Error closing connection: {e}")
        del active_connections[session_id]
    
    # Delete uploaded file
    if session_id in uploaded_files:
        filepath = uploaded_files[session_id]
        if os.path.exists(filepath):
            try:
                os.remove(filepath)
                print(f"Cleaned up file: {filepath}")
            except Exception as e:
                print(f"Error cleaning up file {filepath}: {e}")
        del uploaded_files[session_id]

def cleanup_all_uploads():
    """Clean up all uploaded files and connections on server shutdown"""
    for session_id in list(uploaded_files.keys()):
        cleanup_session_file(session_id)

# Register cleanup function to run on server shutdown
atexit.register(cleanup_all_uploads)

def get_db_handler(session_id):
    """Get or create database handler for session"""
    if session_id in active_connections:
        return active_connections[session_id]
    
    # Get connection params from session
    db_type = session.get('db_type', 'sqlite')
    
    if db_type == 'sqlite':
        connection_params = {
            'db_type': 'sqlite',
            'filepath': session.get('current_db')
        }
    else:
        connection_params = session.get('connection_params', {})
        connection_params['db_type'] = db_type
    
    # Create and connect handler
    try:
        handler = get_database_handler(connection_params)
        handler.connect()
        active_connections[session_id] = handler
        return handler
    except Exception as e:
        raise Exception(f"Failed to connect to database: {str(e)}")

@app.route('/')
def index():
    # Make session permanent (but still expires after PERMANENT_SESSION_LIFETIME)
    session.permanent = True
    
    # Generate a unique session ID if not exists
    if 'session_id' not in session:
        session['session_id'] = secrets.token_hex(16)
    
    return render_template('main.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    """Upload SQLite database file"""
    if 'database' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['database']
    
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    
    if file and allowed_file(file.filename):
        # Clean up previous connection and file if exists
        session_id = session.get('session_id')
        if session_id in uploaded_files:
            cleanup_session_file(session_id)
        
        # Save new file with unique name
        filename = secure_filename(file.filename)
        unique_filename = f"{session_id}_{filename}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
        file.save(filepath)
        
        # Store in session
        session['db_type'] = 'sqlite'
        session['current_db'] = filepath
        session['original_filename'] = filename
        uploaded_files[session_id] = filepath
        
        try:
            # Get table names
            handler = get_db_handler(session_id)
            tables = handler.get_tables()
            
            return jsonify({
                'success': True,
                'filename': filename,
                'tables': tables,
                'db_type': 'sqlite'
            })
        except Exception as e:
            cleanup_session_file(session_id)
            return jsonify({'error': f'Error reading database: {str(e)}'}), 400
    
    return jsonify({'error': 'Invalid file type. Please upload a .db, .sqlite, or .sqlite3 file'}), 400

@app.route('/connect', methods=['POST'])
def connect_database():
    """Connect to remote database (MySQL, PostgreSQL, MongoDB, SQL Server)"""
    data = request.json
    db_type = data.get('db_type', '').lower()
    
    if not db_type:
        return jsonify({'error': 'Database type is required'}), 400
    
    session_id = session.get('session_id')
    if not session_id:
        return jsonify({'error': 'Invalid session'}), 400
    
    # Clean up previous connection
    if session_id in active_connections:
        cleanup_session_file(session_id)
    
    # Build connection parameters
    connection_params = {
        'db_type': db_type,
        'host': data.get('host', 'localhost'),
        'database': data.get('database'),
        'user': data.get('user'),
        'password': data.get('password')
    }
    
    # Add port if specified
    if data.get('port'):
        try:
            connection_params['port'] = int(data['port'])
        except ValueError:
            return jsonify({'error': 'Invalid port number'}), 400
    
    # For MongoDB, also support connection string
    if db_type in ['mongodb', 'mongo'] and data.get('connection_string'):
        connection_params['connection_string'] = data['connection_string']
    
    # Store in session
    session['db_type'] = db_type
    session['connection_params'] = connection_params
    
    try:
        # Test connection and get tables
        handler = get_db_handler(session_id)
        tables = handler.get_tables()
        
        return jsonify({
            'success': True,
            'tables': tables,
            'db_type': db_type,
            'message': f'Successfully connected to {db_type.upper()} database'
        })
    except Exception as e:
        cleanup_session_file(session_id)
        return jsonify({'error': f'Connection failed: {str(e)}'}), 400

@app.route('/tables', methods=['GET'])
def get_tables():
    """Get list of tables"""
    session_id = session.get('session_id')
    if not session_id:
        return jsonify({'error': 'No database loaded'}), 400
    
    try:
        handler = get_db_handler(session_id)
        tables = handler.get_tables()
        return jsonify({'tables': tables})
    except Exception as e:
        return jsonify({'error': str(e)}), 400

@app.route('/table/<table_name>', methods=['GET'])
def get_table_data(table_name):
    """Get paginated table data"""
    session_id = session.get('session_id')
    if not session_id:
        return jsonify({'error': 'No database loaded'}), 400
    
    try:
        page = int(request.args.get('page', 1))
        per_page = int(request.args.get('per_page', 50))
        
        handler = get_db_handler(session_id)
        result = handler.get_table_data(table_name, page, per_page)
        
        return jsonify(result)
    except Exception as e:
        return jsonify({'error': str(e)}), 400

@app.route('/export/<table_name>', methods=['GET'])
def export_table_json(table_name):
    """Export a table as JSON file"""
    session_id = session.get('session_id')
    if not session_id:
        return jsonify({'error': 'No database loaded'}), 400
    
    try:
        handler = get_db_handler(session_id)
        
        # Get table info and all data
        table_info = handler.get_table_info(table_name)
        data = handler.get_all_table_data(table_name)
        
        # Create export data structure
        export_data = {
            'table_name': table_name,
            'database_type': session.get('db_type', 'unknown'),
            'columns': table_info['columns'],
            'row_count': len(data),
            'data': data
        }
        
        # Save to temporary JSON file
        export_filename = f"{session_id}_{table_name}_export.json"
        export_path = os.path.join(app.config['UPLOAD_FOLDER'], export_filename)
        
        with open(export_path, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, indent=2, ensure_ascii=False, default=str)
        
        # Send file and schedule cleanup
        response = send_file(
            export_path,
            as_attachment=True,
            download_name=f"{table_name}_export.json",
            mimetype='application/json'
        )
        
        # Clean up the export file after sending
        @response.call_on_close
        def cleanup_export():
            try:
                if os.path.exists(export_path):
                    os.remove(export_path)
            except Exception as e:
                print(f"Error cleaning up export file: {e}")
        
        return response
        
    except Exception as e:
        return jsonify({'error': str(e)}), 400

@app.route('/disconnect', methods=['POST'])
def disconnect_database():
    """Disconnect from current database"""
    session_id = session.get('session_id')
    if session_id:
        cleanup_session_file(session_id)
        # Clear session data
        session.pop('db_type', None)
        session.pop('current_db', None)
        session.pop('connection_params', None)
        session.pop('original_filename', None)
    return jsonify({'success': True})

@app.route('/cleanup', methods=['POST'])
def cleanup_session():
    """Manually cleanup session files when user closes the tab"""
    session_id = session.get('session_id')
    if session_id:
        cleanup_session_file(session_id)
        session.clear()
    return jsonify({'success': True})

@app.before_request
def check_session():
    """Check if session has expired and cleanup if needed"""
    pass

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)