TypeError: cannot pickle '_thread.RLock' object

create_engine(...)

requests.Session()

boto3.client(...)

s3fs.S3FileSystem(...)

@flow

@flow

from functools import lru_cache
from prefect import flow, task
from sqlalchemy import create_engine

@lru_cache
def get_engine():
    # Created when first used in the subprocess
    return create_engine("postgresql+psycopg2://...")

@task
def do_db_work(...):
    engine = get_engine()
    with engine.connect() as conn:
        ...

@flow
def my_flow(...):
    do_db_work.submit(...)

requests.Session()

boto3.client(...)

get_client()

logging.getLogger(...)

get_run_logger()

import cloudpickle
cloudpickle.dumps(suspect_object)  # should raise if it's the culprit

@flow

import logging

import sentry_sdk
from prefect import task
from sqlalchemy import func
from sqlalchemy import select
from sqlmodel import Session

from app.aws.s3 import get_document_type
from app.aws.s3 import get_policy_document
from app.core.config import settings
from app.core.db import engine
from app.core.sentry import configure_sentry
from app.data_scripts.generate_policy_diff import compare_policies
from app.models import Policy
from app.models import PolicyDiff
from app.utils import get_policy_text_from_authed_s3_link

logging.basicConfig(level=<http://logging.INFO|logging.INFO>)
logger = logging.getLogger(__name__)
configure_sentry()


def query_new_policy_versions(session):
    """
    Returns a dict of parent_policy_id to their two most recent policy versions,
    for policies that need diffs created.

    Uses a window function to find policies with previous versions and checks
    if a diff already exists.
    """

    # Create CTE with window function to get previous policy ID
    # Partition by parent_policy_id, order by date_effective ASC (NULLS LAST to handle null dates)
    base_cte = (
        select(
            Policy.id,
            Policy.parent_policy_id,
            Policy.date_effective,
            Policy.date_expiration,
            func.lag(Policy.id)
            .over(
                partition_by=Policy.parent_policy_id,
                order_by=Policy.date_effective.asc().nulls_last(),
            )
            .label("previous_policy_id"),
        )
    ).cte("base")

    # Main query: select from base CTE, left join with PolicyDiff
    # Filter where previous_policy_id is not null and diff doesn't exist
    query = (
        select(base_cte)
        .outerjoin(
            PolicyDiff,
            PolicyDiff.current_policy_id == base_cte.c.id,
        )
        .where(base_cte.c.previous_policy_id.isnot(None))
        .where(PolicyDiff.current_policy_id.is_(None))
    )

    # Execute query and get policy IDs
    result = session.exec(query).all()
    policy_ids = [row.id for row in result]

    <http://logger.info|logger.info>("%d policies need diffs created", len(policy_ids))

    if not policy_ids:
        return {}

    # Fetch the actual Policy objects and their previous versions
    # Get current policies - use .scalars().all() to ensure we get Policy objects, not Row objects
    current_policies = session.exec(select(Policy).where(Policy.id.in_(policy_ids))).scalars().all()

    # Get previous policy IDs from the result
    policy_id_to_previous = {row.id: row.previous_policy_id for row in result}
    # Filter out None values - previous_policy_id can be None and we don't want to query for those
    previous_policy_ids = [
        policy_id for policy_id in set(policy_id_to_previous.values()) if policy_id is not None
    ]

    # Get previous policies
    if previous_policy_ids:
        previous_policies = (
            session.exec(select(Policy).where(Policy.id.in_(previous_policy_ids))).scalars().all()
        )
    else:
        previous_policies = []

    # Create a dict mapping previous_policy_id to Policy object
    previous_policy_map = {p.id: p for p in previous_policies}

    # Group policies by current_policy_id with current and previous versions
    # Use current_policy.id as key to avoid collisions when multiple policies
    # share the same parent_policy_id
    matched_policy_versions = {}

    for current_policy in current_policies:
        previous_policy_id = policy_id_to_previous[current_policy.id]
        previous_policy = previous_policy_map.get(previous_policy_id)

        if previous_policy:
            # Use current_policy.id as key to ensure uniqueness
            matched_policy_versions[current_policy.id] = {
                "current": current_policy.model_dump(),
                "previous": previous_policy.model_dump(),
            }

    <http://logger.info|logger.info>("%d policy pairs ready for diffing", len(matched_policy_versions.keys()))
    return matched_policy_versions


def fetch_authed_s3_urls(matched_policy_versions):
    """
    Fetches authed S3 URLs and document types for policy pairs. Removes policy pairs
    where URLs or document types cannot be retrieved (e.g., file doesn't exist).
    """
    policy_ids_to_remove = []
    for current_policy_id, policies in matched_policy_versions.items():
        for version_key in ["current", "previous"]:
            version = policies[version_key]
            authed_s3_url = get_policy_document(version["s3_file_name"])
            document_type = get_document_type(version["s3_file_name"])

            if authed_s3_url is None or document_type is None:
                logger.warning(
                    "Failed to get S3 URL or document type for policy_id %s, version %s",
                    current_policy_id,
                    version_key,
                )
                with sentry_sdk.push_scope() as scope:
                    scope.set_tag("task", "run_policy_diff_task")
                    scope.set_tag("environment", settings.ENVIRONMENT)
                    scope.set_tag("error_type", "s3_document_fetch_failure")
                    scope.set_extra("current_policy_id", str(current_policy_id))
                    scope.set_extra("s3_file_name", version["s3_file_name"])
                    scope.set_extra("authed_s3_url", authed_s3_url)
                    scope.set_extra("document_type", document_type)
                    sentry_sdk.capture_message(
                        (
                            f"S3 document fetch failed: Could not retrieve URL or "
                            f"document type for policy {current_policy_id} ({version_key})"
                        ),
                        level="warning",
                    )
                policy_ids_to_remove.append(current_policy_id)
                break

            version["authed_s3_url"] = authed_s3_url
            version["document_type"] = document_type

    # Remove policy pairs that couldn't get URLs/document types
    for policy_id in policy_ids_to_remove:
        matched_policy_versions.pop(policy_id, None)

    if policy_ids_to_remove:
        <http://logger.info|logger.info>(
            "Skipped %d policy pairs due to S3 URL/document type retrieval failures",
            len(policy_ids_to_remove),
        )

    return matched_policy_versions


def fetch_s3_documents(matched_policy_versions):
    """
    Fetches S3 documents for policy pairs. Removes policy pairs where documents
    cannot be fetched (e.g., 404 errors).
    """
    policy_ids_to_remove = []
    for current_policy_id, policies in matched_policy_versions.items():
        current_policy = policies["current"]
        previous_policy = policies["previous"]

        try:
            current_policy["document"] = get_policy_text_from_authed_s3_link(
                current_policy["authed_s3_url"],
                current_policy["document_type"],
                split_lines=True,
            )
        except Exception as e:
            logger.warning(
                "Failed to fetch current policy document for policy_id %s: %s",
                current_policy_id,
                str(e),
            )
            with sentry_sdk.push_scope() as scope:
                scope.set_tag("task", "run_policy_diff_task")
                scope.set_tag("environment", settings.ENVIRONMENT)
                scope.set_extra("current_policy_id", current_policy_id)
                sentry_sdk.capture_exception(e)
            policy_ids_to_remove.append(current_policy_id)
            continue

        try:
            previous_policy["document"] = get_policy_text_from_authed_s3_link(
                previous_policy["authed_s3_url"],
                previous_policy["document_type"],
                split_lines=True,
            )
        except Exception as e:
            logger.warning(
                "Failed to fetch previous policy document for policy_id %s: %s",
                current_policy_id,
                str(e),
            )
            with sentry_sdk.push_scope() as scope:
                scope.set_tag("task", "run_policy_diff_task")
                scope.set_tag("environment", settings.ENVIRONMENT)
                scope.set_extra("current_policy_id", current_policy_id)
                sentry_sdk.capture_exception(e)
            policy_ids_to_remove.append(current_policy_id)
            continue

    # Remove policy pairs that couldn't be fetched
    for policy_id in policy_ids_to_remove:
        matched_policy_versions.pop(policy_id, None)

    if policy_ids_to_remove:
        <http://logger.info|logger.info>(
            "Skipped %d policy pairs due to document fetch failures",
            len(policy_ids_to_remove),
        )

    return matched_policy_versions


def insert_policy_diff_into_db(session, policy_diff):
    try:
        # might need to allow for upsert in the future
        session.add(policy_diff)
        session.commit()
        return
    except Exception as e:
        current_policy_id = getattr(policy_diff, "current_policy_id", None)
        logger.exception("Error inserting policy diff into db for %s", current_policy_id)
        with sentry_sdk.push_scope() as scope:
            scope.set_tag("task", "run_policy_diff_task")
            scope.set_tag("environment", settings.ENVIRONMENT)
            scope.set_extra("current_policy_id", str(policy_diff.current_policy_id))
            scope.set_extra("previous_policy_id", str(policy_diff.previous_policy_id))
            sentry_sdk.capture_exception(e)
        return


@task
def run_policy_diff_task():
    """
    Diffs the two most recent policy versions for each parent policy that had at least one new
    version created today.

    This task:
    - Queries policies created today to identify affected parent_policy_ids.
    - Fetches the two most recent versions for each parent.
    - Extracts the policy text from their PDF S3 files.
    - Performs a diff using the extracted content.
    - Inserts the resulting diffs into the database.
    """

    with Session(engine) as session:
        matched_policy_versions = query_new_policy_versions(session)
        matched_policy_versions_with_s3_urls = fetch_authed_s3_urls(matched_policy_versions)
        matched_policy_versions_with_documents = fetch_s3_documents(
            matched_policy_versions_with_s3_urls
        )
        for current_policy_id, policies in matched_policy_versions_with_documents.items():
            <http://logger.info|logger.info>("Generating policy diff for policy_id %s", current_policy_id)
            print(current_policy_id)
            try:
                policy_diff = compare_policies(
                    policies["current"],
                    policies["previous"],
                )
                insert_policy_diff_into_db(session, policy_diff)
            except Exception as e:
                parent_policy_id = policies["current"].get("parent_policy_id")
                logger.exception(
                    "Error generating policy diff for current_policy_id %s", current_policy_id
                )
                with sentry_sdk.push_scope() as scope:
                    scope.set_tag("task", "run_policy_diff_task")
                    scope.set_tag("environment", settings.ENVIRONMENT)
                    scope.set_extra("current_policy_id", str(current_policy_id))
                    scope.set_extra("parent_policy_id", str(parent_policy_id))
                    sentry_sdk.capture_exception(e)


if __name__ == "__main__":
    run_policy_diff_task()