Data Warehouse Integration for AI#

import snowflake.connector
from contextlib import contextmanager
from typing import Iterator
import pandas as pd

# Educational example — not for clinical use

@contextmanager
def snowflake_connection(
    account: str,
    user: str,
    private_key_path: str,      # Key-pair auth (no passwords in config)
    warehouse: str,
    database: str,
    schema: str,
    role: str,
) -> Iterator[snowflake.connector.SnowflakeConnection]:
    """
    Create a Snowflake connection with key-pair authentication.
    
    Key-pair auth is preferred over password auth for service accounts
    because private keys can be rotated without embedding credentials
    in configuration.
    
    Educational example — not for clinical use.
    """
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.backends import default_backend
    
    with open(private_key_path, "rb") as key_file:
        private_key = serialization.load_pem_private_key(
            key_file.read(),
            password=None,
            backend=default_backend()
        )
    
    private_key_bytes = private_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    
    conn = snowflake.connector.connect(
        account=account,
        user=user,
        private_key=private_key_bytes,
        warehouse=warehouse,
        database=database,
        schema=schema,
        role=role,
    )
    
    try:
        yield conn
    finally:
        conn.close()


def extract_population_for_risk_scoring(
    conn: snowflake.connector.SnowflakeConnection,
    target_date: str,
    risk_model_id: str,
) -> pd.DataFrame:
    """
    Extract patient population for readmission risk scoring.
    
    Uses parameterized queries to prevent SQL injection.
    Applies ROW ACCESS POLICY on the underlying table to enforce
    data access restrictions without application-layer filtering.
    
    Educational example — not for clinical use.
    """
    query = """
        SELECT
            p.patient_id,
            p.age_at_encounter,
            p.primary_diagnosis_code,
            p.comorbidity_count,
            p.prior_admissions_12_months,
            p.length_of_stay_days,
            p.discharge_disposition_code,
            -- De-identification flag: verify with privacy officer before use
            p.is_deidentified
        FROM patient_encounter_facts p
        WHERE
            p.discharge_date = %(target_date)s
            AND p.encounter_class = 'inpatient'
            AND p.is_deidentified = TRUE   -- Only de-identified records for AI batch scoring
        ORDER BY p.patient_id
    """
    
    cursor = conn.cursor()
    cursor.execute(query, {"target_date": target_date})
    
    columns = [desc[0].lower() for desc in cursor.description]
    rows = cursor.fetchall()
    
    return pd.DataFrame(rows, columns=columns)


def write_risk_scores_to_snowflake(
    conn: snowflake.connector.SnowflakeConnection,
    risk_scores: list[dict],
    model_version: str,
) -> int:
    """
    Write AI-generated risk scores back to Snowflake for downstream use.
    Includes model_version for reproducibility and audit.
    
    Educational example — not for clinical use.
    """
    cursor = conn.cursor()
    
    # Use staging table + MERGE to ensure idempotency
    cursor.execute("CREATE TEMP TABLE IF NOT EXISTS risk_score_staging (LIKE ai_risk_scores)")
    
    cursor.executemany(
        """
        INSERT INTO risk_score_staging 
            (patient_id, encounter_id, risk_score, risk_tier, model_version, scored_at)
        VALUES (%(patient_id)s, %(encounter_id)s, %(risk_score)s, %(risk_tier)s, %(model_version)s, CURRENT_TIMESTAMP())
        """,
        [{**score, "model_version": model_version} for score in risk_scores]
    )
    
    cursor.execute("""
        MERGE INTO ai_risk_scores AS target
        USING risk_score_staging AS source
            ON target.patient_id = source.patient_id AND target.encounter_id = source.encounter_id
        WHEN MATCHED THEN UPDATE SET
            risk_score = source.risk_score,
            risk_tier = source.risk_tier,
            model_version = source.model_version,
            scored_at = source.scored_at
        WHEN NOT MATCHED THEN INSERT
            (patient_id, encounter_id, risk_score, risk_tier, model_version, scored_at)
            VALUES (source.patient_id, source.encounter_id, source.risk_score, 
                    source.risk_tier, source.model_version, source.scored_at)
    """)
    
    conn.commit()
    return len(risk_scores)

from google.cloud import bigquery
from typing import Iterator

# Educational example — not for clinical use

def stream_patient_cohort_for_quality_measures(
    project_id: str,
    dataset_id: str,
    measurement_year: int,
) -> Iterator[dict]:
    """
    Stream patient cohort from BigQuery for HEDIS quality measure AI extraction.
    
    Uses streaming reads (client.query iterator) to avoid loading the
    entire cohort into memory — appropriate for large population queries.
    
    Educational example — not for clinical use.
    """
    client = bigquery.Client(project=project_id)
    
    # Parameterized query via BigQuery query parameters
    query = """
        SELECT
            patient_key,
            age_in_years,
            gender_code,
            primary_care_provider_npi,
            continuous_enrollment_flag,
            diabetes_flag,
            hypertension_flag,
            -- Exclude fields that could re-identify
            -- even in this "de-identified" dataset
        FROM `{project}.{dataset}.member_eligibility`
        WHERE
            measurement_year = @measurement_year
            AND continuous_enrollment_flag = TRUE
        ORDER BY patient_key
    """.format(project=project_id, dataset=dataset_id)
    
    job_config = bigquery.QueryJobConfig(
        query_parameters=[
            bigquery.ScalarQueryParameter("measurement_year", "INT64", measurement_year),
        ]
    )
    
    query_job = client.query(query, job_config=job_config)
    
    for row in query_job:
        yield dict(row)


def use_bigquery_ml_for_classification(
    project_id: str,
    dataset_id: str,
    patient_features_table: str,
    model_name: str,
) -> str:
    """
    Use BigQuery ML for in-warehouse inference.
    
    Avoids moving patient data out of BigQuery for scoring —
    the model runs inside BigQuery, results land in a BigQuery table.
    This is the preferred pattern when data cannot leave the warehouse
    for compliance reasons.
    
    Educational example — not for clinical use.
    """
    client = bigquery.Client(project=project_id)
    
    query = f"""
        CREATE OR REPLACE TABLE `{project_id}.{dataset_id}.ai_predictions` AS
        SELECT
            patient_key,
            predicted_readmission_probability,
            predicted_label,
            CURRENT_TIMESTAMP() AS prediction_timestamp,
            '{model_name}' AS model_name
        FROM ML.PREDICT(
            MODEL `{project_id}.{dataset_id}.{model_name}`,
            TABLE `{project_id}.{dataset_id}.{patient_features_table}`
        )
    """
    
    job = client.query(query)
    job.result()    # Wait for completion
    
    return f"{project_id}.{dataset_id}.ai_predictions"