update_query.py

import argparse

import django
from django.db.models import QuerySet

django.setup()
from amcat.models import Query, CodingJob, Project
import logging
logging.basicConfig(level=logging.INFO, format='[%(asctime)s %(name)-12s %(levelname)-7s] %(message)s')

def update_queries(project: Project, queries: QuerySet, dry_run=False):
    # 1. Get all coding jobs and sets
    jobs = {job: aset for (job, aset) in
            CodingJob.objects.filter(project=project).values_list("pk", "articleset_id")}
    logging.info(f"Project contains {len(jobs)} codingjobs: {jobs}")

    # 2. Find queries to update
    logging.info(f"Updating {len(queries)} queries: {[q.name for q in queries]}")

    # 3. Update queries
    for q in queries:
        lastjob = max(q.parameters['codingjobs'])
        jobs_to_add = [j for j in jobs if j > lastjob]
        if jobs_to_add:
            sets_to_add = [jobs[j] for j in jobs_to_add]
            logging.info(f"Updating query {q.id}: {q.name}, adding jobs {jobs_to_add} and sets {sets_to_add}")
            if not dry_run:
                q.parameters["articlesets"] += sets_to_add
                q.parameters["codingjobs"] += jobs_to_add
                q.save()
        else:
            logging.info(f"Query {q.id}: {q.name} already up to date")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("project")
    parser.add_argument("name_pattern")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    project = Project.objects.get(pk=args.project)
    queries = Query.objects.filter(project=project, name__icontains=args.name_pattern)
    update_queries(project, queries, args.dry_run)