text: #
  # Copyright (c) 2019 MagicStack Inc.
  # All rights reserved.
  #
  # See LICENSE for details.
  ##
  
  
  import argparse
  import collections
  import datetime
  import json
  import progress.bar
  
  import sqlalchemy as sa
  import sqlalchemy.orm as orm
  
  import _sqlalchemy.models as m
  
  
  def bar(label, total):
      return progress.bar.Bar(label[:32].ljust(32), max=total)
  
  
  def bulk_insert(db, label, data, into):
      label = f"Creating {len(data)} {label}"
      pbar = bar(label, len(data))
  
      while data:
          chunk = data[:1000]
          data = data[1000:]
          db.execute(sa.insert(into), chunk)
          db.commit()
          pbar.next(len(chunk))
      pbar.finish()
  
  
  def reset_sequence(db, tablename):
      tab = sa.table(tablename, sa.column("id"))
  
      db.execute(
          sa.select(
              sa.func.setval(
                  f"{tablename}_id_seq",
                  sa.select(tab.c.id)
                  .order_by(tab.c.id.desc())
                  .limit(1)
                  .scalar_subquery(),
              )
          )
      )
  
  
  def load_data(filename, engine):
      session_factory = orm.sessionmaker(bind=engine)
      Session = orm.scoped_session(session_factory)
  
      with Session() as db:
  
          # first clear all the existing data
          print(f"purging existing data...")
  
          db.execute(sa.delete(m.Directors))
          db.execute(sa.delete(m.Cast))
          db.execute(sa.delete(m.Review))
          db.execute(sa.delete(m.Movie))
          db.execute(sa.delete(m.Person))
          db.execute(sa.delete(m.User))
          db.commit()
  
      # read the JSON data
      print("loading JSON... ", end="", flush=True)
      with open(filename, "rt") as f:
          records = json.load(f)
      data = collections.defaultdict(list)
      for rec in records:
          rtype = rec["model"].split(".")[-1]
          datum = rec["fields"]
          if "pk" in rec:
              datum["id"] = rec["pk"]
          # convert datetime
          if rtype == "review":
              datum["creation_time"] = datetime.datetime.fromisoformat(
                  datum["creation_time"]
              )
  
          data[rtype].append(datum)
      print("done")
  
      with Session() as db:
  
          # bulk create all the users
          bulk_insert(db, "users", data["user"], m.User)
  
          # bulk create all the people
          bulk_insert(db, "people", data["person"], m.Person)
  
          # bulk create all the movies
          bulk_insert(db, "movies", data["movie"], m.Movie)
  
          # bulk create all the reviews
          bulk_insert(db, "reviews", data["review"], m.Review)
  
          # bulk create all the directors
          bulk_insert(db, "directors", data["directors"], m.Directors)
  
          # bulk create all the cast
          bulk_insert(db, "cast", data["cast"], m.Cast)
  
          # reconcile the autoincrementing indexes with the actual indexes
          reset_sequence(db, "cast")
          reset_sequence(db, "directors")
          reset_sequence(db, "movie")
          reset_sequence(db, "person")
          reset_sequence(db, "user")
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(
          description="Load a specific fixture, old data will be purged."
      )
      parser.add_argument("filename", type=str, help="The JSON dataset file")
  
      args = parser.parse_args()
  
      engine = sa.create_engine(
          "postgresql+asyncpg://sqlalch_bench:edgedbbenchmark@localhost:15432/sqlalch_bench?async_fallback=True"
      )
  
      load_data(args.filename, engine)
  