sqlite2postgres.py

#!/usr/bin/env python3
"""
Migrate OpenModelicaLibraryTesting sqlite3 result databases into PostgreSQL.

The PostgreSQL layout is a 1:1 mirror of the sqlite3 one: one table per branch
with the same columns, plus the [omcversion] and [libversion] lookup tables.
Only the types are adapted (integer -> bigint, real -> double precision).

Only the python standard library is used; the psql client binary does the
talking, so nothing has to be installed on the test machines.  The password is
never passed on the command line - use PGPASSFILE or PGPASSWORD, as understood
by psql.

Typical use:

  export PGPASSFILE=~/.pgpass
  ./sqlite2postgres.py --sqlite dbs/ripper1.db --source ripper1
  ./sqlite2postgres.py --index
  ./sqlite2postgres.py --sqlite dbs/ripper2.db --source ripper2 --skip-existing
  ./sqlite2postgres.py --index

The machines share the tables of the branches they both test, so they are merged
in that order: ripper1 first, then ripper2 with --skip-existing, which keeps the
row already in the database whenever a key collides.  --index in between creates
the unique keys that decide what a collision is - (date, libname, model) for a
branch table - so it has to run before the second machine.

--source only names the machine for the resume bookkeeping; it is not stored in
the data.  Use --pgschema ripper1 to keep a machine in a schema of its own
instead of merging.

The migration is resumable and safe to re-run: how far each table has been
copied is recorded in [migration_progress], and a restart continues from the
last sqlite rowid that made it in.
"""

import argparse
import io
import os
import sqlite3
import subprocess
import sys
import time

# Columns of a per-branch table, in the order used by test.py.  Databases with
# PRAGMA user_version < 3 have no "parsing" column.
#
# The sqlite3 tables declare every column NOT NULL, but tables created before
# that declaration still hold NULLs (libversion.libversion has some), so only
# the key columns are NOT NULL here.
BRANCH_COLUMNS = [
    ("date", "bigint"),
    ("libname", "text"),
    ("model", "text"),
    ("exectime", "double precision"),
    ("frontend", "double precision"),
    ("backend", "double precision"),
    ("simcode", "double precision"),
    ("templates", "double precision"),
    ("compile", "double precision"),
    ("simulate", "double precision"),
    ("verify", "double precision"),
    ("verifyfail", "integer"),
    ("verifytotal", "integer"),
    ("finalphase", "integer"),
    ("parsing", "double precision"),
]

LOOKUP_COLUMNS = {
    "omcversion": [("date", "bigint"), ("branch", "text"), ("omcversion", "text")],
    "libversion": [("date", "bigint"), ("branch", "text"), ("libname", "text"),
                   ("libversion", "text"), ("confighash", "bigint")],
}

# Derived data, no longer generated by all-plots.py; not worth migrating.
SKIP_PREFIXES = ("datelookup_", "sqlite_")


# What identifies a row.  A test run writes one row per model, so a model
# appears once per (date, libname); the lookup tables have one row per run and
# per (run, library).  --index enforces these, which is what lets a second test
# machine be merged into a table that already holds another machine's results.
KEYS = {
    "omcversion": ["date", "branch"],
    "libversion": ["date", "branch", "libname", "confighash"],
}
BRANCH_KEY = ["date", "libname", "model"]

PROGRESS_TABLE = """
CREATE TABLE IF NOT EXISTS "migration_progress" (
  source     text    NOT NULL,
  tbl        text    NOT NULL,
  last_rowid bigint  NOT NULL,
  rows_read  bigint  NOT NULL,
  done       boolean NOT NULL DEFAULT false,
  updated_at timestamptz NOT NULL DEFAULT now(),
  PRIMARY KEY (source, tbl)
)
"""


def lit(s):
  """Quote a python value as an SQL string literal."""
  return "'" + str(s).replace("'", "''") + "'"


def ident(s):
  """Quote an SQL identifier.  Branch names contain '-' and upper case."""
  return '"' + s.replace('"', '""') + '"'


def copy_text(v):
  """Encode one value for COPY's text format.

  Not CSV: an empty CSV field reads back as NULL, and libversion.libversion
  does contain empty strings that have to stay empty strings.  The text format
  spells NULL \\N instead, so both survive.
  """
  if v is None:
    return "\\N"
  return (str(v).replace("\\", "\\\\").replace("\n", "\\n")
          .replace("\r", "\\r").replace("\t", "\\t"))


class Psql:
  """Runs SQL through the psql client, which keeps this script dependency-free."""

  def __init__(self, args):
    self.base = ["psql", "--no-psqlrc", "-w", "-v", "ON_ERROR_STOP=1",
                 "-h", args.host, "-p", str(args.port), "-U", args.user, "-d", args.dbname]
    self.schema = args.pgschema

  def _run(self, argv, stdin=None):
    p = subprocess.run(self.base + argv, input=stdin, capture_output=True, text=True)
    if p.returncode != 0:
      raise RuntimeError("psql failed: %s" % (p.stderr.strip() or p.stdout.strip()))
    return p.stdout

  def query(self, sql):
    """Execute SQL and return the output unaligned and without headers."""
    return self._run(["-tAq", "-c", "SET search_path TO %s; %s" % (ident(self.schema), sql)]).strip()

  def script(self, sql):
    return self._run(["-q", "-f", "-"], stdin="SET search_path TO %s;\n%s" % (ident(self.schema), sql))

  def copy_rows(self, table, columns, rows, also="", skip_existing=None):
    """COPY an iterable of tuples into a table.

    The statements in "also" are committed together with the batch, so that a
    migration killed halfway through never copies the same rows twice.

    With skip_existing set to the key columns of the table, the batch goes
    through a temporary table first and rows that are already there are
    dropped, so the machine migrated first keeps its results.
    """
    cols = ",".join(ident(c) for c in columns)
    target = ident(table)
    buf = io.StringIO()
    buf.write("SET search_path TO %s;\nBEGIN;\n" % ident(self.schema))
    if skip_existing:
      buf.write("CREATE TEMP TABLE batch (LIKE %s) ON COMMIT DROP;\n" % target)
      target = "batch"
    buf.write("COPY %s (%s) FROM STDIN;\n" % (target, cols))
    for r in rows:
      buf.write("\t".join(copy_text(v) for v in r))
      buf.write("\n")
    buf.write("\\.\n")
    if skip_existing:
      key = ",".join(ident(c) for c in skip_existing)
      # DISTINCT ON also drops rows the batch itself holds twice, which the
      # unique index would otherwise reject.
      buf.write("INSERT INTO %s (%s) SELECT DISTINCT ON (%s) %s FROM batch ORDER BY %s"
                " ON CONFLICT DO NOTHING;\n" % (ident(table), cols, key, cols, key))
    if also:
      buf.write(also.rstrip().rstrip(";") + ";\n")
    buf.write("COMMIT;\n")
    if not skip_existing:
      self._run(["-q", "-f", "-"], stdin=buf.getvalue())
      return None
    # Without -q psql reports every statement, so the INSERT says how many rows
    # were new - the interesting number when catching up with a running test.
    out = self._run(["-f", "-"], stdin=buf.getvalue())
    for line in out.split("\n"):
      if line.startswith("INSERT "):
        return int(line.split()[-1])
    return 0


def sqlite_tables(conn):
  """The tables of a testing database, split into (branch tables, lookup tables)."""
  names = [n for (n,) in conn.execute(
      "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")]
  branches, lookups = [], []
  for n in names:
    if any(n.startswith(p) for p in SKIP_PREFIXES):
      continue
    (lookups if n in LOOKUP_COLUMNS else branches).append(n)
  return branches, lookups


def sqlite_columns(conn, tbl):
  return [row[1] for row in conn.execute("PRAGMA table_info(%s)" % ident(tbl))]


def create_table(pg, tbl, columns):
  key = KEYS.get(tbl, BRANCH_KEY)
  cols = ",\n  ".join("%s %s%s" % (ident(c), t, " NOT NULL" if c in key else "")
                      for c, t in columns)
  pg.script("CREATE TABLE IF NOT EXISTS %s (\n  %s\n);" % (ident(tbl), cols))


def read_progress(pg, source, tbl):
  out = pg.query("SELECT last_rowid, rows_read, done FROM migration_progress "
                 "WHERE source=%s AND tbl=%s" % (lit(source), lit(tbl)))
  if not out:
    return 0, 0, False
  last, rows, done = out.split("|")
  return int(last), int(rows), done == "t"


def progress_sql(source, tbl, last_rowid, rows_read, done):
  return ("""INSERT INTO migration_progress (source, tbl, last_rowid, rows_read, done, updated_at)
             VALUES (%s, %s, %d, %d, %s, now())
             ON CONFLICT (source, tbl) DO UPDATE SET
               last_rowid = EXCLUDED.last_rowid, rows_read = EXCLUDED.rows_read,
               done = EXCLUDED.done, updated_at = now()"""
          % (lit(source), lit(tbl), last_rowid, rows_read, "true" if done else "false"))


def migrate_table(pg, sconn, source, tbl, columns, batch, quiet, skip_existing=False,
                  catch_up=False):
  """Copy one sqlite table into its PostgreSQL twin, in resumable batches."""
  names = [c for c, _ in columns]
  have = sqlite_columns(sconn, tbl)
  missing = [c for c in names if c not in have]
  # Older databases lack "parsing"; test.py defaults it to 0.0 as well.
  select = ",".join(ident(c) if c in have else "0" for c in names)

  create_table(pg, tbl, columns)
  last_rowid, rows_read, done = read_progress(pg, source, tbl)
  if done and not catch_up:
    if not quiet:
      print("  %-30s done earlier (%d rows)" % (tbl, rows_read))
    return rows_read, 0

  stored = None
  if catch_up:
    # Read the table from the start and keep the runs the database has never
    # seen.  Not "everything past the rowid we stopped at": VACUUM renumbers
    # the rowids of these tables, and clean-empty-omcversion-dates.py runs one
    # after every test, so that number cannot be trusted between two runs.
    stored = set(int(d) for d in pg.query("SELECT DISTINCT date FROM %s" % ident(tbl)).split("\n") if d)
    last_rowid, rows_read = 0, 0
    skip_existing = True
  key = KEYS.get(tbl, BRANCH_KEY) if skip_existing else None
  if not quiet:
    extra = ""
    if missing:
      extra = ", %s defaulted to 0" % ",".join(missing)
    if catch_up:
      extra += ", %d runs already stored" % len(stored)
    elif last_rowid:
      extra += ", resuming after rowid %d" % last_rowid
    print("  %-30s starting%s" % (tbl, extra))

  t0 = time.time()
  inserted = 0
  date = names.index("date")
  while True:
    chunk = sconn.execute(
        "SELECT rowid,%s FROM %s WHERE rowid > ? ORDER BY rowid LIMIT %d"
        % (select, ident(tbl), batch), (last_rowid,)).fetchall()
    if not chunk:
      break
    last_rowid = chunk[-1][0]
    rows_read += len(chunk)
    rows = [row[1:] for row in chunk]
    if stored is not None:
      rows = [r for r in rows if r[date] not in stored]
    if rows:
      n = pg.copy_rows(tbl, names, rows,
                       also=progress_sql(source, tbl, last_rowid, rows_read, False),
                       skip_existing=key)
      inserted += n if n is not None else len(rows)
    if not quiet:
      sys.stdout.write("\r  %-30s %d rows read, %d new (%.0f rows/s)"
                       % (tbl, rows_read, inserted, rows_read / max(time.time() - t0, 1e-9)))
      sys.stdout.flush()

  pg.query(progress_sql(source, tbl, last_rowid, rows_read, True))
  if not quiet and (inserted or not catch_up):
    print("\r  %-30s %d rows read, %d new%s" % (tbl, rows_read, inserted, " " * 30))
  return rows_read, inserted


def migrate_database(pg, args):
  sconn = sqlite3.connect("file:%s?mode=ro" % os.path.abspath(args.sqlite), uri=True)
  branches, lookups = sqlite_tables(sconn)
  if args.only:
    branches = [b for b in branches if b in args.only]
    lookups = [l for l in lookups if l in args.only]
  print("%s: %d branch tables, lookup tables: %s"
        % (args.sqlite, len(branches), ",".join(lookups) or "none"))

  total = new = 0
  for tbl in lookups + branches:
    cols = LOOKUP_COLUMNS.get(tbl, BRANCH_COLUMNS)
    (r, i) = migrate_table(pg, sconn, args.source, tbl, cols, args.batch, args.quiet,
                           args.skip_existing, args.catch_up)
    total += r
    new += i
  sconn.close()
  print("%s: %d rows read, %d new" % (args.sqlite, total, new))


def create_indexes(pg):
  """The unique key of each table, plus the index the reports need.

  The unique index has to exist before a second test machine is migrated with
  --skip-existing.  It also covers the [date] index the report scripts create
  on the fly in sqlite, since date is its first column.
  """
  tables = [t for t in pg.query(
      "SELECT tablename FROM pg_tables WHERE schemaname=current_schema() ORDER BY tablename").split("\n") if t]

  def index(tbl, cols, unique=False):
    name = ("%s_%s_%s" % ("uq" if unique else "idx", tbl, "_".join(cols)))[:63]
    pg.script("CREATE %sINDEX IF NOT EXISTS %s ON %s (%s);"
              % ("UNIQUE " if unique else "", ident(name), ident(tbl),
                 ",".join(ident(c) for c in cols)))

  for tbl in tables:
    if tbl == "migration_progress":
      continue
    index(tbl, KEYS.get(tbl, BRANCH_KEY), unique=True)
    if tbl == "omcversion":
      index(tbl, ["branch", "date"])
    elif tbl == "libversion":
      index(tbl, ["branch", "libname", "date"])
    else:
      index(tbl, ["libname", "date"])
    print("  indexed %s" % tbl)


def verify(pg, args):
  """Compare the sqlite and PostgreSQL row counts table by table."""
  sconn = sqlite3.connect("file:%s?mode=ro" % os.path.abspath(args.sqlite), uri=True)
  branches, lookups = sqlite_tables(sconn)
  bad = 0
  for tbl in lookups + branches:
    (n,) = sconn.execute("SELECT COUNT(*) FROM %s" % ident(tbl)).fetchone()
    m = int(pg.query("SELECT COUNT(*) FROM %s" % ident(tbl)) or 0)
    flag = "ok" if m >= n else "MISSING"
    if m < n:
      bad += 1
    print("  %-30s sqlite %10d  postgres %10d  %s" % (tbl, n, m, flag))
  sconn.close()
  print("%d of %d tables incomplete" % (bad, len(lookups) + len(branches)))
  return bad


def main():
  parser = argparse.ArgumentParser(description=__doc__,
                                   formatter_class=argparse.RawDescriptionHelpFormatter)
  parser.add_argument("--sqlite", help="sqlite3 database to migrate")
  parser.add_argument("--source", help="machine the database comes from, e.g. ripper1; used for resuming")
  parser.add_argument("--host", default=os.environ.get("PGHOST", "openmodelica.org"))
  parser.add_argument("--port", type=int, default=int(os.environ.get("PGPORT", 5432)))
  parser.add_argument("--user", default=os.environ.get("PGUSER", "om"))
  parser.add_argument("--dbname", default=os.environ.get("PGDATABASE", "omdb"))
  parser.add_argument("--pgschema", default="public", help="PostgreSQL schema to write to (default public)")
  parser.add_argument("--batch", type=int, default=200000, help="rows per COPY batch (default 200000)")
  parser.add_argument("--only", action="append", help="migrate only this table (repeatable)")
  parser.add_argument("--catch-up", action="store_true",
                      help="copy the runs written since the last migration, for a test that was "
                           "still using the sqlite3 database. Reads the whole database but only "
                           "writes the runs it does not have; safe to repeat")
  parser.add_argument("--skip-existing", action="store_true",
                      help="keep the rows already in the database when a key collides; "
                           "use it for every machine after the first one, and run --index before")
  parser.add_argument("--index", action="store_true", help="create the indexes; do this after loading")
  parser.add_argument("--verify", action="store_true", help="compare row counts with --sqlite")
  parser.add_argument("--quiet", action="store_true")
  args = parser.parse_args()

  pg = Psql(args)
  pg.script("CREATE SCHEMA IF NOT EXISTS %s;" % ident(args.pgschema))
  pg.script(PROGRESS_TABLE)

  if args.sqlite and not args.verify:
    if not args.source:
      parser.error("--sqlite needs --source, the machine name, e.g. ripper1")
    migrate_database(pg, args)

  if args.index:
    print("creating indexes")
    create_indexes(pg)

  if args.verify:
    if not args.sqlite:
      parser.error("--verify needs --sqlite")
    sys.exit(1 if verify(pg, args) else 0)


if __name__ == "__main__":
  main()