spliter/split_data/database.py at master · triangular-opensource/spliter

539 lines (470 loc) · 19.2 KB
Database connection and data analysis module for split_data project.
import mysql.connector
from mysql.connector import Error
from typing import List, Dict, Optional
import pandas as pd
from datetime import datetime
import schedule
import time
import threading
# Import expense classifier (optional - only if available)
    from expense_classifier import classify_expense
    CLASSIFIER_AVAILABLE = True
except ImportError:
    CLASSIFIER_AVAILABLE = False
    print("Warning: expense_classifier not available. Tag classification disabled.")
class SplitDataDB:
    """Database connection and operations class for split_data."""
    def __init__(self, host=None, port=None, database=None,
                 user=None, password=None):
        """
        Initialize database connection parameters.
        Uses environment variables if available (for Docker), otherwise defaults.
        Args:
            host: MySQL host (default: from DB_HOST env or localhost)
            port: MySQL port (default: from DB_PORT env or 3307)
            database: Database name (default: from DB_NAME env or split_data)
            user: MySQL username (default: from DB_USER env or split_user)
            password: MySQL password (default: from DB_PASSWORD env or split_password)
        """
        self.host = host or os.getenv('DB_HOST', 'localhost')
        # In Docker, use port 3306 (internal), otherwise 3307 (external)
        default_port = 3306 if os.getenv('DB_HOST') == 'mysql' else 3307
        self.port = port or int(os.getenv('DB_PORT', str(default_port)))
        self.database = database or os.getenv('DB_NAME', 'split_data')
        self.user = user or os.getenv('DB_USER', 'split_user')
        self.password = password or os.getenv('DB_PASSWORD', 'split_password')
        self.connection = None
    def connect(self):
        """Establish connection to MySQL database."""
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                port=self.port,
                database=self.database,
                user=self.user,
                password=self.password
            if self.connection.is_connected():
                print(f"Successfully connected to MySQL database '{self.database}'")
                return True
        except Error as e:
            print(f"Error connecting to MySQL: {e}")
            return False
    def disconnect(self):
        """Close database connection."""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            print("MySQL connection closed")
    def execute_query(self, query: str, params: Optional[tuple] = None) -> List[Dict]:
        """
        Execute a SELECT query and return results as list of dictionaries.
        Args:
            query: SQL query string
            params: Optional tuple of parameters for parameterized queries
        Returns:
            List of dictionaries representing rows
        """
        if not self.connection or not self.connection.is_connected():
            self.connect()
        cursor = self.connection.cursor(dictionary=True)
        try:
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)
            results = cursor.fetchall()
            return results
        except Error as e:
            print(f"Error executing query: {e}")
            return []
        finally:
            cursor.close()
    def execute_update(self, query: str, params: Optional[tuple] = None) -> bool:
        """
        Execute an INSERT, UPDATE, or DELETE query.
        Args:
            query: SQL query string
            params: Optional tuple of parameters for parameterized queries
        Returns:
            True if successful, False otherwise
        """
        if not self.connection or not self.connection.is_connected():
            self.connect()
        cursor = self.connection.cursor()
        try:
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)
            self.connection.commit()
            return True
        except Error as e:
            print(f"Error executing update: {e}")
            self.connection.rollback()
            return False
        finally:
            cursor.close()
    def create_group(self, group_name: str, description: str, user_id: int) -> Optional[int]:
        """
        Create a new group.
        Args:
            group_name: Name of the group
            description: Description of the group
            user_id: User ID who created the group
        Returns:
            Group ID if successful, None otherwise
        """
        query = """
            INSERT INTO `group` (group_name, description, user_id, created_dt, update_dt)
            VALUES (%s, %s, %s, NOW(), NOW())
        """
        cursor = self.connection.cursor()
        try:
            cursor.execute(query, (group_name, description, user_id))
            self.connection.commit()
            group_id = cursor.lastrowid
            return group_id
        except Error as e:
            print(f"Error creating group: {e}")
            self.connection.rollback()
            return None
        finally:
            cursor.close()
    def add_member_to_group(self, group_id: int, member_id: int) -> bool:
        """
        Add a member to a group.
        Args:
            group_id: Group ID
            member_id: User ID to add as member
        Returns:
            True if successful, False otherwise
        """
        query = """
            INSERT INTO member (group_id, member_id)
            VALUES (%s, %s)
            ON DUPLICATE KEY UPDATE group_id=group_id
        """
        return self.execute_update(query, (group_id, member_id))
    def create_split(self, expense_id: int, group_id: int, user_id: int, 
                    amount: float, paid_or_not: Optional[bool] = None) -> Optional[int]:
        """
        Create a split entry for an expense.
        Args:
            expense_id: Expense ID
            group_id: Group ID
            user_id: User ID involved in this split
            amount: Amount (positive = paid, negative = owe)
            paid_or_not: Payment status (None for payer, True/False for debtors)
        Returns:
            Split ID if successful, None otherwise
        """
        query = """
            INSERT INTO split (expense_id, group_id, user_id, amount, paid_or_not, created_dt, updated_dt)
            VALUES (%s, %s, %s, %s, %s, NOW(), NOW())
        """
        cursor = self.connection.cursor()
        try:
            cursor.execute(query, (expense_id, group_id, user_id, amount, paid_or_not))
            self.connection.commit()
            split_id = cursor.lastrowid
            return split_id
        except Error as e:
            print(f"Error creating split: {e}")
            self.connection.rollback()
            return None
        finally:
            cursor.close()
    def get_all_groups(self) -> pd.DataFrame:
        """Get all groups as a pandas DataFrame."""
        query = "SELECT * FROM `group` ORDER BY created_dt DESC"
        results = self.execute_query(query)
        return pd.DataFrame(results)
    def get_group_members(self, group_id: int) -> pd.DataFrame:
        """Get all members of a specific group."""
        query = """
            SELECT m.id, m.group_id, m.member_id, g.group_name
            FROM member m
            JOIN `group` g ON m.group_id = g.id
            WHERE m.group_id = %s
        """
        results = self.execute_query(query, (group_id,))
        return pd.DataFrame(results)
    def get_expenses_by_group(self, group_id: int) -> pd.DataFrame:
        """Get all expenses for a specific group."""
        query = """
            SELECT e.*, g.group_name
            FROM expense e
            JOIN `group` g ON e.group_id = g.id
            WHERE e.group_id = %s
            ORDER BY e.created_dt DESC
        """
        results = self.execute_query(query, (group_id,))
        return pd.DataFrame(results)
    def get_splits_by_expense(self, expense_id: int) -> pd.DataFrame:
        """Get all splits for a specific expense."""
        query = """
            SELECT s.*, e.title as expense_title, e.tag
            FROM split s
            JOIN expense e ON s.expense_id = e.id
            WHERE s.expense_id = %s
        """
        results = self.execute_query(query, (expense_id,))
        return pd.DataFrame(results)
    def get_user_balance(self, user_id: int) -> Dict:
        """
        Calculate user's balance across all groups.
        Returns total paid, total owed, and net balance.
        """
        query = """
            SELECT 
                SUM(CASE WHEN amount > 0 THEN amount ELSE 0 END) as total_paid,
                SUM(CASE WHEN amount < 0 THEN ABS(amount) ELSE 0 END) as total_owed,
                SUM(amount) as net_balance
            FROM split
            WHERE user_id = %s
        """
        results = self.execute_query(query, (user_id,))
        return results[0] if results else {'total_paid': 0, 'total_owed': 0, 'net_balance': 0}
    def get_unsettled_expenses(self) -> pd.DataFrame:
        """Get all expenses that are not yet settled."""
        query = """
            SELECT e.*, g.group_name,
                   COUNT(CASE WHEN s.paid_or_not = FALSE THEN 1 END) as unpaid_count
            FROM expense e
            JOIN `group` g ON e.group_id = g.id
            LEFT JOIN split s ON e.id = s.expense_id
            WHERE e.is_settled = FALSE
            GROUP BY e.id
            ORDER BY e.created_dt DESC
        """
        results = self.execute_query(query)
        return pd.DataFrame(results)
    def get_expense_summary(self) -> pd.DataFrame:
        """Get summary of all expenses with split details."""
        query = """
            SELECT 
                e.id as expense_id,
                e.title,
                e.tag,
                e.is_settled,
                g.group_name,
                e.user_id as paid_by_user,
                COUNT(s.id) as split_count,
                SUM(CASE WHEN s.amount > 0 THEN s.amount ELSE 0 END) as total_paid,
                SUM(CASE WHEN s.amount < 0 THEN ABS(s.amount) ELSE 0 END) as total_owed,
                COUNT(CASE WHEN s.paid_or_not = FALSE THEN 1 END) as unpaid_count
            FROM expense e
            JOIN `group` g ON e.group_id = g.id
            LEFT JOIN split s ON e.id = s.expense_id
            GROUP BY e.id
            ORDER BY e.created_dt DESC
        """
        results = self.execute_query(query)
        return pd.DataFrame(results)
    def get_group_statistics(self, group_id: int) -> Dict:
        """Get statistics for a specific group."""
        query = """
            SELECT 
                COUNT(DISTINCT e.id) as total_expenses,
                COUNT(DISTINCT m.member_id) as total_members,
                SUM(CASE WHEN s.amount > 0 THEN s.amount ELSE 0 END) as total_paid,
                SUM(CASE WHEN s.amount < 0 THEN ABS(s.amount) ELSE 0 END) as total_owed,
                COUNT(CASE WHEN e.is_settled = FALSE THEN 1 END) as unsettled_expenses
            FROM `group` g
            LEFT JOIN expense e ON g.id = e.group_id
            LEFT JOIN member m ON g.id = m.group_id
            LEFT JOIN split s ON e.id = s.expense_id
            WHERE g.id = %s
        """
        results = self.execute_query(query, (group_id,))
        return results[0] if results else {}
    def create_expense_with_classification(self, title: str, group_id: int, 
                                          user_id: int, description: str = "",
                                          tag: Optional[str] = None) -> Optional[int]:
        """
        Create a new expense with automatic tag classification.
        Args:
            title: Expense title (e.g., "McDonald", "Uber ride")
            group_id: Group ID
            user_id: User ID who paid
            description: Optional description
            tag: Optional tag (if provided, classification is skipped)
        Returns:
            Expense ID if successful, None otherwise
        """
        # Auto-classify tag if not provided
        if tag is None and CLASSIFIER_AVAILABLE:
            try:
                tag = classify_expense(title)
            except Exception as e:
                print(f"Warning: Classification failed: {e}. Using 'other' as default.")
                tag = "other"
        elif tag is None:
            tag = "other"
        query = """
            INSERT INTO expense (title, tag, description, is_settled, group_id, user_id, created_dt, updated_dt)
            VALUES (%s, %s, %s, %s, %s, %s, NOW(), NOW())
        """
        cursor = self.connection.cursor()
        try:
            cursor.execute(query, (title, tag, description, False, group_id, user_id))
            self.connection.commit()
            expense_id = cursor.lastrowid
            return expense_id
        except Error as e:
            print(f"Error creating expense: {e}")
            self.connection.rollback()
            return None
        finally:
            cursor.close()
    def analyze_expense_by_tag(self) -> pd.DataFrame:
        """Analyze expenses grouped by tag."""
        query = """
            SELECT 
                e.tag,
                COUNT(DISTINCT e.id) as expense_count,
                SUM(CASE WHEN s.amount > 0 THEN s.amount ELSE 0 END) as total_amount,
                AVG(CASE WHEN s.amount > 0 THEN s.amount ELSE 0 END) as avg_amount,
                COUNT(CASE WHEN e.is_settled = FALSE THEN 1 END) as unsettled_count
            FROM expense e
            LEFT JOIN split s ON e.id = s.expense_id
            WHERE e.tag IS NOT NULL
            GROUP BY e.tag
            ORDER BY total_amount DESC
        """
        results = self.execute_query(query)
        return pd.DataFrame(results)
    def update_settled_expenses(self) -> int:
        """
        Check split table for every expense_id.
        If all splits have paid_or_not as NULL or TRUE (no FALSE values),
        update is_settled from FALSE to TRUE for that expense_id.
        Returns:
            Number of expenses updated
        """
        # Find all expenses that should be marked as settled
        # An expense is settled if all its splits have paid_or_not as NULL or TRUE
        query = """
            UPDATE expense e
            SET e.is_settled = TRUE
            WHERE e.is_settled = FALSE
            AND NOT EXISTS (
                SELECT 1
                FROM split s
                WHERE s.expense_id = e.id
                AND s.paid_or_not = FALSE
        """
        cursor = self.connection.cursor()
        try:
            cursor.execute(query)
            rows_updated = cursor.rowcount
            self.connection.commit()
            if rows_updated > 0:
                print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
                      f"Updated {rows_updated} expense(s) to settled status")
            return rows_updated
        except Error as e:
            print(f"Error updating settled expenses: {e}")
            self.connection.rollback()
            return 0
        finally:
            cursor.close()
    def start_settlement_checker(self, interval_minutes: int = 10):
        """
        Start a background scheduler that checks and updates settlement status
        every specified number of minutes.
        Args:
            interval_minutes: Interval in minutes between checks (default: 10)
        """
        if not self.connection or not self.connection.is_connected():
            self.connect()
        def run_check():
            """Wrapper function to run the settlement check."""
            try:
                if not self.connection or not self.connection.is_connected():
                    self.connect()
                self.update_settled_expenses()
            except Exception as e:
                print(f"Error in settlement checker: {e}")
        # Schedule the job
        schedule.every(interval_minutes).minutes.do(run_check)
        print(f"Settlement checker started. Will check every {interval_minutes} minutes.")
        print("Press Ctrl+C to stop.")
        # Run the scheduler in a separate thread
        def run_scheduler():
            while True:
                schedule.run_pending()
                time.sleep(60)  # Check every minute for pending jobs
        scheduler_thread = threading.Thread(target=run_scheduler, daemon=True)
        scheduler_thread.start()
        return scheduler_thread
def main():
    """Example usage and data analysis."""
    db = SplitDataDB()
        # Connect to database
        if not db.connect():
            print("Failed to connect to database")
            return
        print("\n" + "="*60)
        print("SPLIT DATA ANALYSIS")
        print("="*60)
        # 1. Get all groups
        print("\n1. All Groups:")
        print("-" * 60)
        groups_df = db.get_all_groups()
        print(groups_df.to_string(index=False))
        # 2. Get expense summary
        print("\n2. Expense Summary:")
        print("-" * 60)
        expense_summary = db.get_expense_summary()
        print(expense_summary.to_string(index=False))
        # 3. Get unsettled expenses
        print("\n3. Unsettled Expenses:")
        print("-" * 60)
        unsettled = db.get_unsettled_expenses()
        if not unsettled.empty:
            print(unsettled.to_string(index=False))
        else:
            print("No unsettled expenses")
        # 4. Analyze expenses by tag
        print("\n4. Expenses by Tag:")
        print("-" * 60)
        tag_analysis = db.analyze_expense_by_tag()
        print(tag_analysis.to_string(index=False))
        # 5. User balances
        print("\n5. User Balances:")
        print("-" * 60)
        for user_id in [1, 2, 3, 4]:
            balance = db.get_user_balance(user_id)
            print(f"User {user_id}: Paid=${balance['total_paid']:.2f}, "
                  f"Owed=${balance['total_owed']:.2f}, "
                  f"Net=${balance['net_balance']:.2f}")
        # 6. Group statistics
        print("\n6. Group Statistics:")
        print("-" * 60)
        for group_id in [1, 2]:
            stats = db.get_group_statistics(group_id)
            print(f"\nGroup {group_id}:")
            print(f"  Total Expenses: {stats.get('total_expenses', 0)}")
            print(f"  Total Members: {stats.get('total_members', 0)}")
            print(f"  Total Paid: ${stats.get('total_paid', 0):.2f}")
            print(f"  Total Owed: ${stats.get('total_owed', 0):.2f}")
            print(f"  Unsettled Expenses: {stats.get('unsettled_expenses', 0)}")
        # 7. Detailed splits for expense 1
        print("\n7. Detailed Splits for Expense 1 (McDonald):")
        print("-" * 60)
        splits_df = db.get_splits_by_expense(1)
        print(splits_df.to_string(index=False))
    except Exception as e:
        print(f"Error: {e}")
    finally:
        db.disconnect()
if __name__ == "__main__":
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

database.py

Latest commit

History

database.py

File metadata and controls
