database.py

import os
import duckdb
import sqlite3
from datetime import datetime
import logging
import pprint

# Get logger (configuration is handled in config.py)
logger = logging.getLogger(__name__)

class Database:
    """
    Base database class.
    """
    def __init__(self, config):
        self.config = config
        self.conn = None
        
        # Get debug_writes flag from config, default to False if not present
        # Check both direct access and nested access to handle different config structures
        if isinstance(config, dict) and "database" in config and isinstance(config["database"], dict):
            self.debug_writes = config["database"].get("debug_writes", False)
        else:
            self.debug_writes = config.get("debug_writes", False)
        
        # Log the debug_writes setting
        logger.info(f"Database debug_writes flag initialized to: {self.debug_writes}")
    
    def connect(self):
        """
        Connect to the database.
        """
        raise NotImplementedError("Subclasses must implement connect()")
    
    def create_table(self, table_name, schema, recreate=False):
        """
        Create a table if it doesn't exist.
        
        Args:
            table_name: The name of the table to create
            schema: The schema definition for the table
            recreate: If True, drop the table if it exists before creating it
        """
        raise NotImplementedError("Subclasses must implement create_table()")
    
    def upsert(self, table_name, data, primary_key=None):
        """
        Upsert data into a table.
        """
        raise NotImplementedError("Subclasses must implement upsert()")
    
    def insert(self, table_name, data):
        """
        Insert data into a table.
        """
        raise NotImplementedError("Subclasses must implement insert()")
    
    def query(self, sql, params=None):
        """
        Execute a query and return the results.
        """
        raise NotImplementedError("Subclasses must implement query()")
    
    def close(self):
        """
        Close the database connection.
        """
        if self.conn:
            self.conn.close()
            self.conn = None
    
    def _debug_print_data(self, operation, table_name, data, primary_key=None):
        """
        Print debug information about data being written to the database.
        
        Args:
            operation: The operation being performed (e.g., "INSERT", "UPSERT")
            table_name: The name of the table
            data: The data being written
            primary_key: The primary key column (for upserts)
        """
        if not self.debug_writes:
            return
        
        logger.info(f"DEBUG WRITES: {operation} to table '{table_name}'")
        if primary_key:
            logger.info(f"Primary key: {primary_key}")
        
        logger.info(f"Number of rows: {len(data)}")
        
        # Print a sample of the data (up to 5 rows)
        sample_size = min(5, len(data))
        logger.info(f"Sample data (showing {sample_size} of {len(data)} rows):")
        
        pp = pprint.PrettyPrinter(indent=2)
        for i, item in enumerate(data[:sample_size]):
            logger.info(f"Row {i+1}:")
            logger.info(pp.pformat(item))
        
        if len(data) > sample_size:
            logger.info(f"... and {len(data) - sample_size} more rows")

class MotherDuckDatabase(Database):
    """
    MotherDuck database implementation.
    """
    def connect(self):
        """
        Connect to MotherDuck.
        """
        motherduck_token = os.getenv("MOTHERDUCK_TOKEN")
        if not motherduck_token:
            raise ValueError("MOTHERDUCK_TOKEN environment variable is not set")
        
        db_name = self.config["database"]["connection"]["motherduck"]["database"]
        connection_string = f"md:{db_name}?motherduck_token={motherduck_token}"
        
        logger.info(f"Connecting to MotherDuck database: {db_name}")
        self.conn = duckdb.connect(database=connection_string, read_only=False)
        return self.conn
    
    def create_table(self, table_name, schema, recreate=False):
        """
        Create a table if it doesn't exist.
        
        Args:
            table_name: The name of the table to create
            schema: The schema definition for the table
            recreate: If True, drop the table if it exists before creating it
        """
        if not self.conn:
            self.connect()
        
        # If recreate is True, drop the table if it exists
        if recreate:
            logger.info(f"Dropping table {table_name} if it exists (recreate=True)")
            self.conn.execute(f"""
            DROP TABLE IF EXISTS {table_name}
            """)
        
        # Create the table
        self.conn.execute(f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            {schema}
        )
        """)
        
        # Check if updated_at column exists, if not add it
        column_check = self.conn.execute(f"""
        SELECT column_name
        FROM information_schema.columns
        WHERE table_name = '{table_name}' AND column_name = 'updated_at'
        """).fetchall()
        
        if not column_check:
            logger.info(f"Adding updated_at column to table {table_name}")
            self.conn.execute(f"""
            ALTER TABLE {table_name} ADD COLUMN updated_at TIMESTAMP
            """)
    
    def upsert(self, table_name, data, primary_key=None):
        """
        Upsert data into a table using a true merge approach.
        This implementation updates existing records and inserts new ones without deleting any data.
        """
        if not self.conn:
            self.connect()
        
        if not data:
            return 0
        
        # Debug print the data being upserted
        self._debug_print_data("UPSERT", table_name, data, primary_key)
        
        # Create a temporary table for the new data
        columns = list(data[0].keys())
        if "updated_at" not in columns:
            columns.append("updated_at")
        
        column_defs = ", ".join([f"{col} {self._get_column_type(data[0][col]) if col in data[0] else 'TIMESTAMP'}" for col in columns])
        
        self.conn.execute(f"""
        CREATE TEMPORARY TABLE temp_data (
            {column_defs}
        )
        """)
        
        # Get table schema to check column types
        try:
            schema_query = f"""
            SELECT column_name, data_type
            FROM information_schema.columns
            WHERE table_name = '{table_name}'
            """
            schema_info = self.conn.execute(schema_query).fetchall()
            column_types = {col[0]: col[1] for col in schema_info}
            
            logger.info(f"DEBUG: Table schema for {table_name}: {column_types}")
        except Exception as e:
            logger.warning(f"DEBUG: Could not get schema info: {e}")
            column_types = {}
        
        # Insert data into the temporary table
        current_time = datetime.utcnow().isoformat()
        logger.info(f"DEBUG: Using UTC current_time for updated_at: {current_time}")
        
        for item in data:
            try:
                # Prepare values with updated_at and convert timestamps
                values = []
                for col in columns:
                    if col == "updated_at":
                        values.append(current_time)
                    elif col in column_types and "TIMESTAMP" in column_types.get(col, "").upper() and item.get(col):
                        # Convert string timestamps to datetime objects for TIMESTAMP columns
                        try:
                            timestamp_str = item.get(col)
                            logger.info(f"DEBUG: Converting timestamp: {timestamp_str}, type: {type(timestamp_str)}")
                            
                            # If it's already a datetime object, use it directly
                            if isinstance(timestamp_str, datetime):
                                values.append(timestamp_str)
                            # If it's a string in ISO format, parse it
                            elif isinstance(timestamp_str, str):
                                # Try to parse the timestamp string
                                try:
                                    # Try ISO format with Z (UTC)
                                    dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%SZ")
                                    values.append(dt)
                                    logger.info(f"DEBUG: Converted to datetime: {dt}")
                                except ValueError:
                                    try:
                                        # Try ISO format without Z
                                        dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%S")
                                        values.append(dt)
                                        logger.info(f"DEBUG: Converted to datetime: {dt}")
                                    except ValueError:
                                        # If all parsing fails, pass the string as is
                                        logger.warning(f"DEBUG: Could not parse timestamp: {timestamp_str}")
                                        values.append(timestamp_str)
                            else:
                                # If it's neither a datetime nor a string, use as is
                                values.append(timestamp_str)
                        except Exception as e:
                            logger.error(f"DEBUG: Error converting timestamp {item.get(col)}: {e}")
                            values.append(item.get(col))
                    else:
                        values.append(item.get(col))
                
                placeholders = ", ".join(["?" for _ in range(len(values))])
                insert_sql = f"""
                INSERT INTO temp_data ({', '.join(columns)})
                VALUES ({placeholders})
                """
                
                if self.debug_writes:
                    logger.info(f"DEBUG: Executing SQL: {insert_sql}")
                    logger.info(f"DEBUG: With values: {values}")
                
                self.conn.execute(insert_sql, values)
            except Exception as e:
                logger.error(f"DEBUG: Error inserting row into temp table: {e}")
                logger.error(f"DEBUG: Row data: {item}")
                raise
        
        # Check if the target table exists
        table_exists_query = f"""
        SELECT COUNT(*) FROM information_schema.tables
        WHERE table_name = '{table_name}'
        """
        table_exists = self.conn.execute(table_exists_query).fetchone()[0] > 0
        
        if not table_exists:
            # If table doesn't exist, just insert all records
            logger.info(f"Table {table_name} does not exist yet, creating and inserting all records")
            
            # Create the table with the same schema as temp_data
            create_table_sql = f"""
            CREATE TABLE {table_name} AS SELECT * FROM temp_data WHERE 1=0
            """
            self.conn.execute(create_table_sql)
            
            # Insert all records
            insert_all_sql = f"""
            INSERT INTO {table_name} ({', '.join(columns)})
            SELECT {', '.join(columns)} FROM temp_data
            """
            self.conn.execute(insert_all_sql)
        else:
            # If primary key is provided, use it for true upserting (merge)
            if primary_key:
                # Get list of primary keys in the new data
                pk_values = [item[primary_key] for item in data]
                pk_list = ", ".join([f"'{pk}'" for pk in pk_values])
                
                if self.debug_writes:
                    # Log the records being updated
                    select_sql = f"SELECT * FROM {table_name} WHERE {primary_key} IN ({pk_list})"
                    logger.info(f"DEBUG: Executing SQL: {select_sql}")
                    existing_records = self.query(select_sql)
                    
                    if existing_records:
                        logger.info(f"DEBUG WRITES: Found {len(existing_records)} existing records in '{table_name}' with {primary_key} IN ({pk_list})")
                        logger.info(f"DEBUG: Will update these records and insert new ones")
                        
                        # Log the first existing record
                        if len(existing_records) > 0:
                            logger.info(f"DEBUG: First existing record: {existing_records[0]}")
                            if hasattr(existing_records[0], 'keys'):
                                logger.info(f"DEBUG: First existing record updated_at: {existing_records[0].get('updated_at', 'N/A')}")
                
                # Create a temporary table with existing records that match the primary keys
                create_existing_sql = f"""
                CREATE TEMPORARY TABLE existing_records AS
                SELECT * FROM {table_name} WHERE {primary_key} IN ({pk_list})
                """
                
                if self.debug_writes:
                    logger.info(f"DEBUG: Executing SQL: {create_existing_sql}")
                
                self.conn.execute(create_existing_sql)
                
                # Update existing records
                non_pk_columns = [col for col in columns if col != primary_key]
                set_clauses = [f"{col} = temp_data.{col}" for col in non_pk_columns]
                
                update_sql = f"""
                UPDATE {table_name}
                SET {', '.join(set_clauses)}
                FROM temp_data
                WHERE {table_name}.{primary_key} = temp_data.{primary_key}
                """
                
                if self.debug_writes:
                    logger.info(f"DEBUG: Executing SQL: {update_sql}")
                
                self.conn.execute(update_sql)
                
                # Insert new records that don't exist in the target table
                insert_new_sql = f"""
                INSERT INTO {table_name} ({', '.join(columns)})
                SELECT {', '.join(columns)} FROM temp_data
                WHERE NOT EXISTS (
                    SELECT 1 FROM {table_name}
                    WHERE {table_name}.{primary_key} = temp_data.{primary_key}
                )
                """
                
                if self.debug_writes:
                    logger.info(f"DEBUG: Executing SQL: {insert_new_sql}")
                
                self.conn.execute(insert_new_sql)
                
                # Drop the temporary table with existing records
                self.conn.execute("DROP TABLE existing_records")
            else:
                # If no primary key is provided, just insert all records
                # This might create duplicates, but that's expected behavior without a primary key
                insert_all_sql = f"""
                INSERT INTO {table_name} ({', '.join(columns)})
                SELECT {', '.join(columns)} FROM temp_data
                """
                
                if self.debug_writes:
                    logger.info(f"DEBUG: Executing SQL: {insert_all_sql}")
                
                self.conn.execute(insert_all_sql)
        
        # Drop the temporary table
        self.conn.execute("DROP TABLE temp_data")
        
        return len(data)
    
    def insert(self, table_name, data):
        """
        Insert data into a table without upserting.
        """
        if not self.conn:
            self.connect()
        
        if not data:
            return 0
        
        # Debug print the data being inserted
        self._debug_print_data("INSERT", table_name, data)
        
        # Get column names from the first data item
        columns = list(data[0].keys())
        if "updated_at" not in columns:
            columns.append("updated_at")
        
        # Get table schema to check column types
        try:
            schema_query = f"""
            SELECT column_name, data_type
            FROM information_schema.columns
            WHERE table_name = '{table_name}'
            """
            schema_info = self.conn.execute(schema_query).fetchall()
            column_types = {col[0]: col[1] for col in schema_info}
            
            logger.info(f"DEBUG: Table schema for {table_name}: {column_types}")
        except Exception as e:
            logger.warning(f"DEBUG: Could not get schema info: {e}")
            column_types = {}
        
        # Insert records
        count = 0
        current_time = datetime.utcnow().isoformat()
        
        for item in data:
            try:
                # Prepare values with updated_at and convert timestamps
                values = []
                for col in columns:
                    if col == "updated_at":
                        values.append(current_time)
                    elif col in column_types and "TIMESTAMP" in column_types[col].upper() and item.get(col):
                        # Convert string timestamps to datetime objects for TIMESTAMP columns
                        try:
                            timestamp_str = item.get(col)
                            logger.info(f"DEBUG: Converting timestamp: {timestamp_str}, type: {type(timestamp_str)}")
                            
                            # If it's already a datetime object, use it directly
                            if isinstance(timestamp_str, datetime):
                                values.append(timestamp_str)
                            # If it's a string in ISO format, parse it
                            elif isinstance(timestamp_str, str):
                                # Try to parse the timestamp string
                                try:
                                    # Try ISO format with Z (UTC)
                                    dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%SZ")
                                    values.append(dt)
                                    logger.info(f"DEBUG: Converted to datetime: {dt}")
                                except ValueError:
                                    try:
                                        # Try ISO format without Z
                                        dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%S")
                                        values.append(dt)
                                        logger.info(f"DEBUG: Converted to datetime: {dt}")
                                    except ValueError:
                                        # If all parsing fails, pass the string as is
                                        logger.warning(f"DEBUG: Could not parse timestamp: {timestamp_str}")
                                        values.append(timestamp_str)
                            else:
                                # If it's neither a datetime nor a string, use as is
                                values.append(timestamp_str)
                        except Exception as e:
                            logger.error(f"DEBUG: Error converting timestamp {item.get(col)}: {e}")
                            values.append(item.get(col))
                    else:
                        values.append(item.get(col))
                
                placeholders = ", ".join(["?" for _ in range(len(values))])
                insert_sql = f"""
                INSERT INTO {table_name} ({', '.join(columns)})
                VALUES ({placeholders})
                """
                
                if self.debug_writes:
                    logger.info(f"DEBUG: Executing SQL: {insert_sql}")
                    logger.info(f"DEBUG: With values: {values}")
                
                self.conn.execute(insert_sql, values)
                count += 1
            except Exception as e:
                logger.error(f"DEBUG: Error inserting row {count+1}: {e}")
                logger.error(f"DEBUG: Row data: {item}")
                raise
        
        return count
    
    def query(self, sql, params=None):
        """
        Execute a query and return the results.
        """
        if not self.conn:
            self.connect()
        
        if params:
            return self.conn.execute(sql, params).fetchall()
        else:
            return self.conn.execute(sql).fetchall()
    
    def _get_column_type(self, value):
        """
        Get the SQL type for a value.
        """
        if isinstance(value, int):
            return "INTEGER"
        elif isinstance(value, float):
            return "FLOAT"
        elif isinstance(value, bool):
            return "BOOLEAN"
        elif isinstance(value, (datetime, str)):
            return "TEXT"
        else:
            return "TEXT"

class SQLiteDatabase(Database):
    """
    SQLite database implementation.
    """
    def connect(self):
        """
        Connect to SQLite.
        """
        sqlite_config = self.config["database"]["connection"]["sqlite"]
        
        logger.info(f"Connecting to SQLite database: {sqlite_config['path']}")
        self.conn = sqlite3.connect(sqlite_config["path"])
        
        # Enable foreign keys
        self.conn.execute("PRAGMA foreign_keys = ON")
        
        # Configure SQLite to return rows as dictionaries
        self.conn.row_factory = sqlite3.Row
        
        return self.conn
    
    def create_table(self, table_name, schema, recreate=False):
        """
        Create a table if it doesn't exist.
        
        Args:
            table_name: The name of the table to create
            schema: The schema definition for the table
            recreate: If True, drop the table if it exists before creating it
        """
        if not self.conn:
            self.connect()
        
        # If recreate is True, drop the table if it exists
        if recreate:
            logger.info(f"Dropping table {table_name} if it exists (recreate=True)")
            self.conn.execute(f"""
            DROP TABLE IF EXISTS {table_name}
            """)
            self.conn.commit()
        
        # Create the table
        self.conn.execute(f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            {schema}
        )
        """)
        
        # Check if updated_at column exists, if not add it
        cursor = self.conn.execute(f"PRAGMA table_info({table_name})")
        columns = [row[1] for row in cursor.fetchall()]
        
        if "updated_at" not in columns:
            logger.info(f"Adding updated_at column to table {table_name}")
            self.conn.execute(f"""
            ALTER TABLE {table_name} ADD COLUMN updated_at TEXT
            """)
        
        self.conn.commit()
    
    def upsert(self, table_name, data, primary_key=None):
        """
        Upsert data into a table using a true merge approach.
        This implementation updates existing records and inserts new ones without deleting any data.
        """
        if not self.conn:
            self.connect()
        
        if not data:
            return 0
        
        # Debug print the data being upserted
        self._debug_print_data("UPSERT", table_name, data, primary_key)
        
        # Get column names from the first data item
        columns = list(data[0].keys())
        if "updated_at" not in columns:
            columns.append("updated_at")
        
        # Begin transaction
        self.conn.execute("BEGIN TRANSACTION")
        
        try:
            # Check if the table exists
            cursor = self.conn.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='{table_name}'")
            table_exists = cursor.fetchone() is not None
            
            if not table_exists:
                # If table doesn't exist, create it and insert all records
                logger.info(f"Table {table_name} does not exist yet, creating and inserting all records")
                
                # Create the table with columns from the data
                column_defs = []
                for col in columns:
                    if col == primary_key:
                        column_defs.append(f"{col} TEXT PRIMARY KEY")
                    elif col == "updated_at":
                        column_defs.append(f"{col} TEXT")
                    else:
                        column_defs.append(f"{col} TEXT")
                
                create_table_sql = f"""
                CREATE TABLE {table_name} (
                    {', '.join(column_defs)}
                )
                """
                self.conn.execute(create_table_sql)
                
                # Insert all records
                current_time = datetime.utcnow().isoformat()
                for item in data:
                    # Prepare values with updated_at
                    values = [item.get(col) for col in columns if col != "updated_at"]
                    values.append(current_time)  # Add updated_at
                    
                    placeholders = ", ".join(["?" for _ in range(len(values))])
                    self.conn.execute(f"""
                    INSERT INTO {table_name} ({', '.join(columns)})
                    VALUES ({placeholders})
                    """, values)
            else:
                # If primary key is provided, use it for true upserting (merge)
                if primary_key:
                    current_time = datetime.utcnow().isoformat()
                    
                    for item in data:
                        # Check if record exists
                        cursor = self.conn.execute(f"SELECT * FROM {table_name} WHERE {primary_key} = ?", (item[primary_key],))
                        existing_record = cursor.fetchone()
                        
                        if existing_record:
                            # Update existing record
                            if self.debug_writes:
                                logger.info(f"DEBUG WRITES: Updating existing record in '{table_name}' with {primary_key} = '{item[primary_key]}'")
                            
                            # Prepare SET clause for update
                            set_clauses = []
                            update_values = []
                            
                            for col in columns:
                                if col != primary_key:  # Skip primary key in SET clause
                                    if col == "updated_at":
                                        set_clauses.append(f"{col} = ?")
                                        update_values.append(current_time)
                                    else:
                                        set_clauses.append(f"{col} = ?")
                                        update_values.append(item.get(col))
                            
                            # Add primary key value for WHERE clause
                            update_values.append(item[primary_key])
                            
                            update_sql = f"""
                            UPDATE {table_name}
                            SET {', '.join(set_clauses)}
                            WHERE {primary_key} = ?
                            """
                            
                            if self.debug_writes:
                                logger.info(f"DEBUG: Executing SQL: {update_sql}")
                                logger.info(f"DEBUG: With values: {update_values}")
                            
                            self.conn.execute(update_sql, update_values)
                        else:
                            # Insert new record
                            if self.debug_writes:
                                logger.info(f"DEBUG WRITES: Inserting new record into '{table_name}' with {primary_key} = '{item[primary_key]}'")
                            
                            # Prepare values with updated_at
                            values = [item.get(col) for col in columns if col != "updated_at"]
                            values.append(current_time)  # Add updated_at
                            
                            placeholders = ", ".join(["?" for _ in range(len(values))])
                            insert_sql = f"""
                            INSERT INTO {table_name} ({', '.join(columns)})
                            VALUES ({placeholders})
                            """
                            
                            if self.debug_writes:
                                logger.info(f"DEBUG: Executing SQL: {insert_sql}")
                                logger.info(f"DEBUG: With values: {values}")
                            
                            self.conn.execute(insert_sql, values)
                else:
                    # If no primary key is provided, just insert all records
                    # This might create duplicates, but that's expected behavior without a primary key
                    current_time = datetime.utcnow().isoformat()
                    for item in data:
                        # Prepare values with updated_at
                        values = [item.get(col) for col in columns if col != "updated_at"]
                        values.append(current_time)  # Add updated_at
                        
                        placeholders = ", ".join(["?" for _ in range(len(values))])
                        self.conn.execute(f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        VALUES ({placeholders})
                        """, values)
            
            # Commit transaction
            self.conn.commit()
            return len(data)
        
        except Exception as e:
            # Rollback transaction on error
            self.conn.rollback()
            logger.error(f"Error upserting data: {e}")
            raise
    
    def insert(self, table_name, data):
        """
        Insert data into a table without upserting.
        """
        if not self.conn:
            self.connect()
        
        if not data:
            return 0
        
        # Debug print the data being inserted
        self._debug_print_data("INSERT", table_name, data)
        
        # Get column names from the first data item
        columns = list(data[0].keys())
        if "updated_at" not in columns:
            columns.append("updated_at")
        
        # Get table schema to check column types
        try:
            cursor = self.conn.execute(f"PRAGMA table_info({table_name})")
            schema_info = cursor.fetchall()
            column_types = {row[1]: row[2] for row in schema_info}
            
            logger.info(f"DEBUG: Table schema for {table_name}: {column_types}")
        except Exception as e:
            logger.warning(f"DEBUG: Could not get schema info: {e}")
            column_types = {}
        
        # Begin transaction
        self.conn.execute("BEGIN TRANSACTION")
        
        try:
            # Insert new records
            current_time = datetime.utcnow().isoformat()
            count = 0
            
            for item in data:
                try:
                    # Prepare values with updated_at and convert timestamps
                    values = []
                    for col in columns:
                        if col == "updated_at":
                            values.append(current_time)
                        elif col in column_types and "TIMESTAMP" in column_types.get(col, "").upper() and item.get(col):
                            # Convert string timestamps to datetime objects for TIMESTAMP columns
                            try:
                                timestamp_str = item.get(col)
                                logger.info(f"DEBUG: Converting timestamp: {timestamp_str}, type: {type(timestamp_str)}")
                                
                                # If it's already a datetime object, use it directly
                                if isinstance(timestamp_str, datetime):
                                    values.append(timestamp_str.isoformat())
                                # If it's a string in ISO format, parse it
                                elif isinstance(timestamp_str, str):
                                    # Try to parse the timestamp string
                                    try:
                                        # Try ISO format with Z (UTC)
                                        dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%SZ")
                                        values.append(dt.isoformat())
                                        logger.info(f"DEBUG: Converted to datetime: {dt}")
                                    except ValueError:
                                        try:
                                            # Try ISO format without Z
                                            dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%S")
                                            values.append(dt.isoformat())
                                            logger.info(f"DEBUG: Converted to datetime: {dt}")
                                        except ValueError:
                                            # If all parsing fails, pass the string as is
                                            logger.warning(f"DEBUG: Could not parse timestamp: {timestamp_str}")
                                            values.append(timestamp_str)
                                else:
                                    # If it's neither a datetime nor a string, use as is
                                    values.append(timestamp_str)
                            except Exception as e:
                                logger.error(f"DEBUG: Error converting timestamp {item.get(col)}: {e}")
                                values.append(item.get(col))
                        else:
                            values.append(item.get(col))
                    
                    placeholders = ", ".join(["?" for _ in range(len(values))])
                    insert_sql = f"""
                    INSERT INTO {table_name} ({', '.join(columns)})
                    VALUES ({placeholders})
                    """
                    
                    if self.debug_writes:
                        logger.info(f"DEBUG: Executing SQL: {insert_sql}")
                        logger.info(f"DEBUG: With values: {values}")
                    
                    self.conn.execute(insert_sql, values)
                    count += 1
                except Exception as e:
                    logger.error(f"DEBUG: Error inserting row {count+1}: {e}")
                    logger.error(f"DEBUG: Row data: {item}")
                    raise
            
            # Commit transaction
            self.conn.commit()
            return count
        
        except Exception as e:
            # Rollback transaction on error
            self.conn.rollback()
            logger.error(f"Error inserting data: {e}")
            raise
    
    def query(self, sql, params=None):
        """
        Execute a query and return the results.
        """
        if not self.conn:
            self.connect()
        
        cursor = self.conn.cursor()
        if params:
            cursor.execute(sql, params)
        else:
            cursor.execute(sql)
        
        return cursor.fetchall()

try:
    import psycopg2
    import psycopg2.extras
    
    class PostgreSQLDatabase(Database):
        """
        PostgreSQL database implementation.
        """
        def connect(self):
            """
            Connect to PostgreSQL.
            """
            pg_config = self.config["database"]["connection"]["postgresql"]
            
            logger.info(f"Connecting to PostgreSQL database: {pg_config['host']}/{pg_config['database']}")
            self.conn = psycopg2.connect(
                host=pg_config["host"],
                port=pg_config["port"],
                database=pg_config["database"],
                user=pg_config["username"],
                password=pg_config["password"]
            )
            
            return self.conn
        
        def create_table(self, table_name, schema, recreate=False):
            """
            Create a table if it doesn't exist.
            
            Args:
                table_name: The name of the table to create
                schema: The schema definition for the table
                recreate: If True, drop the table if it exists before creating it
            """
            if not self.conn:
                self.connect()
            
            # Create the table
            with self.conn.cursor() as cursor:
                # If recreate is True, drop the table if it exists
                if recreate:
                    logger.info(f"Dropping table {table_name} if it exists (recreate=True)")
                    cursor.execute(f"""
                    DROP TABLE IF EXISTS {table_name}
                    """)
                
                cursor.execute(f"""
                CREATE TABLE IF NOT EXISTS {table_name} (
                    {schema}
                )
                """)
                
                # Check if updated_at column exists, if not add it
                cursor.execute(f"""
                SELECT column_name
                FROM information_schema.columns
                WHERE table_name = '{table_name}' AND column_name = 'updated_at'
                """)
                
                if not cursor.fetchone():
                    logger.info(f"Adding updated_at column to table {table_name}")
                    cursor.execute(f"""
                    ALTER TABLE {table_name} ADD COLUMN updated_at TIMESTAMP
                    """)
                
                self.conn.commit()
        
        def upsert(self, table_name, data, primary_key=None):
            """
            Upsert data into a table using a true merge approach.
            This implementation updates existing records and inserts new ones without deleting any data.
            """
            if not self.conn:
                self.connect()
            
            if not data:
                return 0
            
            # Debug print the data being upserted
            self._debug_print_data("UPSERT", table_name, data, primary_key)
            
            # Get column names from the first data item
            columns = list(data[0].keys())
            if "updated_at" not in columns:
                columns.append("updated_at")
            
            # Begin transaction
            with self.conn.cursor() as cursor:
                # Check if the table exists
                cursor.execute(f"""
                SELECT EXISTS (
                    SELECT FROM information_schema.tables
                    WHERE table_name = '{table_name}'
                )
                """)
                table_exists = cursor.fetchone()[0]
                
                if not table_exists:
                    # If table doesn't exist, create it and insert all records
                    logger.info(f"Table {table_name} does not exist yet, creating and inserting all records")
                    
                    # Create the table with columns from the data
                    column_defs = []
                    for col in columns:
                        if col == primary_key:
                            column_defs.append(f"{col} TEXT PRIMARY KEY")
                        elif col == "updated_at":
                            column_defs.append(f"{col} TIMESTAMP")
                        else:
                            column_defs.append(f"{col} TEXT")
                    
                    cursor.execute(f"""
                    CREATE TABLE {table_name} (
                        {', '.join(column_defs)}
                    )
                    """)
                    
                    # Insert all records
                    current_time = datetime.utcnow()
                    for item in data:
                        # Prepare values with updated_at
                        values = [item.get(col) for col in columns if col != "updated_at"]
                        values.append(current_time)  # Add updated_at
                        
                        placeholders = ", ".join(["%s" for _ in range(len(values))])
                        cursor.execute(f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        VALUES ({placeholders})
                        """, values)
                else:
                    # If primary key is provided, use it for true upserting (merge)
                    if primary_key:
                        current_time = datetime.utcnow()
                        
                        # PostgreSQL supports ON CONFLICT DO UPDATE for true upserts
                        for item in data:
                            # Prepare values with updated_at
                            values = [item.get(col) for col in columns if col != "updated_at"]
                            values.append(current_time)  # Add updated_at
                            
                            # Log the operation if debug_writes is enabled
                            if self.debug_writes:
                                cursor.execute(f"SELECT * FROM {table_name} WHERE {primary_key} = %s", (item[primary_key],))
                                existing_record = cursor.fetchone()
                                if existing_record:
                                    logger.info(f"DEBUG WRITES: Updating existing record in '{table_name}' with {primary_key} = '{item[primary_key]}'")
                                else:
                                    logger.info(f"DEBUG WRITES: Inserting new record into '{table_name}' with {primary_key} = '{item[primary_key]}'")
                            
                            # Prepare the ON CONFLICT clause
                            non_pk_columns = [col for col in columns if col != primary_key]
                            set_clauses = [f"{col} = EXCLUDED.{col}" for col in non_pk_columns]
                            
                            placeholders = ", ".join(["%s" for _ in range(len(values))])
                            upsert_sql = f"""
                            INSERT INTO {table_name} ({', '.join(columns)})
                            VALUES ({placeholders})
                            ON CONFLICT ({primary_key}) DO UPDATE
                            SET {', '.join(set_clauses)}
                            """
                            
                            if self.debug_writes:
                                logger.info(f"DEBUG: Executing SQL: {upsert_sql}")
                                logger.info(f"DEBUG: With values: {values}")
                            
                            cursor.execute(upsert_sql, values)
                    else:
                        # If no primary key is provided, just insert all records
                        # This might create duplicates, but that's expected behavior without a primary key
                        current_time = datetime.utcnow()
                        for item in data:
                            # Prepare values with updated_at
                            values = [item.get(col) for col in columns if col != "updated_at"]
                            values.append(current_time)  # Add updated_at
                            
                            placeholders = ", ".join(["%s" for _ in range(len(values))])
                            cursor.execute(f"""
                            INSERT INTO {table_name} ({', '.join(columns)})
                            VALUES ({placeholders})
                            """, values)
                
                self.conn.commit()
                return len(data)
        
        def insert(self, table_name, data):
            """
            Insert data into a table without upserting.
            """
            if not self.conn:
                self.connect()
            
            if not data:
                return 0
            
            # Debug print the data being inserted
            self._debug_print_data("INSERT", table_name, data)
            
            # Get column names from the first data item
            columns = list(data[0].keys())
            if "updated_at" not in columns:
                columns.append("updated_at")
            
            # Get table schema to check column types
            try:
                with self.conn.cursor() as cursor:
                    cursor.execute(f"""
                    SELECT column_name, data_type
                    FROM information_schema.columns
                    WHERE table_name = '{table_name}'
                    """)
                    schema_info = cursor.fetchall()
                    column_types = {col[0]: col[1] for col in schema_info}
                    
                    logger.info(f"DEBUG: Table schema for {table_name}: {column_types}")
            except Exception as e:
                logger.warning(f"DEBUG: Could not get schema info: {e}")
                column_types = {}
            
            # Begin transaction
            with self.conn.cursor() as cursor:
                # Insert new records
                current_time = datetime.utcnow()
                count = 0
                
                for item in data:
                    try:
                        # Prepare values with updated_at and convert timestamps
                        values = []
                        for col in columns:
                            if col == "updated_at":
                                values.append(current_time)
                            elif col in column_types and "TIMESTAMP" in column_types.get(col, "").upper() and item.get(col):
                                # Convert string timestamps to datetime objects for TIMESTAMP columns
                                try:
                                    timestamp_str = item.get(col)
                                    logger.info(f"DEBUG: Converting timestamp: {timestamp_str}, type: {type(timestamp_str)}")
                                    
                                    # If it's already a datetime object, use it directly
                                    if isinstance(timestamp_str, datetime):
                                        values.append(timestamp_str)
                                    # If it's a string in ISO format, parse it
                                    elif isinstance(timestamp_str, str):
                                        # Try to parse the timestamp string
                                        try:
                                            # Try ISO format with Z (UTC)
                                            dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%SZ")
                                            values.append(dt)
                                            logger.info(f"DEBUG: Converted to datetime: {dt}")
                                        except ValueError:
                                            try:
                                                # Try ISO format without Z
                                                dt = datetime.strptime(timestamp_str, "%Y-%m-%dT%H:%M:%S")
                                                values.append(dt)
                                                logger.info(f"DEBUG: Converted to datetime: {dt}")
                                            except ValueError:
                                                # If all parsing fails, pass the string as is
                                                logger.warning(f"DEBUG: Could not parse timestamp: {timestamp_str}")
                                                values.append(timestamp_str)
                                    else:
                                        # If it's neither a datetime nor a string, use as is
                                        values.append(timestamp_str)
                                except Exception as e:
                                    logger.error(f"DEBUG: Error converting timestamp {item.get(col)}: {e}")
                                    values.append(item.get(col))
                            else:
                                values.append(item.get(col))
                        
                        placeholders = ", ".join(["%s" for _ in range(len(values))])
                        insert_sql = f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        VALUES ({placeholders})
                        """
                        
                        if self.debug_writes:
                            logger.info(f"DEBUG: Executing SQL: {insert_sql}")
                            logger.info(f"DEBUG: With values: {values}")
                        
                        cursor.execute(insert_sql, values)
                        count += 1
                    except Exception as e:
                        logger.error(f"DEBUG: Error inserting row {count+1}: {e}")
                        logger.error(f"DEBUG: Row data: {item}")
                        raise
                
                self.conn.commit()
                return count
        
        def query(self, sql, params=None):
            """
            Execute a query and return the results.
            """
            if not self.conn:
                self.connect()
            
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
                if params:
                    cursor.execute(sql, params)
                else:
                    cursor.execute(sql)
                
                return cursor.fetchall()
    
    class RDSDatabase(PostgreSQLDatabase):
        """
        Amazon RDS database implementation.
        Inherits from PostgreSQLDatabase since RDS can use PostgreSQL.
        """
        def connect(self):
            """
            Connect to Amazon RDS.
            """
            rds_config = self.config["database"]["connection"]["rds"]
            
            logger.info(f"Connecting to RDS database: {rds_config['host']}/{rds_config['database']}")
            self.conn = psycopg2.connect(
                host=rds_config["host"],
                port=rds_config["port"],
                database=rds_config["database"],
                user=rds_config["username"],
                password=rds_config["password"]
            )
            
            return self.conn
except ImportError:
    # If psycopg2 is not installed, create placeholder classes
    class PostgreSQLDatabase(Database):
        def __init__(self, *args, **kwargs):
            raise ImportError("psycopg2 is not installed. Please install it to use PostgreSQL.")
    
    class RDSDatabase(Database):
        def __init__(self, *args, **kwargs):
            raise ImportError("psycopg2 is not installed. Please install it to use RDS.")

def get_database(config):
    """
    Get a database instance based on the configuration.
    """
    db_type = config["database"]["type"]
    
    if db_type == "motherduck":
        return MotherDuckDatabase(config)
    elif db_type == "postgresql":
        return PostgreSQLDatabase(config)
    elif db_type == "sqlite":
        return SQLiteDatabase(config)
    elif db_type == "rds":
        return RDSDatabase(config)
    else:
        raise ValueError(f"Unsupported database type: {db_type}")