data_layer.py

import duckdb
import pandas as pd
import os
from pathlib import Path
from typing import Dict, List, Any, Optional
import logging
from test_logger import test_logger

logger = logging.getLogger(__name__)

class OlistDataLayer:
    """Data layer for managing Olist CSV files with DuckDB"""
    
    def __init__(self, data_dir: str = "data"):
        self.data_dir = Path(data_dir)
        self.conn = duckdb.connect(":memory:")
        self.tables = {}
        self.schema_info = {}
        self._setup_tables()
    
    def _setup_tables(self):
        """Load all CSV files into DuckDB tables"""
        csv_files = {
            'customers': 'olist_customers_dataset.csv',
            'geolocation': 'olist_geolocation_dataset.csv', 
            'orders': 'olist_orders_dataset.csv',
            'order_items': 'olist_order_items_dataset.csv',
            'payments': 'olist_order_payments_dataset.csv',
            'reviews': 'olist_order_reviews_dataset.csv',
            'products': 'olist_products_dataset.csv',
            'sellers': 'olist_sellers_dataset.csv',
            'category_translation': 'product_category_name_translation.csv'
        }
        
        for table_name, csv_file in csv_files.items():
            csv_path = self.data_dir / csv_file
            if csv_path.exists():
                try:
                    # Load CSV into DuckDB
                    self.conn.execute(f"""
                        CREATE TABLE {table_name} AS 
                        SELECT * FROM read_csv_auto('{csv_path}')
                    """)
                    
                    # Store table info
                    self.tables[table_name] = str(csv_path)
                    
                    # Get schema information
                    schema_result = self.conn.execute(f"DESCRIBE {table_name}").fetchall()
                    self.schema_info[table_name] = {
                        'columns': [row[0] for row in schema_result],
                        'types': [row[1] for row in schema_result],
                        'path': str(csv_path)
                    }
                    
                    logger.info(f"Loaded table '{table_name}' from {csv_file}")
                except Exception as e:
                    logger.error(f"Failed to load {csv_file}: {e}")
            else:
                logger.warning(f"CSV file not found: {csv_path}")
    
    def get_schema_info(self) -> Dict[str, Any]:
        """Get complete schema information for all tables"""
        return {
            'tables': self.schema_info,
            'relationships': self._get_relationships()
        }
    
    def _get_relationships(self) -> Dict[str, List[str]]:
        """Define table relationships based on foreign keys"""
        return {
            'orders': {
                'primary_key': 'order_id',
                'foreign_keys': ['customer_id'],
                'referenced_by': ['order_items', 'payments', 'reviews']
            },
            'customers': {
                'primary_key': 'customer_id', 
                'foreign_keys': ['customer_zip_code_prefix'],
                'referenced_by': ['orders']
            },
            'order_items': {
                'primary_key': ['order_id', 'order_item_id'],
                'foreign_keys': ['order_id', 'product_id', 'seller_id'],
                'referenced_by': []
            },
            'products': {
                'primary_key': 'product_id',
                'foreign_keys': ['product_category_name'],
                'referenced_by': ['order_items']
            },
            'sellers': {
                'primary_key': 'seller_id',
                'foreign_keys': ['seller_zip_code_prefix'], 
                'referenced_by': ['order_items']
            },
            'payments': {
                'primary_key': ['order_id', 'payment_sequential'],
                'foreign_keys': ['order_id'],
                'referenced_by': []
            },
            'reviews': {
                'primary_key': 'review_id',
                'foreign_keys': ['order_id'],
                'referenced_by': []
            },
            'geolocation': {
                'primary_key': 'geolocation_zip_code_prefix',
                'foreign_keys': [],
                'referenced_by': ['customers', 'sellers']
            },
            'category_translation': {
                'primary_key': 'product_category_name',
                'foreign_keys': [],
                'referenced_by': ['products']
            }
        }
    
    def execute_query(self, query: str) -> Dict[str, Any]:
        """Execute SQL query and return results"""
        try:
            # Log query execution start
            test_logger.log_data_layer_operation("query_execution_start", {"query": query})
            
            result = self.conn.execute(query).fetchall()
            columns = [desc[0] for desc in self.conn.description]
            
            query_result = {
                'success': True,
                'data': result,
                'columns': columns,
                'row_count': len(result),
                'query': query,
                'limited': 'LIMIT' in query.upper(),
                'preview': result[:5] if result else []
            }
            
            # Log successful query execution
            test_logger.log_data_layer_operation("query_execution_success", {
                "query": query,
                "row_count": len(result),
                "columns": columns
            })
            
            return query_result
            
        except Exception as e:
            logger.error(f"Query execution failed: {e}")
            error_result = {
                'success': False,
                'error': str(e),
                'query': query,
                'data': [],
                'columns': [],
                'row_count': 0
            }
            
            # Log query execution error
            test_logger.log_error("query_execution", str(e), {"query": query})
            
            return error_result
    
    def get_sample_data(self, table_name: str, limit: int = 5) -> Dict[str, Any]:
        """Get sample data from a table"""
        if table_name not in self.tables:
            return {'success': False, 'error': f'Table {table_name} not found'}
        
        return self.execute_query(f"SELECT * FROM {table_name} LIMIT {limit}")
    
    def get_table_stats(self, table_name: str) -> Dict[str, Any]:
        """Get basic statistics for a table"""
        if table_name not in self.tables:
            return {'success': False, 'error': f'Table {table_name} not found'}
        
        try:
            count_result = self.conn.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()
            return {
                'success': True,
                'table_name': table_name,
                'row_count': count_result[0],
                'columns': self.schema_info[table_name]['columns'],
                'column_types': self.schema_info[table_name]['types']
            }
        except Exception as e:
            return {'success': False, 'error': str(e)}
    
    def validate_query(self, query: str) -> Dict[str, Any]:
        """Validate SQL query without executing it"""
        try:
            # Use EXPLAIN to validate query syntax
            self.conn.execute(f"EXPLAIN {query}")
            return {'valid': True, 'error': None}
        except Exception as e:
            return {'valid': False, 'error': str(e)}
    
    def close(self):
        """Close database connection"""
        if self.conn:
            self.conn.close()