unclecode--crawl4ai/crawl4ai/async_database.py

import os
import time
from pathlib import Path
import aiosqlite
import asyncio
from typing import Optional, Dict
from contextlib import asynccontextmanager
import json
from .models import CrawlResult, MarkdownGenerationResult, StringCompatibleMarkdown
import aiofiles
from .async_logger import AsyncLogger

from .utils import ensure_content_dirs, generate_content_hash
from .utils import VersionManager
from .utils import get_error_context, create_box_message

base_directory = DB_PATH = os.path.join(
    os.getenv("CRAWL4_AI_BASE_DIRECTORY", Path.home()), ".crawl4ai"
)
os.makedirs(DB_PATH, exist_ok=True)
DB_PATH = os.path.join(base_directory, "crawl4ai.db")


class AsyncDatabaseManager:
    def __init__(self, pool_size: int = 10, max_retries: int = 3):
        self.db_path = DB_PATH
        self.content_paths = ensure_content_dirs(os.path.dirname(DB_PATH))
        self.pool_size = pool_size
        self.max_retries = max_retries
        self.connection_pool: Dict[int, aiosqlite.Connection] = {}
        self.pool_lock = asyncio.Lock()
        self.init_lock = asyncio.Lock()
        self.connection_semaphore = asyncio.Semaphore(pool_size)
        self._initialized = False
        self.version_manager = VersionManager()
        self.logger = AsyncLogger(
            log_file=os.path.join(base_directory, ".crawl4ai", "crawler_db.log"),
            verbose=False,
            tag_width=10,
        )

    async def initialize(self):
        """Initialize the database and connection pool"""
        try:
            self.logger.info("Initializing database", tag="INIT")
            # Ensure the database file exists
            os.makedirs(os.path.dirname(self.db_path), exist_ok=True)

            # Check if version update is needed
            needs_update = self.version_manager.needs_update()

            # Always ensure base table exists
            await self.ainit_db()

            # Verify the table exists
            async with aiosqlite.connect(self.db_path, timeout=30.0) as db:
                async with db.execute(
                    "SELECT name FROM sqlite_master WHERE type='table' AND name='crawled_data'"
                ) as cursor:
                    result = await cursor.fetchone()
                    if not result:
                        raise Exception("crawled_data table was not created")

            # If version changed or fresh install, run updates
            if needs_update:
                self.logger.info("New version detected, running updates", tag="INIT")
                await self.update_db_schema()
                from .migrations import (
                    run_migration,
                )  # Import here to avoid circular imports

                await run_migration()
                self.version_manager.update_version()  # Update stored version after successful migration
                self.logger.success(
                    "Version update completed successfully", tag="COMPLETE"
                )
            else:
                self.logger.success(
                    "Database initialization completed successfully", tag="COMPLETE"
                )

        except Exception as e:
            self.logger.error(
                message="Database initialization error: {error}",
                tag="ERROR",
                params={"error": str(e)},
            )
            self.logger.info(
                message="Database will be initialized on first use", tag="INIT"
            )

            raise

    async def cleanup(self):
        """Cleanup connections when shutting down"""
        async with self.pool_lock:
            for conn in self.connection_pool.values():
                await conn.close()
            self.connection_pool.clear()

    @asynccontextmanager
    async def get_connection(self):
        """Connection pool manager with enhanced error handling"""
        if not self._initialized:
            async with self.init_lock:
                if not self._initialized:
                    try:
                        await self.initialize()
                        self._initialized = True
                    except Exception as e:
                        import sys

                        error_context = get_error_context(sys.exc_info())
                        self.logger.error(
                            message="Database initialization failed:\n{error}\n\nContext:\n{context}\n\nTraceback:\n{traceback}",
                            tag="ERROR",
                            force_verbose=True,
                            params={
                                "error": str(e),
                                "context": error_context["code_context"],
                                "traceback": error_context["full_traceback"],
                            },
                        )
                        raise

        await self.connection_semaphore.acquire()
        task_id = id(asyncio.current_task())

        try:
            async with self.pool_lock:
                if task_id not in self.connection_pool:
                    try:
                        conn = await aiosqlite.connect(self.db_path, timeout=30.0)
                        await conn.execute("PRAGMA journal_mode = WAL")
                        await conn.execute("PRAGMA busy_timeout = 5000")

                        # Verify database structure
                        async with conn.execute(
                            "PRAGMA table_info(crawled_data)"
                        ) as cursor:
                            columns = await cursor.fetchall()
                            column_names = [col[1] for col in columns]
                            expected_columns = {
                                "url",
                                "html",
                                "cleaned_html",
                                "markdown",
                                "extracted_content",
                                "success",
                                "media",
                                "links",
                                "metadata",
                                "screenshot",
                                "response_headers",
                                "downloaded_files",
                            }
                            missing_columns = expected_columns - set(column_names)
                            if missing_columns:
                                raise ValueError(
                                    f"Database missing columns: {missing_columns}"
                                )

                        self.connection_pool[task_id] = conn
                    except Exception as e:
                        import sys

                        error_context = get_error_context(sys.exc_info())
                        error_message = (
                            f"Unexpected error in db get_connection at line {error_context['line_no']} "
                            f"in {error_context['function']} ({error_context['filename']}):\n"
                            f"Error: {str(e)}\n\n"
                            f"Code context:\n{error_context['code_context']}"
                        )
                        self.logger.error(
                            message="{error}",
                            tag="ERROR",
                            params={"error": str(error_message)},
                            boxes=["error"],
                        )

                        raise

            yield self.connection_pool[task_id]

        except Exception as e:
            import sys

            error_context = get_error_context(sys.exc_info())
            error_message = (
                f"Unexpected error in db get_connection at line {error_context['line_no']} "
                f"in {error_context['function']} ({error_context['filename']}):\n"
                f"Error: {str(e)}\n\n"
                f"Code context:\n{error_context['code_context']}"
            )
            self.logger.error(
                message="{error}",
                tag="ERROR",
                params={"error": str(error_message)},
                boxes=["error"],
            )
            raise
        finally:
            async with self.pool_lock:
                if task_id in self.connection_pool:
                    await self.connection_pool[task_id].close()
                    del self.connection_pool[task_id]
            self.connection_semaphore.release()

    async def execute_with_retry(self, operation, *args):
        """Execute database operations with retry logic"""
        for attempt in range(self.max_retries):
            try:
                async with self.get_connection() as db:
                    result = await operation(db, *args)
                    await db.commit()
                    return result
            except Exception as e:
                if attempt == self.max_retries - 1:
                    self.logger.error(
                        message="Operation failed after {retries} attempts: {error}",
                        tag="ERROR",
                        force_verbose=True,
                        params={"retries": self.max_retries, "error": str(e)},
                    )
                    raise
                await asyncio.sleep(1 * (attempt + 1))  # Exponential backoff

    async def ainit_db(self):
        """Initialize database schema"""
        async with aiosqlite.connect(self.db_path, timeout=30.0) as db:
            await db.execute(
                """
                CREATE TABLE IF NOT EXISTS crawled_data (
                    url TEXT PRIMARY KEY,
                    html TEXT,
                    cleaned_html TEXT,
                    markdown TEXT,
                    extracted_content TEXT,
                    success BOOLEAN,
                    media TEXT DEFAULT "{}",
                    links TEXT DEFAULT "{}",
                    metadata TEXT DEFAULT "{}",
                    screenshot TEXT DEFAULT "",
                    response_headers TEXT DEFAULT "{}",
                    downloaded_files TEXT DEFAULT "{}"  -- New column added
                )
            """
            )
            await db.commit()

    async def update_db_schema(self):
        """Update database schema if needed"""
        async with aiosqlite.connect(self.db_path, timeout=30.0) as db:
            cursor = await db.execute("PRAGMA table_info(crawled_data)")
            columns = await cursor.fetchall()
            column_names = [column[1] for column in columns]

            # List of new columns to add
            new_columns = [
                "media",
                "links",
                "metadata",
                "screenshot",
                "response_headers",
                "downloaded_files",
                # Smart cache validation columns (added in 0.8.x)
                "etag",
                "last_modified",
                "head_fingerprint",
                "cached_at",
            ]

            for column in new_columns:
                if column not in column_names:
                    await self.aalter_db_add_column(column, db)
            await db.commit()

    async def aalter_db_add_column(self, new_column: str, db):
        """Add new column to the database"""
        if new_column == "response_headers":
            await db.execute(
                f'ALTER TABLE crawled_data ADD COLUMN {new_column} TEXT DEFAULT "{{}}"'
            )
        elif new_column == "cached_at":
            # Timestamp column for cache validation
            await db.execute(
                f"ALTER TABLE crawled_data ADD COLUMN {new_column} REAL DEFAULT 0"
            )
        else:
            await db.execute(
                f'ALTER TABLE crawled_data ADD COLUMN {new_column} TEXT DEFAULT ""'
            )
        self.logger.info(
            message="Added column '{column}' to the database",
            tag="INIT",
            params={"column": new_column},
        )

    async def aget_cached_url(self, url: str) -> Optional[CrawlResult]:
        """Retrieve cached URL data as CrawlResult"""

        async def _get(db):
            async with db.execute(
                "SELECT * FROM crawled_data WHERE url = ?", (url,)
            ) as cursor:
                row = await cursor.fetchone()
                if not row:
                    return None

                # Get column names
                columns = [description[0] for description in cursor.description]
                # Create dict from row data
                row_dict = dict(zip(columns, row))

                # Load content from files using stored hashes
                content_fields = {
                    "html": row_dict["html"],
                    "cleaned_html": row_dict["cleaned_html"],
                    "markdown": row_dict["markdown"],
                    "extracted_content": row_dict["extracted_content"],
                    "screenshot": row_dict["screenshot"],
                    "screenshots": row_dict["screenshot"],
                }

                for field, hash_value in content_fields.items():
                    if hash_value:
                        content = await self._load_content(
                            hash_value,
                            field.split("_")[0],  # Get content type from field name
                        )
                        row_dict[field] = content or ""
                    else:
                        row_dict[field] = ""

                # Parse JSON fields
                json_fields = [
                    "media",
                    "links",
                    "metadata",
                    "response_headers",
                    "markdown",
                ]
                for field in json_fields:
                    try:
                        row_dict[field] = (
                            json.loads(row_dict[field]) if row_dict[field] else {}
                        )
                    except json.JSONDecodeError:
                        # Very UGLY, never mention it to me please
                        if field == "markdown" and isinstance(row_dict[field], str):
                            row_dict[field] = MarkdownGenerationResult(
                                raw_markdown=row_dict[field] or "",
                                markdown_with_citations="",
                                references_markdown="",
                                fit_markdown="",
                                fit_html="",
                            )
                        else:
                            row_dict[field] = {}

                if isinstance(row_dict["markdown"], Dict):
                    if row_dict["markdown"].get("raw_markdown"):
                        row_dict["markdown"] = row_dict["markdown"]["raw_markdown"]

                # Parse downloaded_files
                try:
                    row_dict["downloaded_files"] = (
                        json.loads(row_dict["downloaded_files"])
                        if row_dict["downloaded_files"]
                        else []
                    )
                except json.JSONDecodeError:
                    row_dict["downloaded_files"] = []

                # Remove any fields not in CrawlResult model
                valid_fields = CrawlResult.__annotations__.keys()
                filtered_dict = {k: v for k, v in row_dict.items() if k in valid_fields}
                filtered_dict["markdown"] = row_dict["markdown"]
                return CrawlResult(**filtered_dict)

        try:
            return await self.execute_with_retry(_get)
        except Exception as e:
            self.logger.error(
                message="Error retrieving cached URL: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )
            return None

    async def aget_cache_metadata(self, url: str) -> Optional[Dict]:
        """
        Retrieve only cache validation metadata for a URL (lightweight query).

        Returns dict with: url, etag, last_modified, head_fingerprint, cached_at, response_headers
        This is used for cache validation without loading full content.
        """
        async def _get_metadata(db):
            async with db.execute(
                """SELECT url, etag, last_modified, head_fingerprint, cached_at, response_headers
                   FROM crawled_data WHERE url = ?""",
                (url,)
            ) as cursor:
                row = await cursor.fetchone()
                if not row:
                    return None

                columns = [description[0] for description in cursor.description]
                row_dict = dict(zip(columns, row))

                # Parse response_headers JSON
                try:
                    row_dict["response_headers"] = (
                        json.loads(row_dict["response_headers"])
                        if row_dict["response_headers"] else {}
                    )
                except json.JSONDecodeError:
                    row_dict["response_headers"] = {}

                return row_dict

        try:
            return await self.execute_with_retry(_get_metadata)
        except Exception as e:
            self.logger.error(
                message="Error retrieving cache metadata: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )
            return None

    async def aupdate_cache_metadata(
        self,
        url: str,
        etag: Optional[str] = None,
        last_modified: Optional[str] = None,
        head_fingerprint: Optional[str] = None,
    ):
        """
        Update only the cache validation metadata for a URL.
        Used to update etag/last_modified after a successful validation.
        """
        async def _update(db):
            updates = []
            values = []

            if etag is not None:
                updates.append("etag = ?")
                values.append(etag)
            if last_modified is not None:
                updates.append("last_modified = ?")
                values.append(last_modified)
            if head_fingerprint is not None:
                updates.append("head_fingerprint = ?")
                values.append(head_fingerprint)

            if not updates:
                return

            values.append(url)
            await db.execute(
                f"UPDATE crawled_data SET {', '.join(updates)} WHERE url = ?",
                tuple(values)
            )

        try:
            await self.execute_with_retry(_update)
        except Exception as e:
            self.logger.error(
                message="Error updating cache metadata: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )

    async def acache_url(self, result: CrawlResult):
        """Cache CrawlResult data"""
        # Store content files and get hashes
        content_map = {
            "html": (result.html, "html"),
            "cleaned_html": (result.cleaned_html or "", "cleaned"),
            "markdown": None,
            "extracted_content": (result.extracted_content or "", "extracted"),
            "screenshot": (result.screenshot or "", "screenshots"),
        }

        try:
            if isinstance(result.markdown, StringCompatibleMarkdown):
                content_map["markdown"] = (
                    result.markdown,
                    "markdown",
                )
            elif isinstance(result.markdown, MarkdownGenerationResult):
                content_map["markdown"] = (
                    result.markdown.model_dump_json(),
                    "markdown",
                )
            elif isinstance(result.markdown, str):
                markdown_result = MarkdownGenerationResult(raw_markdown=result.markdown)
                content_map["markdown"] = (
                    markdown_result.model_dump_json(),
                    "markdown",
                )
            else:
                content_map["markdown"] = (
                    MarkdownGenerationResult().model_dump_json(),
                    "markdown",
                )
        except Exception as e:
            self.logger.warning(
                message=f"Error processing markdown content: {str(e)}", tag="WARNING"
            )
            # Fallback to empty markdown result
            content_map["markdown"] = (
                MarkdownGenerationResult().model_dump_json(),
                "markdown",
            )

        content_hashes = {}
        for field, (content, content_type) in content_map.items():
            content_hashes[field] = await self._store_content(content, content_type)

        # Extract cache validation headers from response
        response_headers = result.response_headers or {}
        etag = response_headers.get("etag") or response_headers.get("ETag") or ""
        last_modified = response_headers.get("last-modified") or response_headers.get("Last-Modified") or ""
        # head_fingerprint is set by caller via result attribute (if available)
        head_fingerprint = getattr(result, "head_fingerprint", None) or ""
        cached_at = time.time()

        async def _cache(db):
            await db.execute(
                """
                INSERT INTO crawled_data (
                    url, html, cleaned_html, markdown,
                    extracted_content, success, media, links, metadata,
                    screenshot, response_headers, downloaded_files,
                    etag, last_modified, head_fingerprint, cached_at
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(url) DO UPDATE SET
                    html = excluded.html,
                    cleaned_html = excluded.cleaned_html,
                    markdown = excluded.markdown,
                    extracted_content = excluded.extracted_content,
                    success = excluded.success,
                    media = excluded.media,
                    links = excluded.links,
                    metadata = excluded.metadata,
                    screenshot = excluded.screenshot,
                    response_headers = excluded.response_headers,
                    downloaded_files = excluded.downloaded_files,
                    etag = excluded.etag,
                    last_modified = excluded.last_modified,
                    head_fingerprint = excluded.head_fingerprint,
                    cached_at = excluded.cached_at
            """,
                (
                    result.url,
                    content_hashes["html"],
                    content_hashes["cleaned_html"],
                    content_hashes["markdown"],
                    content_hashes["extracted_content"],
                    result.success,
                    json.dumps(result.media),
                    json.dumps(result.links),
                    json.dumps(result.metadata or {}),
                    content_hashes["screenshot"],
                    json.dumps(result.response_headers or {}),
                    json.dumps(result.downloaded_files or []),
                    etag,
                    last_modified,
                    head_fingerprint,
                    cached_at,
                ),
            )

        try:
            await self.execute_with_retry(_cache)
        except Exception as e:
            self.logger.error(
                message="Error caching URL: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )

    async def aget_total_count(self) -> int:
        """Get total number of cached URLs"""

        async def _count(db):
            async with db.execute("SELECT COUNT(*) FROM crawled_data") as cursor:
                result = await cursor.fetchone()
                return result[0] if result else 0

        try:
            return await self.execute_with_retry(_count)
        except Exception as e:
            self.logger.error(
                message="Error getting total count: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )
            return 0

    async def aclear_db(self):
        """Clear all data from the database"""

        async def _clear(db):
            await db.execute("DELETE FROM crawled_data")

        try:
            await self.execute_with_retry(_clear)
        except Exception as e:
            self.logger.error(
                message="Error clearing database: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )

    async def aflush_db(self):
        """Drop the entire table"""

        async def _flush(db):
            await db.execute("DROP TABLE IF EXISTS crawled_data")

        try:
            await self.execute_with_retry(_flush)
        except Exception as e:
            self.logger.error(
                message="Error flushing database: {error}",
                tag="ERROR",
                force_verbose=True,
                params={"error": str(e)},
            )

    async def _store_content(self, content: str, content_type: str) -> str:
        """Store content in filesystem and return hash"""
        if not content:
            return ""

        content_hash = generate_content_hash(content)
        file_path = os.path.join(self.content_paths[content_type], content_hash)

        # Only write if file doesn't exist
        if not os.path.exists(file_path):
            async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
                await f.write(content)

        return content_hash

    async def _load_content(
        self, content_hash: str, content_type: str
    ) -> Optional[str]:
        """Load content from filesystem by hash"""
        if not content_hash:
            return None

        file_path = os.path.join(self.content_paths[content_type], content_hash)
        try:
            async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
                return await f.read()
        except:
            self.logger.error(
                message="Failed to load content: {file_path}",
                tag="ERROR",
                force_verbose=True,
                params={"file_path": file_path},
            )
            return None


# Create a singleton instance
async_db_manager = AsyncDatabaseManager()