main.py

import re
from requests import Session
from requests.exceptions import ConnectionError, Timeout
from lxml.html import fromstring
from time import sleep
import os
import psycopg2
from psycopg2.extras import execute_batch
from dotenv import load_dotenv
import click


def configure_session(
    session,
):
    params = {
        "_countries": [],
        "_grantingAuthorityRegions": [],
        "grantingAuthorityRegions": [],
        "resetSearch": "true",
        "_selectAll": [
            "",
            "",
        ],
        "countries": [
            "CountryAUT",
            "CountryBEL",
            "CountryBGR",
            "CountryHRV",
            "CountryCYP",
            "CountryCZE",
            "CountryDNK",
            "CountryEST",
            "CountryEIB",
            "CountryFIN",
            "CountryFRA",
            "CountryDEU",
            "CountryGRC",
            "CountryHUN",
            "CountryISL",
            "CountryIRL",
            "CountryITA",
            "CountryLVA",
            "CountryLTU",
            "CountryLUX",
            "CountryMLT",
            "CountryNLD",
            "CountryPOL",
            "CountryPRT",
            "CountryROM",
            "CountrySVK",
            "CountrySVN",
            "CountryESP",
            "CountrySWE",
            "CountryGBR",
        ],
        "lang": "en",
    }

    data = [
        ("resetSearch", "true"),
        ("_countries", ""),
        ("countries", "CountryAUT"),
        ("_countries", ""),
        ("countries", "CountryBEL"),
        ("_countries", ""),
        ("countries", "CountryBGR"),
        ("_countries", ""),
        ("countries", "CountryHRV"),
        ("_countries", ""),
        ("countries", "CountryCYP"),
        ("_countries", ""),
        ("countries", "CountryCZE"),
        ("_countries", ""),
        ("countries", "CountryDNK"),
        ("_countries", ""),
        ("countries", "CountryEST"),
        ("_countries", ""),
        ("countries", "CountryEIB"),
        ("_countries", ""),
        ("countries", "CountryFIN"),
        ("_countries", ""),
        ("countries", "CountryFRA"),
        ("_countries", ""),
        ("countries", "CountryDEU"),
        ("_countries", ""),
        ("countries", "CountryGRC"),
        ("_countries", ""),
        ("countries", "CountryHUN"),
        ("_countries", ""),
        ("countries", "CountryISL"),
        ("_countries", ""),
        ("countries", "CountryIRL"),
        ("_countries", ""),
        ("countries", "CountryITA"),
        ("_countries", ""),
        ("countries", "CountryLVA"),
        ("_countries", ""),
        ("countries", "CountryLTU"),
        ("_countries", ""),
        ("countries", "CountryLUX"),
        ("_countries", ""),
        ("countries", "CountryMLT"),
        ("_countries", ""),
        ("countries", "CountryNLD"),
        ("_countries", ""),
        ("countries", "CountryPOL"),
        ("_countries", ""),
        ("countries", "CountryPRT"),
        ("_countries", ""),
        ("countries", "CountryROM"),
        ("_countries", ""),
        ("countries", "CountrySVK"),
        ("_countries", ""),
        ("countries", "CountrySVN"),
        ("_countries", ""),
        ("countries", "CountryESP"),
        ("_countries", ""),
        ("countries", "CountrySWE"),
        ("_countries", ""),
        ("countries", "CountryGBR"),
    ]

    backoff = 30
    while True:
        try:
            response = session.get(
                "https://webgate.ec.europa.eu/competition/transparency/public",
                params=params,
                timeout=60,
            )

            cookie_header = response.headers.get("set-cookie")
            match = re.search("LB_TRANSPARENCY=([^;]+)", cookie_header) if cookie_header else None
            if match:
                session.cookies.set("LB_TRANSPARENCY", match.group(1))
            else:
                click.echo("No LB_TRANSPARENCY cookie found, retrying...")
                raise ConnectionError("Missing LB_TRANSPARENCY cookie")

            session.post(
                "https://webgate.ec.europa.eu/competition/transparency/public/search",
                data=data,
                timeout=60,
            )
            break

        except (ConnectionError, Timeout) as e:
            click.echo(f"configure_session failed: {e}")
            click.echo(f"Retrying in {backoff}s...")
            sleep(backoff)
            backoff = min(backoff * 2, 600)


def scrape_and_process(db_params):
    """Scrape data and directly process it without writing to disk."""
    data = {
        "resetSearch": "true",
        "countries": [
            "CountryAUT",
            "CountryBEL",
            "CountryBGR",
            "CountryHRV",
            "CountryCYP",
            "CountryCZE",
            "CountryDNK",
            "CountryEST",
            "CountryEIB",
            "CountryFIN",
            "CountryFRA",
            "CountryDEU",
            "CountryGRC",
            "CountryHUN",
            "CountryISL",
            "CountryIRL",
            "CountryITA",
            "CountryLVA",
            "CountryLTU",
            "CountryLUX",
            "CountryMLT",
            "CountryNLD",
            "CountryPOL",
            "CountryPRT",
            "CountryROM",
            "CountrySVK",
            "CountrySVN",
            "CountryESP",
            "CountrySWE",
            "CountryGBR",
        ],
        "grantingAuthorityRegions": [],
        "aidMeasureTitle": "",
        "aidMeasureCaseNumber": "",
        "refNo": "",
        "beneficiaryMs-input": "",
        "thirdPartyNonEuCountry-input": "",
        "beneficiaryNationalId": "",
        "beneficiaryName": "",
        "beneficiaryTypes": "",
        "regions-input": "",
        "sectors-input": "",
        "aidInstruments": "",
        "objectives": "",
        "nominalAmountFrom": "",
        "nominalAmountTo": "",
        "grantedAmountFrom": "",
        "grantedAmountTo": "",
        "dateGrantedFrom": "",
        "dateGrantedTo": "",
        "grantingAuthorityNames": "",
        "entrustedEntities": "",
        "financialIntermediaries": "",
        "offset": int(os.getenv("START_PAGE", "0")) * 100,
        "max": 100,
    }

    if data["offset"] > 0:
        click.echo(f"Starting from page {data['offset'] // 100} (offset {data['offset']})")

    session = Session()

    # Connect to the database
    conn, cursor = setup_database(db_params)

    # Get initial row count
    cursor.execute("SELECT COUNT(*) FROM state_aid_awards")
    initial_row_count = cursor.fetchone()[0]

    processed_rows = 0
    total_pages = 0

    while True:
        click.echo(f"Processing offset {data['offset']}...")

        backoff = 30
        while True:
            try:
                response = session.post(
                    "https://webgate.ec.europa.eu/competition/transparency/public/search/results?sort=serverReference&order=desc",
                    data=data,
                    timeout=60,
                )
                break
            except (ConnectionError, Timeout) as e:
                click.echo(f"Request failed: {e}")
                click.echo(f"Retrying in {backoff}s...")
                sleep(backoff)
                backoff = min(backoff * 2, 600)
                click.echo("Reconfiguring session...")
                configure_session(session)

        if (
            "Please choose a language" in response.text
            or "currently unable to handle the request" in response.text
        ):
            click.echo("Cookie expired or overload, retrying configuration...")
            configure_session(session)
        else:
            # Parse the HTML content directly
            tree = fromstring(response.text)

            # Extract data from the HTML
            data_rows = extract_data_from_html_content(tree)
            processed_rows += len(data_rows)
            total_pages += 1

            # Insert into database
            insert_data(conn, cursor, data_rows, f"page_{data['offset']}")

            data["offset"] += 100

            # Get max offset to stop at some point
            max = int(tree.findall(".//a[@class='step']")[-1].text) * 100

            if data["offset"] > max:
                break

    # Get final row count to determine how many new rows were inserted
    cursor.execute("SELECT COUNT(*) FROM state_aid_awards")
    final_row_count = cursor.fetchone()[0]
    new_rows_inserted = final_row_count - initial_row_count

    cursor.close()
    conn.close()

    click.echo(
        f"Completed processing {total_pages} pages with {processed_rows} total rows."
    )
    click.echo(
        f"Inserted {new_rows_inserted} new records, skipped {processed_rows - new_rows_inserted} duplicates."
    )


def clean_text(text):
    """Remove excess whitespace and newlines from text."""
    if text is None:
        return None
    return re.sub(r"\s+", " ", text).strip()


def extract_data_from_html_content(tree):
    """Extract data from an HTML tree containing state aid transparency data."""
    # Find the data table using lxml
    data_table = tree.xpath("//table[@id='resultsTable']")
    if not data_table:
        click.echo("No data table found in HTML content")
        return []

    # Extract headers using lxml
    headers = []
    header_elements = tree.xpath("//table[@id='resultsTable']/thead/tr/th")
    for th in header_elements:
        # Get the text inside the <a> tag if it exists, otherwise get the th text
        a_tag = th.xpath("./a")
        if a_tag:
            header = clean_text(a_tag[0].text_content())
        else:
            header = clean_text(th.text_content())
        headers.append(header)

    # Extract rows using lxml
    rows = []
    row_elements = tree.xpath("//table[@id='resultsTable']/tbody/tr")
    for tr in row_elements:
        row = {}
        cells = tr.xpath("./td")

        for i, cell in enumerate(cells):
            if i < len(headers):
                header = headers[i]

                # Check if the cell has a title attribute (for truncated values)
                title = cell.get("title")
                if title:
                    value = clean_text(title)
                else:
                    value = clean_text(cell.text_content())

                # Special handling for National ID column
                if header == "National ID" and title:
                    row[header] = clean_text(cell.text_content())
                    row[header + " Type"] = clean_text(title)
                # Special handling for amount values
                elif (
                    "Amount" in header
                    or header == "Aid element, expressed as full amount"
                ):
                    # Try to extract the numeric part and currency
                    amount_match = (
                        re.match(r"([\d,.]+)\s*([A-Z]{3})?", value) if value else None
                    )
                    if amount_match:
                        amount = amount_match.group(1)
                        # Remove commas and convert to integer
                        amount = int(float(amount.replace(",", "")))
                        row[header] = amount
                    else:
                        row[header] = None
                else:
                    row[header] = value

        if row:  # Only add non-empty rows
            rows.append(row)

    return rows


def setup_database(db_params):
    """Set up the PostgreSQL database and table."""
    conn = psycopg2.connect(**db_params)
    cursor = conn.cursor()

    # Create table if it doesn't exist
    cursor.execute(
        """
    CREATE TABLE IF NOT EXISTS state_aid_awards (
        id SERIAL PRIMARY KEY,
        country TEXT,
        aid_measure_title TEXT,
        sa_number TEXT,
        ref_no TEXT,
        national_id TEXT,
        national_id_type TEXT,
        beneficiary_name TEXT,
        beneficiary_type TEXT,
        region TEXT,
        sector TEXT,
        aid_instrument TEXT,
        aid_objectives TEXT,
        nominal_amount NUMERIC,
        aid_element NUMERIC,
        date_granted DATE,
        granting_authority_name TEXT,
        entrusted_entity TEXT,
        financial_intermediaries TEXT,
        published_date DATE,
        beneficiary_ms TEXT,
        third_party_non_eu_country TEXT,
        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
        UNIQUE (sa_number, ref_no, national_id, beneficiary_name, date_granted)
    );
    """
    )

    conn.commit()
    return conn, cursor


def insert_data(conn, cursor, data_rows, file_name):
    """Insert extracted data into the database. Returns number of newly inserted rows."""
    if not data_rows:
        return 0

    # Map the column names from HTML to database columns
    column_mapping = {
        "Country": "country",
        "Aid Measure Title": "aid_measure_title",
        "SA.Number": "sa_number",
        "Ref-no.": "ref_no",
        "National ID": "national_id",
        "National ID Type": "national_id_type",
        "Name of the beneficiary": "beneficiary_name",
        "Beneficiary Type": "beneficiary_type",
        "Region": "region",
        "Sector (NACE)": "sector",
        "Aid Instrument": "aid_instrument",
        "Objectives of the Aid": "aid_objectives",
        "Nominal Amount, expressed as full amount": "nominal_amount",
        "Aid element, expressed as full amount": "aid_element",
        "Date of granting": "date_granted",
        "Granting Authority Name": "granting_authority_name",
        "Entrusted Entity": "entrusted_entity",
        "Financial Intermediaries": "financial_intermediaries",
        "Published Date": "published_date",
        "Another Beneficiary Member State": "beneficiary_ms",
        "Third country outside of the EU": "third_party_non_eu_country",
    }

    # Prepare data for insertion
    insert_data = []
    for row in data_rows:
        db_row = {}

        for html_col, db_col in column_mapping.items():
            if html_col in row:
                value = row[html_col]

                # Convert dates to proper format
                if db_col in ["date_granted", "published_date"] and value:
                    # Try to parse date in format DD/MM/YYYY
                    try:
                        parts = value.split("/")
                        if len(parts) == 3:
                            value = f"{parts[2]}-{parts[1]}-{parts[0]}"  # Convert to YYYY-MM-DD
                    except:
                        pass  # Keep original value if parsing fails

                db_row[db_col] = value

        # Ensure key fields are not NULL - use placeholder values for missing data
        # These placeholders help maintain uniqueness while preventing NULL values in the key
        if not db_row.get("sa_number"):
            db_row["sa_number"] = "UNKNOWN"
        if not db_row.get("ref_no"):
            db_row["ref_no"] = "UNKNOWN"
        if not db_row.get("national_id"):
            db_row["national_id"] = "UNKNOWN"
        if not db_row.get("beneficiary_name"):
            db_row["beneficiary_name"] = "UNKNOWN"
        if not db_row.get("date_granted"):
            db_row["date_granted"] = "1900-01-01"  # Default date for missing dates

        insert_data.append(db_row)

    # Prepare SQL statement with all possible columns
    columns = list(column_mapping.values())
    placeholders = [f"%({col})s" for col in columns]

    sql = f"""
    INSERT INTO state_aid_awards ({', '.join(columns)})
    VALUES ({', '.join(placeholders)})
    ON CONFLICT (sa_number, ref_no, national_id, beneficiary_name, date_granted) DO NOTHING
    """

    # Count rows before insert
    cursor.execute("SELECT COUNT(*) FROM state_aid_awards")
    rows_before = cursor.fetchone()[0]

    # Execute batch insert
    execute_batch(cursor, sql, insert_data, page_size=100)
    conn.commit()

    # Count rows after insert
    cursor.execute("SELECT COUNT(*) FROM state_aid_awards")
    rows_after = cursor.fetchone()[0]

    new_rows = rows_after - rows_before
    print(f"Processed {len(insert_data)} rows from {file_name}: {new_rows} new, {len(insert_data) - new_rows} duplicates")

    return new_rows


@click.group()
def cli():
    """CLI for EU State Aid Transparency Register."""
    # Load environment variables
    load_dotenv()


def get_db_params():
    """Get database connection parameters from environment variables."""
    return {
        "dbname": os.getenv("DB_NAME", "state_aid_db"),
        "user": os.getenv("DB_USER", "postgres"),
        "password": os.getenv("DB_PASSWORD", ""),
        "host": os.getenv("DB_HOST", "localhost"),
        "port": os.getenv("DB_PORT", "5432"),
    }


@cli.command()
def run():
    """Scrape data from the EU State Aid Transparency Register and directly import to the database."""
    db_params = get_db_params()

    click.echo("Fetching and importing data directly to database...")
    scrape_and_process(db_params)
    click.echo("Process completed!")


if __name__ == "__main__":
    cli()