Сеанс SQLAlchemy очищается в задании сельдерея и функции on_success - PullRequest
Новая
       23

Сеанс SQLAlchemy очищается в задании сельдерея и функции on_success

0 голосов
/ 28 июня 2019

Я создаю инструмент, который извлекает данные из другой базы данных, преобразует их и сохраняет в моей собственной базе данных. Я мигрирую из APScheduler в Celery, но столкнулся со следующей проблемой:

Я использую класс, который я называю JobRecords, чтобы сохранить, когда задание было выполнено, было ли оно успешным и с какими ошибками оно встречалось. Я использую это, чтобы не слишком оглядываться на обновленные записи, тем более что в некоторых таблицах несколько миллионов строк.

Поскольку система одинакова для всех заданий, я создал подкласс из объекта сельдерея Task. Я проверяю, выполняется ли задание в контексте приложения Flask, и извлекаю последний раз, когда это задание было успешно завершено. Также я регистрирую значение для now, чтобы избежать проблем с синхронизацией между запросом базы данных и добавлением записи задания. 

class RecordedTask(Task):
  """
  Task sublass that uses JobRecords to get the last run date
  and add new JobRecords on completion
  """
  now: datetime = None
  ignore_result = True

  _session: scoped_session = None
  success: bool = True
  info: dict = None

  @property
  def session(self) -> Session:
    """Making sure we have one global session instance"""
    if self._session is None:
      from app.extensions import db
      self._session = db.session
    return self._session

  def __call__(self, *args, **kwargs):
    from app.models import JobRecord

    kwargs['last_run'] = (
        self.session.query(func.max(JobRecord.run_at_))
        .filter(JobRecord.job_id == self.name, JobRecord.success)
        .first()
    )[0] or datetime.min
    self.now = kwargs['now'] = datetime.utcnow()

    with app.app_context():
      super(RecordedTask, self).__call__(*args, **kwargs)

  def on_failure(self, exc, task_id, args: list, kwargs: dict, einfo):
    self.session.rollback()
    self.success = False
    self.info = dict(
        args=args,
        kwargs=kwargs,
        error=exc.args,
        exc=format_exception(exc.__class__, exc, exc.__traceback__),
    )
    app.logger.error(f"Error executing job '{self.name}': {exc}")

  def on_success(self, retval, task_id, args: list, kwargs: dict):
    app.logger.info(f"Executed job '{self.name}' successfully, adding JobRecord")

    for entry in self.to_trigger:
      if len(entry) == 2:
        job, kwargs = entry
      else:
        job, = entry
        kwargs = {}
      app.logger.info(f"Scheduling job '{job}'")
      current_celery_app.signature(job, **kwargs).delay()

  def after_return(self, *args, **kwargs):
    from app.models import JobRecord
    record = JobRecord(
        job_id=self.name,
        run_at_=self.now,
        info=self.info,
        success=self.success
    )
    self.session.add(record)
    self.session.commit()
    self.session.remove()

Я добавил пример задания по обновлению модели под названием Location, но таких заданий много, как этот. 

@celery.task(bind=True, name="update_locations")
def update_locations(self, last_run: datetime = datetime.min, **_):
  """Get the locations from the external database and check for updates"""
  locations: List[ExternalLocation] = ExternalLocation.query.filter(
      ExternalLocation.updated_at_ >= last_run
  ).order_by(ExternalLocation.id).all()

  app.logger.info(f"ExternalLocation: collected {len(locations)} updated locations")
  for update_location in locations:
    existing_location: Location = Location.query.filter(
        Location.external_id == update_location.id
    ).first()

    if existing_location is None:
      self.session.add(Location.from_worker(update_location))
    else:
      existing_location.update_from_worker(update_location)

Проблема в том, что когда я запускаю это задание, объекты Location не фиксируются с JobRecord, поэтому создается только последнее. Если я отслеживаю его с помощью отладчика, Location.query.count() возвращает правильное значение внутри функции, но как только он входит в обратный вызов on_success, он возвращается к 0, а self._session.new возвращает пустой dict.

Я уже пытался добавить сеанс как свойство, чтобы убедиться, что он везде одинаковый, но проблема все еще сохраняется. Может быть, это как-то связано с scoped_session из-за Flask-SQLAlchemy?

Извините за большой объем кода, я пытался убрать как можно больше. Любая помощь приветствуется!

1 Ответ

0 голосов
/ 30 июня 2019

Я обнаружил, что виновником была комбинация scoped_session и контекст приложения Flask. Как и любой контекстный менеджер, запуск кода with app.app_context() вызвал функцию __exit__ при выходе, что, в свою очередь, привело к очистке ScopedRegistry, где хранился scoped_session. Затем был создан новый сеанс, к нему были добавлены JobRecords, и этот сеанс был зафиксирован. Поэтому местоположения не будут записываться в базу данных.

Есть два возможных решения. Если у вас нет сеансов в других файлах, чем в вашей задаче, вы можете добавить свойство сеанса в задачу. Таким образом, вы полностью избегаете scoped_session и можете выполнить очистку в функции after_return. 

 @property 
 def session(self):
   if self._session is None:
      from dashboard.extensions import db
      self._session = db.create_session(options={})()
    return self._session

Однако я также обращался к сеансу в моих файлах определения модели через from extensions import db. Поэтому я использовал две разные сессии. Я закончил тем, что использовал app.app_context().push() вместо contextmanager, таким образом избегая функции __exit__ 

  app.app_context().push()
  super(RecordedTask, self).__call__(*args, **kwargs)
...