Exemple #1
0
class PyPGQueueTestCase(TestCase):
  def setUp(self):
    self.da = ModelAccess(get_pg_core(CONNECTION_STRING)).open(autocommit=True)
    self._clean_tables()

  def tearDown(self):
    self._clean_tables()

  def _clean_tables(self):
    self.da.update(M.SerializationKey.table_name, dict(active_job_id=None))
    self.da.delete(M.Job.table_name)
    self.da.delete(M.ScheduledJob.table_name)
    self.da.delete(M.SerializationKey.table_name)

  def _table_counts(self):
    return (self.da.count(M.Job.table_name),
            self.da.count(M.ScheduledJob.table_name),
            self.da.count(M.SerializationKey.table_name))

  def _get_ts(self, seconds_in_future=0):
    return utc_now() + timedelta(seconds=seconds_in_future)

  def test_queue_job(self):
    self.assertEqual(self._table_counts(), (0, 0, 0))
    pypgq.queue_job(self.da, "Job A", dict(name="Trey", status="Start"))
    pypgq.queue_job(self.da, "Job B")
    pypgq.queue_job(self.da, "Job C", dict(name="Julie", status="In Progress"), "tenant/1")
    pypgq.queue_job(self.da, "Job D", None, "tenant/1")
    pypgq.queue_job(self.da, "Job E", None, "tenant/2")
    self.assertEqual(self._table_counts(), (5, 0, 2))

  def test_schedule_job(self):
    self.assertEqual(self._table_counts(), (0, 0, 0))
    pypgq.queue_job(self.da, "Job A", dict(name="Trey", status="Start"), scheduled_at=self._get_ts())
    pypgq.queue_job(self.da, "Job B", scheduled_at=self._get_ts())
    pypgq.queue_job(self.da, "Job C", dict(name="Julie", status="In Progress"), "tenant/1", scheduled_at=self._get_ts())
    pypgq.queue_job(self.da, "Job D", None, "tenant/1", scheduled_at=self._get_ts(2))
    pypgq.queue_job(self.da, "Job E", None, "tenant/2", scheduled_at=self._get_ts(2))
    self.assertEqual(self._table_counts(), (0, 5, 2))

    queue = pypgq.Queue(self.da, schedule_frequency=1)
    self.assertEqual(queue._schedule_jobs(), 3)
    self.assertEqual(self._table_counts(), (3, 2, 2))
    sleep(2)
    self.assertEqual(queue._schedule_jobs(), 2)
    self.assertEqual(self._table_counts(), (5, 0, 2))
    self.assertIsNone(queue._schedule_jobs())

  def test_run_jobs(self):
    qda = ModelAccess(get_pg_core(CONNECTION_STRING)).open(autocommit=True)
    queue = pypgq.Queue(
      qda,
      worker_count=2,
      cooperative=None)#pypgq.Cooperative(pypgq.Cooperative.advisory_lock, (1, 1)))
    queue._sleep_time = 1.0
    queue.add_handler("sleep_job", sleep_job)
    queue.add_handler("exception_job", exception_job)

    self.assertEqual(queue.status(),
                     {"waiting_jobs": 0,
                      "running_jobs": 0,
                      "running_job_ids": (),
                      "completed_jobs": 0,
                      "stop_mode": pypgq.StopMode.never})

    j1, _ = pypgq.queue_job(self.da, "sleep_job", {"seconds": 5})
    j2, _ = pypgq.queue_job(self.da, "sleep_job", {"seconds": 5})
    j3, _ = pypgq.queue_job(self.da, "exception_job", {"seconds": 5})

    sqt = Thread(target=queue.start)
    sqt.start()

    # NOTE: this is kind of tricy. We need to wait for the right intervals to check the status, and
    #       it's not easy to get right. The better result would be to have some kind of event system
    #       on the queue to call functions when it starts, stops, etc. jobs.
    sleep(2.5)
    self.assertEqual(queue.status(),
                     {"waiting_jobs": 1,
                      "running_jobs": 2,
                      "running_job_ids": (j1.id, j2.id),
                      "completed_jobs": 0,
                      "stop_mode": pypgq.StopMode.never})
    sleep(3.5)
    self.assertEqual(queue.status(),
                     {"waiting_jobs": 0,
                      "running_jobs": 1,
                      "running_job_ids": (j3.id, ),
                      "completed_jobs": 2,
                      "stop_mode": pypgq.StopMode.never})
    sleep(5.0)
    self.assertEqual(queue.status(),
                     {"waiting_jobs": 0,
                      "running_jobs": 0,
                      "running_job_ids": (),
                      "completed_jobs": 3,
                      "stop_mode": pypgq.StopMode.never})

    queue.stop(None, pypgq.StopMode.when_all_done)
    sqt.join()

  def test_job_serialization(self):
    qda = ModelAccess(get_pg_core(CONNECTION_STRING)).open(autocommit=True)
    queue = pypgq.Queue(qda, worker_count=2)
    queue._sleep_time = 1.0
    queue.add_handler("sleep_job", sleep_job)
    queue.add_handler("exception_job", exception_job)

    j1, _ = pypgq.queue_job(self.da, "sleep_job", {"seconds": 3}, "serialize")
    j2, _ = pypgq.queue_job(self.da, "sleep_job", {"seconds": 3}, "serialize")

    sqt = Thread(target=queue.start)
    sqt.start()

    sleep(1.5)
    self.assertEqual(queue.status(),
                     {"waiting_jobs": 1,
                      "running_jobs": 1,
                      "running_job_ids": (j1.id, ),
                      "completed_jobs": 0,
                      "stop_mode": pypgq.StopMode.never})
    sleep(3.0)
    self.assertEqual(queue.status(),
                     {"waiting_jobs": 0,
                     "running_jobs": 1,
                     "running_job_ids": (j2.id, ),
                     "completed_jobs": 1,
                     "stop_mode": pypgq.StopMode.never})
    sleep(2.0)
    self.assertEqual(queue.status(),
                     {"waiting_jobs": 0,
                      "running_jobs": 0,
                      "running_job_ids": (),
                      "completed_jobs": 2,
                      "stop_mode": pypgq.StopMode.never})

    queue.stop(None, pypgq.StopMode.when_all_done)
    sqt.join()
Exemple #2
0
class TransferTracker(object):
  """The TransferTracker class keeps track of data being transferred from one database or table to
  another. Specifically, it makes it possible to record the old and new ids for a relation, and
  query for the new ID.
  """
  def __init__(self, data_access_core):
    """Initialize the TransferTracker class."""
    self.data_access = ModelAccess(data_access_core)

  def open(self, autocommit=False):
    """Call-through to data_access.open."""
    self.data_access.open(autocommit=autocommit)
    return self

  def close(self, commit=True):
    """Call-through to data_access.close."""
    self.data_access.close(commit=commit)
    return self

  def commit(self):
    """Call-through to data_access.commit."""
    self.data_access.commit()
    return self

  def rollback(self):
    """Call-through to data_access.rollback."""
    self.data_access.rollback()
    return self

  def setup(self):
    """Creates the default relations and transfers tables. The SQL used may not work on all
    databases. (It was written for SQLite3)
    """
    cmds = [
      """
        create table if not exists relations (
          id integer not null primary key,
          name text not null unique,
          completed_at datetime
        );
      """,
      """
        create table if not exists transfers (
          relation_id integer not null references relations (id) on delete cascade,
          old_id text not null,
          new_id text,
          primary key (relation_id, old_id)
        );
      """,
      """
        create index if not exists transfers_relation_id_idx on transfers (relation_id);
      """
    ]

    for cmd in cmds:
      self.data_access.execute(cmd)
    self.data_access.commit()
    return self

  def cleanup(self):
    """Cleanup the database, calling vacuum and analyze."""
    self.data_access.execute("vacuum")
    self.data_access.execute("analyze")

  def reset(self, relation_name=None):
    """Reset the transfer info for a particular relation, or if none is given, for all relations.
    """
    if relation_name is not None:
      self.data_access.delete("relations", dict(name=relation_name))
    else:
      self.data_access.delete("relations", "1=1")
    return self

  def start_transfer(self, relation_name):
    """Write records to the data source indicating that a transfer has been started for a particular
    relation.
    """
    self.reset(relation_name)
    relation = Relation(name=relation_name)
    self.data_access.insert_model(relation)
    return relation

  def register_transfer(self, relation, old_id, new_id):
    """Register the old and new ids for a particular record in a relation."""
    transfer = Transfer(relation_id=relation.id, old_id=old_id, new_id=new_id)
    self.data_access.insert_model(transfer)
    return transfer

  def complete_transfer(self, relation, cleanup=True):
    """Write records to the data source indicating that a transfer has been completed for a
    particular relation.
    """
    relation.completed_at = utc_now().isoformat()
    self.data_access.update_model(relation)
    if cleanup:
      self.cleanup()
    return relation

  def is_transfer_complete(self, relation_name):
    """Checks to see if a tansfer has been completed."""
    phold = self.data_access.sql_writer.to_placeholder()
    return self.data_access.find_model(
      Relation,
      ("name = {0} and completed_at is not null".format(phold), [relation_name])) is not None

  def get_new_id(self, relation_name, old_id, strict=False):
    """Given a relation name and its old ID, get the new ID for a relation. If strict is true, an
    error is thrown if no record is found for the relation and old ID.
    """
    record = self.data_access.find(
      "relations as r inner join transfers as t on r.id = t.relation_id",
      (("r.name", relation_name), ("t.old_id", old_id)),
      columns="new_id")
    if record:
      return record[0]
    else:
      if strict:
        raise KeyError("{0} with id {1} not found".format(relation_name, old_id))

  def id_getter(self, relation_name, strict=False):
    """Returns a function that accepts an old_id and returns the new ID for the enclosed relation
    name."""
    def get_id(old_id):
      """Get the new ID for the enclosed relation, given an old ID."""
      return self.get_new_id(relation_name, old_id, strict)
    return get_id