Upload database.py

database.py

@@ -0,0 +1,432 @@
+import os
+import re
+from datetime import datetime, timedelta, timezone
+from typing import List, Dict
+
+import psycopg2
+from psycopg2.extras import Json
+
+# ── Configuration ────────────────────────────────────────────────────────────
+
+# Pull DATABASE_URL from the environment, falling back to a local .env file
+# so the same Postgres credentials work for both this scraper and the
+# Next.js / Drizzle side without having to re-export the variable.
+try:
+    from dotenv import load_dotenv
+
+    _here = os.path.dirname(os.path.abspath(__file__))
+    for _candidate in (
+        os.path.join(_here, ".env"),
+        os.path.join(_here, "components", ".env"),
+    ):
+        if os.path.exists(_candidate):
+            load_dotenv(_candidate, override=False)
+except ImportError:
+    pass
+
+DATABASE_URL = os.environ.get("DATABASE_URL")
+if not DATABASE_URL:
+    raise RuntimeError(
+        "DATABASE_URL is not set. Define it in your environment or in a "
+        ".env file at the project root (or in components/.env)."
+    )
+
+
+def _connect():
+    """Open a fresh Postgres connection.
+
+    We open/close per call instead of holding a global. With Neon's pooler
+    endpoint (the one in DATABASE_URL) each `connect` is effectively a
+    cheap checkout from their server-side pool, so this stays simple and
+    plays nicely with FastAPI's threadpool — no shared cursors between
+    request handlers.
+    """
+    return psycopg2.connect(DATABASE_URL)
+
+
+# ── Schema bootstrap ─────────────────────────────────────────────────────────
+
+
+def init_database():
+    """Initialize the database with required tables.
+
+    Idempotent: every statement uses IF NOT EXISTS, so it's safe to run on
+    every import and it won't fight with Drizzle's migrations if they've
+    already created the table.
+    """
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                CREATE TABLE IF NOT EXISTS scraped_jobs (
+                    id              SERIAL PRIMARY KEY,
+                    title           TEXT,
+                    company         TEXT,
+                    location        TEXT,
+                    link            TEXT UNIQUE,
+                    posted          TEXT,
+                    posted_datetime TIMESTAMP,
+                    programs        JSONB DEFAULT '[]'::jsonb,
+                    schools         JSONB DEFAULT '[]'::jsonb,
+                    scraped_at      TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    keywords        TEXT,
+                    freshness       TEXT,
+                    work_types      JSONB DEFAULT '[]'::jsonb,
+                    job_types       JSONB DEFAULT '[]'::jsonb,
+                    company_rating  REAL,
+                    contact_details JSONB DEFAULT '[]'::jsonb
+                )
+                """
+            )
+            # Migrate existing tables that don't yet have company_rating or contact_details
+            cur.execute(
+                """
+                ALTER TABLE scraped_jobs
+                    ADD COLUMN IF NOT EXISTS company_rating REAL
+                """
+            )
+            cur.execute(
+                """
+                ALTER TABLE scraped_jobs
+                    ADD COLUMN IF NOT EXISTS contact_details JSONB DEFAULT '[]'::jsonb
+                """
+            )
+            cur.execute(
+                "CREATE INDEX IF NOT EXISTS idx_scraped_at "
+                "ON scraped_jobs(scraped_at)"
+            )
+            cur.execute(
+                "CREATE INDEX IF NOT EXISTS idx_link ON scraped_jobs(link)"
+            )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+# ── Internal helpers ─────────────────────────────────────────────────────────
+
+
+def _coerce_posted_datetime(posted: str, raw_posted_dt, now_dt: datetime) -> datetime:
+    """Normalize the scraper's posted_datetime into a real `datetime`.
+
+    The scraper sometimes only gives us a relative label like "3 hours
+    ago" or a bare date string. SQLite stored these as TEXT, but Postgres
+    has a real TIMESTAMP column, so we resolve everything to a `datetime`
+    here before insertion. Falls back to `now_dt` when neither source is
+    usable. Behavior matches the original SQLite version.
+    """
+    needs_recompute = (
+        not raw_posted_dt
+        or (isinstance(raw_posted_dt, str) and len(raw_posted_dt) <= 10)
+    )
+
+    if posted and needs_recompute:
+        posted_lower = posted.lower()
+        match = re.search(r"(\d+)", posted_lower)
+        if match:
+            num = int(match.group(1))
+            if "minute" in posted_lower:
+                return now_dt - timedelta(minutes=num)
+            if "hour" in posted_lower:
+                return now_dt - timedelta(hours=num)
+            if "day" in posted_lower:
+                return now_dt - timedelta(days=num)
+            if "week" in posted_lower:
+                return now_dt - timedelta(weeks=num)
+            if "month" in posted_lower:
+                return now_dt - timedelta(days=num * 30)
+            return now_dt
+
+    if isinstance(raw_posted_dt, datetime):
+        return raw_posted_dt
+    if isinstance(raw_posted_dt, str):
+        try:
+            return datetime.fromisoformat(raw_posted_dt)
+        except ValueError:
+            pass
+
+    return now_dt
+
+
+def _row_to_job(row) -> Dict:
+    """Translate a SELECT row tuple into the dict shape main.py / the UI expect.
+
+    Notes:
+      * `programs` / `schools` come back as native Python lists because
+        psycopg2 auto-parses JSONB — no `json.loads` needed.
+      * `posted_datetime` / `scraped_at` come back as `datetime` objects;
+        we ISO-format them so the result survives `json.dumps` in FastAPI
+        without extra encoders.
+    """
+    return {
+        "title": row[0],
+        "company": row[1],
+        "location": row[2],
+        "link": row[3],
+        "posted": row[4],
+        "posted_datetime": row[5].replace(tzinfo=timezone.utc).isoformat() if row[5] else None,
+        "programs": row[6] if row[6] is not None else [],
+        "schools": row[7] if row[7] is not None else [],
+        "scraped_at": row[8].replace(tzinfo=timezone.utc).isoformat() if row[8] else None,
+        "keywords": row[9],
+        "company_rating": row[10],  # float or None
+        "contact_details": row[11] if len(row) > 11 and row[11] is not None else [],
+    }
+
+
+# ── Public API (signatures preserved from the SQLite version) ────────────────
+
+
+def save_scraped_jobs(jobs: List[Dict], search_params: Dict):
+    """Save scraped jobs to database."""
+    conn = _connect()
+    saved_count = 0
+    try:
+        with conn.cursor() as cur:
+            now_dt = datetime.now(timezone.utc)
+
+            for job in jobs:
+                posted = job.get("posted", "")
+                posted_dt = _coerce_posted_datetime(
+                    posted, job.get("posted_datetime"), now_dt
+                )
+
+                cur.execute(
+                    """
+                    INSERT INTO scraped_jobs (
+                        title, company, location, link, posted,
+                        posted_datetime, programs, schools, keywords,
+                        freshness, work_types, job_types, scraped_at, contact_details
+                    )
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                    ON CONFLICT (link) DO UPDATE SET
+                        title           = EXCLUDED.title,
+                        company         = EXCLUDED.company,
+                        location        = EXCLUDED.location,
+                        posted          = EXCLUDED.posted,
+                        posted_datetime = EXCLUDED.posted_datetime,
+                        programs        = EXCLUDED.programs,
+                        schools         = EXCLUDED.schools,
+                        keywords        = EXCLUDED.keywords,
+                        freshness       = EXCLUDED.freshness,
+                        work_types      = EXCLUDED.work_types,
+                        job_types       = EXCLUDED.job_types,
+                        scraped_at      = EXCLUDED.scraped_at,
+                        contact_details = EXCLUDED.contact_details
+                    """,
+                    (
+                        job.get("title"),
+                        job.get("company"),
+                        job.get("location"),
+                        job.get("link"),
+                        posted,
+                        posted_dt,
+                        Json(job.get("programs", []) or []),
+                        Json(job.get("schools", []) or []),
+                        search_params.get("keywords"),
+                        search_params.get("freshness"),
+                        Json(search_params.get("work_types", []) or []),
+                        Json(search_params.get("job_types", []) or []),
+                        now_dt,
+                        Json(job.get("contact_details", []) or []),
+                    ),
+                )
+                saved_count += 1
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        conn.close()
+    return saved_count
+
+
+def get_existing_jobs_contacts(links: List[str]) -> Dict[str, list]:
+    """Get existing contact_details for jobs already in the database."""
+    if not links:
+        return {}
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            cur.execute(
+                "SELECT link, contact_details FROM scraped_jobs WHERE link = ANY(%s)",
+                (links,)
+            )
+            rows = cur.fetchall()
+            return {row[0]: row[1] if row[1] is not None else [] for row in rows}
+    finally:
+        conn.close()
+
+
+def get_companies_with_ratings(companies: List[str]) -> List[str]:
+    """Get companies that already have a company_rating in the database."""
+    if not companies:
+        return []
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            lower_companies = [c.lower() for c in companies]
+            cur.execute(
+                "SELECT DISTINCT company FROM scraped_jobs WHERE LOWER(company) = ANY(%s) AND company_rating IS NOT NULL",
+                (lower_companies,)
+            )
+            rows = cur.fetchall()
+            return [row[0].lower() for row in rows]
+    finally:
+        conn.close()
+
+
+def get_jobs_by_timeframe(hours_ago: int) -> List[Dict]:
+    """Get jobs scraped within the specified timeframe."""
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            time_threshold = datetime.now(timezone.utc) - timedelta(hours=hours_ago)
+            cur.execute(
+                """
+                SELECT title, company, location, link, posted, posted_datetime,
+                       programs, schools, scraped_at, keywords, company_rating, contact_details
+                FROM scraped_jobs
+                WHERE posted_datetime >= %s
+                ORDER BY posted_datetime DESC
+                """,
+                (time_threshold,),
+            )
+            rows = cur.fetchall()
+    finally:
+        conn.close()
+    return [_row_to_job(row) for row in rows]
+
+
+def get_jobs_in_timeframe(start_hours: int, end_hours: int) -> List[Dict]:
+    """Get jobs scraped between start_hours and end_hours ago."""
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            now = datetime.now(timezone.utc)
+            start_time = now - timedelta(hours=start_hours)
+            end_time = now - timedelta(hours=end_hours)
+            cur.execute(
+                """
+                SELECT title, company, location, link, posted, posted_datetime,
+                       programs, schools, scraped_at, keywords, company_rating, contact_details
+                FROM scraped_jobs
+                WHERE posted_datetime >= %s AND posted_datetime <= %s
+                ORDER BY posted_datetime DESC
+                """,
+                (start_time, end_time),
+            )
+            rows = cur.fetchall()
+    finally:
+        conn.close()
+    return [_row_to_job(row) for row in rows]
+
+
+def cleanup_old_jobs(days_old: int = 7):
+    """Remove jobs older than specified days."""
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            cutoff_time = datetime.now(timezone.utc) - timedelta(days=days_old)
+            cur.execute(
+                "DELETE FROM scraped_jobs WHERE posted_datetime < %s",
+                (cutoff_time,),
+            )
+            deleted_count = cur.rowcount
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        conn.close()
+    return deleted_count
+
+
+def get_binned_jobs() -> Dict[str, Dict]:
+    """Fetch jobs from the last 30 days in one query and bin them by timeframe."""
+    conn = _connect()
+    try:
+        with conn.cursor() as cur:
+            now = datetime.now(timezone.utc)
+            start_time = now - timedelta(hours=720)
+            cur.execute(
+                """
+                SELECT title, company, location, link, posted, posted_datetime,
+                       programs, schools, scraped_at, keywords, company_rating, contact_details
+                FROM scraped_jobs
+                WHERE posted_datetime >= %s
+                ORDER BY posted_datetime DESC
+                """,
+                (start_time,),
+            )
+            rows = cur.fetchall()
+
+            t_1, t_2, t_5, t_24, t_168, t_720 = [], [], [], [], [], []
+            for row in rows:
+                dt = row[5]
+                job = _row_to_job(row)
+                if not dt:
+                    continue
+
+                dt_aware = dt.replace(tzinfo=timezone.utc)
+                hours = (now - dt_aware).total_seconds() / 3600
+                if hours <= 1:
+                    t_1.append(job)
+                elif hours <= 2:
+                    t_2.append(job)
+                elif hours <= 5:
+                    t_5.append(job)
+                elif hours <= 24:
+                    t_24.append(job)
+                elif hours <= 168:
+                    t_168.append(job)
+                elif hours <= 720:
+                    t_720.append(job)
+
+    finally:
+        conn.close()
+
+    return {
+        "1_hour": {"jobs": t_1, "count": len(t_1)},
+        "2_hours": {"jobs": t_2, "count": len(t_2)},
+        "5_hours": {"jobs": t_5, "count": len(t_5)},
+        "24_hours": {"jobs": t_24, "count": len(t_24)},
+        "1_week": {"jobs": t_168, "count": len(t_168)},
+        "1_month": {"jobs": t_720, "count": len(t_720)},
+    }
+
+
+init_database()
+
+
+def update_company_ratings(ratings: Dict[str, float]) -> int:
+    """Bulk-update company_rating on all scraped_jobs rows by company name.
+
+    Uses a single multi-row UPDATE so the number of DB round-trips equals one
+    per unique company (not one per job row). Matching is case-insensitive.
+    Returns the total number of rows updated.
+    """
+    if not ratings:
+        return 0
+    conn = _connect()
+    updated = 0
+    try:
+        with conn.cursor() as cur:
+            for company, rating in ratings.items():
+                cur.execute(
+                    """
+                    UPDATE scraped_jobs
+                       SET company_rating = %s
+                     WHERE LOWER(company) = LOWER(%s)
+                    """,
+                    (float(rating), company),
+                )
+                updated += cur.rowcount
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        conn.close()
+    return updated