database_handler.py

"""
Database Handler - Abstraction layer for multiple database types
Supports: SQLite, MySQL, PostgreSQL, MongoDB, SQL Server
"""

import sqlite3
import json
from typing import Dict, List, Any, Tuple

class DatabaseHandler:
    """Base class for database handlers"""
    
    def __init__(self, connection_params: Dict[str, Any]):
        self.connection_params = connection_params
        self.db_type = connection_params.get('db_type', 'sqlite')
        self.connection = None
        
    def connect(self):
        """Establish connection to database"""
        raise NotImplementedError
        
    def disconnect(self):
        """Close database connection"""
        if self.connection:
            self.connection.close()
            
    def get_tables(self) -> List[str]:
        """Get list of all tables"""
        raise NotImplementedError
        
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        """Get table structure (columns, types)"""
        raise NotImplementedError
        
    def get_table_data(self, table_name: str, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        """Get paginated table data"""
        raise NotImplementedError
        
    def get_all_table_data(self, table_name: str) -> List[Dict[str, Any]]:
        """Get all data from a table (for export)"""
        raise NotImplementedError


class SQLiteHandler(DatabaseHandler):
    """Handler for SQLite databases"""
    
    def connect(self):
        filepath = self.connection_params.get('filepath')
        self.connection = sqlite3.connect(filepath)
        self.connection.row_factory = sqlite3.Row
        
    def get_tables(self) -> List[str]:
        cursor = self.connection.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;")
        return [row[0] for row in cursor.fetchall()]
        
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        cursor = self.connection.cursor()
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [row[1] for row in cursor.fetchall()]
        return {'columns': columns}
        
    def get_table_data(self, table_name: str, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        cursor = self.connection.cursor()
        offset = (page - 1) * per_page
        
        # Get columns
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [row[1] for row in cursor.fetchall()]
        
        # Get total count
        cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
        total_count = cursor.fetchone()[0]
        
        # Get data
        cursor.execute(f"SELECT * FROM {table_name} LIMIT ? OFFSET ?", (per_page, offset))
        rows = cursor.fetchall()
        
        data = [dict(row) for row in rows]
        
        return {
            'columns': columns,
            'data': data,
            'total_count': total_count,
            'page': page,
            'per_page': per_page,
            'total_pages': (total_count + per_page - 1) // per_page if total_count > 0 else 1
        }
        
    def get_all_table_data(self, table_name: str) -> List[Dict[str, Any]]:
        cursor = self.connection.cursor()
        cursor.execute(f"SELECT * FROM {table_name}")
        rows = cursor.fetchall()
        return [dict(row) for row in rows]


class MySQLHandler(DatabaseHandler):
    """Handler for MySQL databases"""
    
    def connect(self):
        import pymysql
        self.connection = pymysql.connect(
            host=self.connection_params.get('host', 'localhost'),
            port=self.connection_params.get('port', 3306),
            user=self.connection_params.get('user'),
            password=self.connection_params.get('password'),
            database=self.connection_params.get('database'),
            cursorclass=pymysql.cursors.DictCursor
        )
        
    def get_tables(self) -> List[str]:
        with self.connection.cursor() as cursor:
            cursor.execute("SHOW TABLES")
            results = cursor.fetchall()
            # MySQL returns tables in a dict with key like 'Tables_in_dbname'
            return [list(row.values())[0] for row in results]
            
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        with self.connection.cursor() as cursor:
            cursor.execute(f"DESCRIBE `{table_name}`")
            results = cursor.fetchall()
            columns = [row['Field'] for row in results]
            return {'columns': columns}
            
    def get_table_data(self, table_name: str, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        offset = (page - 1) * per_page
        
        with self.connection.cursor() as cursor:
            # Get columns
            cursor.execute(f"DESCRIBE `{table_name}`")
            results = cursor.fetchall()
            columns = [row['Field'] for row in results]
            
            # Get total count
            cursor.execute(f"SELECT COUNT(*) as count FROM `{table_name}`")
            total_count = cursor.fetchone()['count']
            
            # Get data
            cursor.execute(f"SELECT * FROM `{table_name}` LIMIT %s OFFSET %s", (per_page, offset))
            data = cursor.fetchall()
            
            return {
                'columns': columns,
                'data': data,
                'total_count': total_count,
                'page': page,
                'per_page': per_page,
                'total_pages': (total_count + per_page - 1) // per_page if total_count > 0 else 1
            }
            
    def get_all_table_data(self, table_name: str) -> List[Dict[str, Any]]:
        with self.connection.cursor() as cursor:
            cursor.execute(f"SELECT * FROM `{table_name}`")
            return cursor.fetchall()


class PostgreSQLHandler(DatabaseHandler):
    """Handler for PostgreSQL databases"""
    
    def connect(self):
        import psycopg2
        import psycopg2.extras
        self.connection = psycopg2.connect(
            host=self.connection_params.get('host', 'localhost'),
            port=self.connection_params.get('port', 5432),
            user=self.connection_params.get('user'),
            password=self.connection_params.get('password'),
            database=self.connection_params.get('database')
        )
        
    def get_tables(self) -> List[str]:
        import psycopg2.extras
        with self.connection.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            cursor.execute("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public' 
                ORDER BY table_name
            """)
            return [row['table_name'] for row in cursor.fetchall()]
            
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        import psycopg2.extras
        with self.connection.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            cursor.execute(f"""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = %s
                ORDER BY ordinal_position
            """, (table_name,))
            columns = [row['column_name'] for row in cursor.fetchall()]
            return {'columns': columns}
            
    def get_table_data(self, table_name: str, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        import psycopg2.extras
        offset = (page - 1) * per_page
        
        with self.connection.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            # Get columns
            cursor.execute(f"""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = %s
                ORDER BY ordinal_position
            """, (table_name,))
            columns = [row['column_name'] for row in cursor.fetchall()]
            
            # Get total count
            cursor.execute(f'SELECT COUNT(*) as count FROM "{table_name}"')
            total_count = cursor.fetchone()['count']
            
            # Get data
            cursor.execute(f'SELECT * FROM "{table_name}" LIMIT %s OFFSET %s', (per_page, offset))
            rows = cursor.fetchall()
            data = [dict(row) for row in rows]
            
            return {
                'columns': columns,
                'data': data,
                'total_count': total_count,
                'page': page,
                'per_page': per_page,
                'total_pages': (total_count + per_page - 1) // per_page if total_count > 0 else 1
            }
            
    def get_all_table_data(self, table_name: str) -> List[Dict[str, Any]]:
        import psycopg2.extras
        with self.connection.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            cursor.execute(f'SELECT * FROM "{table_name}"')
            rows = cursor.fetchall()
            return [dict(row) for row in rows]


class MongoDBHandler(DatabaseHandler):
    """Handler for MongoDB databases"""
    
    def connect(self):
        from pymongo import MongoClient
        
        # Build connection string
        if 'connection_string' in self.connection_params:
            connection_string = self.connection_params['connection_string']
        else:
            host = self.connection_params.get('host', 'localhost')
            port = self.connection_params.get('port', 27017)
            user = self.connection_params.get('user', '')
            password = self.connection_params.get('password', '')
            
            if user and password:
                connection_string = f"mongodb://{user}:{password}@{host}:{port}/"
            else:
                connection_string = f"mongodb://{host}:{port}/"
        
        self.client = MongoClient(connection_string)
        database_name = self.connection_params.get('database')
        self.connection = self.client[database_name]
        
    def disconnect(self):
        if hasattr(self, 'client'):
            self.client.close()
            
    def get_tables(self) -> List[str]:
        """In MongoDB, tables are called collections"""
        return sorted(self.connection.list_collection_names())
        
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        """Get columns (fields) from first document as sample"""
        collection = self.connection[table_name]
        sample_doc = collection.find_one()
        
        if sample_doc:
            columns = list(sample_doc.keys())
        else:
            columns = []
            
        return {'columns': columns}
        
    def get_table_data(self, table_name: str, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        collection = self.connection[table_name]
        skip = (page - 1) * per_page
        
        # Get total count
        total_count = collection.count_documents({})
        
        # Get data
        cursor = collection.find().skip(skip).limit(per_page)
        data = []
        columns = set()
        
        for doc in cursor:
            # Convert ObjectId to string
            if '_id' in doc:
                doc['_id'] = str(doc['_id'])
            # Collect all unique fields
            columns.update(doc.keys())
            data.append(doc)
        
        columns = sorted(list(columns))
        
        return {
            'columns': columns,
            'data': data,
            'total_count': total_count,
            'page': page,
            'per_page': per_page,
            'total_pages': (total_count + per_page - 1) // per_page if total_count > 0 else 1
        }
        
    def get_all_table_data(self, table_name: str) -> List[Dict[str, Any]]:
        collection = self.connection[table_name]
        data = []
        
        for doc in collection.find():
            # Convert ObjectId to string
            if '_id' in doc:
                doc['_id'] = str(doc['_id'])
            data.append(doc)
            
        return data


class SQLServerHandler(DatabaseHandler):
    """Handler for SQL Server databases"""
    
    def connect(self):
        import pyodbc
        
        connection_string = (
            f"DRIVER={{ODBC Driver 17 for SQL Server}};"
            f"SERVER={self.connection_params.get('host', 'localhost')};"
            f"DATABASE={self.connection_params.get('database')};"
            f"UID={self.connection_params.get('user')};"
            f"PWD={self.connection_params.get('password')}"
        )
        
        self.connection = pyodbc.connect(connection_string)
        
    def get_tables(self) -> List[str]:
        cursor = self.connection.cursor()
        cursor.execute("""
            SELECT TABLE_NAME 
            FROM INFORMATION_SCHEMA.TABLES 
            WHERE TABLE_TYPE = 'BASE TABLE'
            ORDER BY TABLE_NAME
        """)
        return [row[0] for row in cursor.fetchall()]
        
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        cursor = self.connection.cursor()
        cursor.execute(f"""
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_NAME = ?
            ORDER BY ORDINAL_POSITION
        """, (table_name,))
        columns = [row[0] for row in cursor.fetchall()]
        return {'columns': columns}
        
    def get_table_data(self, table_name: str, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        offset = (page - 1) * per_page
        cursor = self.connection.cursor()
        
        # Get columns
        cursor.execute(f"""
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_NAME = ?
            ORDER BY ORDINAL_POSITION
        """, (table_name,))
        columns = [row[0] for row in cursor.fetchall()]
        
        # Get total count
        cursor.execute(f"SELECT COUNT(*) FROM [{table_name}]")
        total_count = cursor.fetchone()[0]
        
        # Get data (using OFFSET/FETCH for pagination)
        cursor.execute(f"""
            SELECT * FROM [{table_name}]
            ORDER BY (SELECT NULL)
            OFFSET ? ROWS
            FETCH NEXT ? ROWS ONLY
        """, (offset, per_page))
        
        rows = cursor.fetchall()
        data = []
        for row in rows:
            row_dict = {}
            for i, col in enumerate(columns):
                row_dict[col] = row[i]
            data.append(row_dict)
        
        return {
            'columns': columns,
            'data': data,
            'total_count': total_count,
            'page': page,
            'per_page': per_page,
            'total_pages': (total_count + per_page - 1) // per_page if total_count > 0 else 1
        }
        
    def get_all_table_data(self, table_name: str) -> List[Dict[str, Any]]:
        cursor = self.connection.cursor()
        
        # Get columns
        cursor.execute(f"""
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_NAME = ?
            ORDER BY ORDINAL_POSITION
        """, (table_name,))
        columns = [row[0] for row in cursor.fetchall()]
        
        # Get all data
        cursor.execute(f"SELECT * FROM [{table_name}]")
        rows = cursor.fetchall()
        
        data = []
        for row in rows:
            row_dict = {}
            for i, col in enumerate(columns):
                row_dict[col] = row[i]
            data.append(row_dict)
            
        return data


def get_database_handler(connection_params: Dict[str, Any]) -> DatabaseHandler:
    """Factory function to get appropriate database handler"""
    db_type = connection_params.get('db_type', 'sqlite').lower()
    
    handlers = {
        'sqlite': SQLiteHandler,
        'mysql': MySQLHandler,
        'postgresql': PostgreSQLHandler,
        'postgres': PostgreSQLHandler,
        'mongodb': MongoDBHandler,
        'mongo': MongoDBHandler,
        'sqlserver': SQLServerHandler,
        'mssql': SQLServerHandler
    }
    
    handler_class = handlers.get(db_type)
    if not handler_class:
        raise ValueError(f"Unsupported database type: {db_type}")
        
    return handler_class(connection_params)