ferrero-opentext/Python-Version/scripts/shared/database.py

"""
Database Client - PostgreSQL Operations
Handles tracking IDs, master assets, and all-done checks
Compatible with Python 3.6+
"""

import psycopg2
import psycopg2.pool
import json
import random
import string
import logging

logger = logging.getLogger('Database')

class Database:
    def __init__(self, config):
        self.config = config['database']

        # Create connection pool
        try:
            self.pool = psycopg2.pool.ThreadedConnectionPool(
                minconn=1,
                maxconn=10,
                host=self.config['host'],
                port=self.config['port'],
                database=self.config['database'],
                user=self.config['user'],
                password=self.config['password']
            )
            logger.info("Database connection pool created")
        except Exception as e:
            logger.error("Database connection failed: {}".format(str(e)))
            raise

    def get_connection(self):
        """Get connection from pool"""
        return self.pool.getconn()

    def put_connection(self, conn):
        """Return connection to pool"""
        self.pool.putconn(conn)

    def generate_unique_tracking_id(self):
        """
        Generate unique 6-character tracking ID

        Returns:
            str: 6-character alphanumeric ID
        """
        conn = self.get_connection()
        try:
            cursor = conn.cursor()

            for attempt in range(100):
                # Generate random 6-char ID
                tracking_id = ''.join(random.choices(
                    string.ascii_letters + string.digits, k=6
                ))

                # Check uniqueness
                cursor.execute(
                    "SELECT COUNT(*) FROM master_assets WHERE tracking_id = %s",
                    (tracking_id,)
                )

                count = cursor.fetchone()[0]

                if count == 0:
                    logger.info("Generated tracking ID: {}".format(tracking_id))
                    return tracking_id

            raise Exception("Failed to generate unique tracking ID after 100 attempts")

        finally:
            cursor.close()
            self.put_connection(conn)

    def store_master_asset(self, tracking_id, opentext_id, asset_data, box_file_id, box_url, upload_folder_id):
        """
        Store master asset with FULL metadata in JSONB column

        Args:
            tracking_id: 6-char tracking ID
            opentext_id: DAM asset ID
            asset_data: Complete DAM asset JSON
            box_file_id: Box file ID
            box_url: Box URL
            upload_folder_id: Final Assets folder ID for upload

        Returns:
            dict with success boolean
        """
        conn = self.get_connection()
        try:
            cursor = conn.cursor()

            # Extract basic info
            name = asset_data.get('name', 'unknown')
            name_parts = name.rsplit('.', 1)
            filename = name_parts[0]
            extension = '.' + name_parts[1] if len(name_parts) > 1 else ''

            # Store complete metadata as JSONB
            full_metadata_json = json.dumps(asset_data)

            # Description with Box info
            description = "Box File ID: {}\nBox URL: {}\nDAM Asset ID: {}".format(
                box_file_id, box_url, opentext_id
            )

            # Insert or update
            cursor.execute("""
                INSERT INTO master_assets (
                    tracking_id, opentext_id, original_filename, file_extension,
                    file_size_bytes, mime_type, upload_directory,
                    description, full_metadata, status
                ) VALUES (
                    %s, %s, %s, %s, %s, %s, %s, %s, %s, 'active'
                )
                ON CONFLICT (tracking_id) DO UPDATE SET
                    upload_directory = EXCLUDED.upload_directory,
                    description = EXCLUDED.description,
                    full_metadata = EXCLUDED.full_metadata,
                    updated_at = CURRENT_TIMESTAMP
            """, (
                tracking_id,
                opentext_id,
                filename,
                extension,
                asset_data.get('file_size'),
                asset_data.get('mime_type'),
                upload_folder_id,
                description,
                full_metadata_json
            ))

            conn.commit()
            logger.info("Stored master asset: {}".format(tracking_id))

            return {'success': True, 'tracking_id': tracking_id}

        except Exception as e:
            conn.rollback()
            logger.error("Failed to store master asset: {}".format(str(e)))
            return {'success': False, 'error': str(e)}

        finally:
            cursor.close()
            self.put_connection(conn)

    def get_master_asset(self, tracking_id):
        """
        Get master asset by tracking ID

        Returns:
            dict with tracking_id, opentext_id, upload_directory, full_metadata
        """
        conn = self.get_connection()
        try:
            cursor = conn.cursor()

            cursor.execute("""
                SELECT tracking_id, opentext_id, upload_directory, full_metadata, description
                FROM master_assets
                WHERE tracking_id = %s AND status = 'active'
            """, (tracking_id,))

            row = cursor.fetchone()

            if not row:
                return None

            # Parse JSONB as dict
            full_metadata = row[3] if isinstance(row[3], dict) else json.loads(row[3])

            return {
                'tracking_id': row[0],
                'opentext_id': row[1],
                'upload_directory': row[2],
                'full_metadata': full_metadata,
                'description': row[4]
            }

        finally:
            cursor.close()
            self.put_connection(conn)

    def check_campaign_upload_complete(self, campaign_id):
        """
        Check if ALL master assets for a campaign have been uploaded

        Args:
            campaign_id: Campaign ID

        Returns:
            bool: True if all assets uploaded, False otherwise
        """
        conn = self.get_connection()
        try:
            cursor = conn.cursor()

            # Count total master assets for this campaign
            cursor.execute("""
                SELECT COUNT(DISTINCT tracking_id)
                FROM master_assets
                WHERE campaign_id = %s AND status = 'active'
            """, (campaign_id,))

            total_masters = cursor.fetchone()[0]

            if total_masters == 0:
                return False

            # Count how many have been uploaded (exist in derivative_assets)
            cursor.execute("""
                SELECT COUNT(DISTINCT ma.tracking_id)
                FROM master_assets ma
                INNER JOIN derivative_assets da ON ma.tracking_id = da.tracking_id
                WHERE ma.campaign_id = %s AND ma.status = 'active'
                AND da.upload_status = 'completed'
            """, (campaign_id,))

            uploaded_count = cursor.fetchone()[0]

            all_done = uploaded_count == total_masters

            logger.info("Campaign {} upload status: {}/{} assets uploaded{}".format(
                campaign_id, uploaded_count, total_masters,
                " - ALL DONE" if all_done else ""
            ))

            return all_done

        finally:
            cursor.close()
            self.put_connection(conn)

    def store_derivative_asset(self, tracking_id, master_asset_id, dam_asset_id, filename):
        """
        Store derivative asset record after upload

        Args:
            tracking_id: Master asset tracking ID
            master_asset_id: Master asset DB ID
            dam_asset_id: DAM asset ID of uploaded derivative
            filename: Clean filename of derivative

        Returns:
            dict with success boolean
        """
        conn = self.get_connection()
        try:
            cursor = conn.cursor()

            cursor.execute("""
                INSERT INTO derivative_assets (
                    tracking_id, master_asset_id, dam_asset_id,
                    derivative_filename, upload_status
                ) VALUES (%s, %s, %s, %s, 'completed')
                ON CONFLICT (tracking_id, derivative_filename) DO UPDATE SET
                    dam_asset_id = EXCLUDED.dam_asset_id,
                    uploaded_at = CURRENT_TIMESTAMP,
                    upload_status = 'completed'
            """, (tracking_id, master_asset_id, dam_asset_id, filename))

            conn.commit()
            logger.info("Stored derivative asset: {} → {}".format(tracking_id, dam_asset_id))

            return {'success': True}

        except Exception as e:
            conn.rollback()
            logger.error("Failed to store derivative: {}".format(str(e)))
            return {'success': False, 'error': str(e)}

        finally:
            cursor.close()
            self.put_connection(conn)

    def get_campaign_asset_count(self, campaign_id):
        """Get total master asset count for campaign"""
        conn = self.get_connection()
        try:
            cursor = conn.cursor()

            cursor.execute("""
                SELECT COUNT(*) FROM master_assets
                WHERE campaign_id = %s AND status = 'active'
            """, (campaign_id,))

            return cursor.fetchone()[0]

        finally:
            cursor.close()
            self.put_connection(conn)

    def test_connection(self):
        """Test database connection"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            cursor.execute("SELECT 1")
            cursor.close()
            self.put_connection(conn)
            logger.info("Database connection OK")
            return True
        except Exception as e:
            logger.error("Database connection failed: {}".format(str(e)))
            return False

    def close(self):
        """Close all connections in pool"""
        if self.pool:
            self.pool.closeall()