Source code for operators.bedrock_scraper.website_scraper_operator

from typing import Literal
from airflow.utils.decorators import apply_defaults
from operators.bedrock_scraper.website_converter.website_converter import (
    WebsiteToHTML,
    WebsiteToPDf,
)
from operators.bedrock_scraper.website_reader.sitemap_reader import GeneralSiteMapReader
from operators.bedrock_scraper.website_reader.govtrack_reader import GovtrackReader
from operators.update_knowledge_base import UpdateKnowledgeBaseOperator
from util import Util
import logging
from airflow.exceptions import AirflowSkipException




class WebsiteScraperOperator(UpdateKnowledgeBaseOperator):
    """
    Airflow operator for scraping websites using sitemaps and updating knowledge bases.

    This operator performs the following tasks:
    1. Reads and parses a sitemap from a given URL
    2. Converts web pages to PDF or HTML format
    3. Stores the converted files in S3
    4. Updates the specified knowledge base with the new content

    Attributes:
        ui_color (str): Color code for the operator in Airflow UI
        sitemap_url (str): URL of the sitemap to process
        s3_prefix (str): Prefix for S3 storage path
        s3_bucket_name (str): Name of the S3 bucket for storage
        data_source_id (str): Identifier for the data source
        sitemap_class_type (str): Type of sitemap reader to use ("GeneralSiteMapReader" or "GovtrackReader")
        knowledge_base_id (str): ID of the knowledge base to update
        kendra_index_id (str): ID of the Kendra index
        export_format (str): Format to export pages as ("pdf" or "html")
    """

    ui_color = "#A6E6A6"
    logger = logging.getLogger(__name__)

    @apply_defaults
    def __init__(
        self,
        sitemap_url,
        s3_prefix,
        s3_bucket_name,
        data_source_id,
        sitemap_class_type: Literal[
            "GeneralSiteMapReader", "GovtrackReader"
        ] = "GeneralSiteMapReader",
        knowledge_base_id=None,
        kendra_index_id=None,
        export_format="pdf",
        *args,
        **kwargs,
    ):
        super(WebsiteScraperOperator, self).__init__(
            data_source_id=data_source_id,
            knowledge_base_id=knowledge_base_id,
            kendra_index_id=kendra_index_id,
            *args,
            **kwargs,
        )
        self.sitemap_url = sitemap_url
        self.s3_prefix = s3_prefix
        self.s3_bucket_name = s3_bucket_name
        self.data_source_id = data_source_id
        self.knowledge_base_id = knowledge_base_id
        self.map_name = None
        self.export_format = export_format
        self.sitemap_class_type = sitemap_class_type



    def execute(self, context):
        """
        Handles the website scraping logic:
            - Initializes the appropriate sitemap reader
            - Parses the sitemap to extract URLs
            - Processes each URL by converting it to PDF or HTML and saving to S3
            - Handles checking for existing files to avoid duplicate processing
        """
        self.log.info(
            f"Starting website scraping operation - Type: {self.sitemap_class_type}"
        )
        self.log.info(
            f"Configuration: Sitemap URL: {self.sitemap_url}, Export Format: {self.export_format}"
        )

        # Initialize sitemap reader
        self.log.info(f"Initializing {self.sitemap_class_type}...")
        if self.sitemap_class_type == "GeneralSiteMapReader":
            self.sitemap_class = GeneralSiteMapReader(self.sitemap_url)
        elif self.sitemap_class_type == "GovtrackReader":
            self.sitemap_class = GovtrackReader(self.sitemap_url)

        self.log.info("Parsing sitemap...")
        try:
            self.sitemap_class.parse_sitemap()
            self.log.info(
                f"Sitemap parsed successfully - Found {len(self.sitemap_class.urls)} URLs"
            )
        except Exception as e:
            self.log.error(f"Failed to parse sitemap: {str(e)}")
            raise AirflowSkipException(f"Sitemap parsing failed: {str(e)}")

        # Update context with domain information
        context["params"]["map_name"] = self.sitemap_class.domain
        context["map_index_template"] = "{{ params.map_name }}"
        self.log.info(f"Set map name in context: {self.sitemap_class.domain}")

        # Get existing S3 files
        s3_keys = Util.get_existing_s3_files(
            self.s3_bucket_name,
            self.s3_prefix,
            full_key=True,
            file_type=f".{self.export_format}",
        )
        self.log.info(
            f"Found {len(s3_keys)} existing files in S3 bucket: {self.s3_bucket_name}"
        )

        processed_count = 0
        skipped_count = 0

        # Process URLs
        total_urls = len(self.sitemap_class.urls)
        for url in self.sitemap_class.urls:
            try:
                if self.export_format == "pdf":
                    site = WebsiteToPDf(url, self.s3_prefix, self.s3_bucket_name)
                elif self.export_format == "html":
                    site = WebsiteToHTML(url, self.s3_prefix, self.s3_bucket_name)
                else:
                    raise ValueError(f"Unsupported export format: {self.export_format}")

                key = f"{site.get_key()}.{self.export_format}"
                if key in s3_keys:
                    self.log.debug(
                        f"Skipping existing file: s3://{self.s3_bucket_name}/{key}"
                    )
                    skipped_count += 1
                    continue

                self.log.info(
                    f"Processing URL [{processed_count + 1}/{total_urls}]: {url}"
                )
                site.save()
                processed_count += 1

            except Exception as e:
                self.log.error(f"Failed to process URL {url}: {str(e)}")
                raise

        self.log.info(
            f"Website scraping complete - Processed: {processed_count}, Skipped: {skipped_count}"
        )
        self.log.info(
            "Starting the knowledge base update via parent class UpdateKnowledgeBaseOperator..."
        )
        super().execute(context)
        self.log.info("Knowledge base update completed successfully")