`src/db.js`

87 lines
import pg from 'pg';

// Connect to the realestate DB. Prefer DATABASE_URL; otherwise let libpq env vars
// (PGHOST/PGUSER/...) apply and just set the database name to REALESTATE_DB.
const { Pool } = pg;
export const pool = process.env.DATABASE_URL
  ? new Pool({ connectionString: process.env.DATABASE_URL })
  : new Pool({ database: process.env.REALESTATE_DB || 'realestate' });

export async function q(text, params) {
  const client = await pool.connect();
  try {
    return await client.query(text, params);
  } finally {
    client.release();
  }
}

// Generic chunked upsert.
//   table    — target table name
//   conflict — array of PK column names for ON CONFLICT
//   rows     — array of plain objects; keys are column names, one key may be `raw`
//              (stringified to jsonb). Columns are taken from the first row.
// Returns the number of rows sent.
export async function upsert(table, conflict, rows) {
  if (!rows.length) return 0;

  // Postgres rejects an INSERT ... ON CONFLICT that targets the same conflict key
  // twice in one statement ("cannot affect row a second time"). Some datasets
  // repeat a key within a single page (e.g. one code-enforcement apno or one film
  // permit_no spanning multiple rows), so collapse by conflict key first — last wins.
  if (rows.length > 1) {
    const seen = new Map();
    for (const r of rows) seen.set(conflict.map((k) => r[k]).join(''), r);
    if (seen.size !== rows.length) rows = [...seen.values()];
  }

  const cols = Object.keys(rows[0]);
  const updatable = cols.filter((c) => !conflict.includes(c));
  const setClause = updatable.length
    ? 'DO UPDATE SET ' + updatable.map((c) => `"${c}" = EXCLUDED."${c}"`).join(', ')
    : 'DO NOTHING';

  const JSONB = new Set(['raw', 'geom', 'attrs']); // columns stored as jsonb
  const CHUNK = 500;
  let sent = 0;
  for (let i = 0; i < rows.length; i += CHUNK) {
    const slice = rows.slice(i, i + CHUNK);
    const values = [];
    const tuples = slice.map((row) => {
      const ph = cols.map((c) => {
        let v = row[c];
        if (JSONB.has(c) && v != null && typeof v !== 'string') v = JSON.stringify(v);
        values.push(v === undefined ? null : v);
        return `$${values.length}${JSONB.has(c) ? '::jsonb' : ''}`;
      });
      return `(${ph.join(',')})`;
    });
    const sql =
      `INSERT INTO ${table} (${cols.map((c) => `"${c}"`).join(',')}) ` +
      `VALUES ${tuples.join(',')} ` +
      `ON CONFLICT (${conflict.map((c) => `"${c}"`).join(',')}) ${setClause}`;
    await q(sql, values);
    sent += slice.length;
  }
  return sent;
}

export async function getState(source) {
  const { rows } = await q('SELECT * FROM la_ingest_state WHERE source = $1', [source]);
  return rows[0] || null;
}

export async function setState(source, { dataset_id, last_cursor, rows_upserted, last_status }) {
  await q(
    `INSERT INTO la_ingest_state (source, dataset_id, last_cursor, last_run, rows_upserted, last_status)
     VALUES ($1,$2,$3, now(), $4, $5)
     ON CONFLICT (source) DO UPDATE SET
       dataset_id = EXCLUDED.dataset_id,
       last_cursor = COALESCE(EXCLUDED.last_cursor, la_ingest_state.last_cursor),
       last_run = EXCLUDED.last_run,
       rows_upserted = la_ingest_state.rows_upserted + EXCLUDED.rows_upserted,
       last_status = EXCLUDED.last_status`,
    [source, dataset_id, last_cursor, rows_upserted, last_status]
  );
}