Search code examples
pythonpeewee

peewee check automatically created id not in result of subquery


I have next data structure:

from enum import IntEnum, unique
from pathlib import Path
from datetime import datetime
from peewee import *

@unique
class Status(IntEnum):
    CREATED = 0
    FAIL = -1
    SUCCESS = 1

db_path = Path(__file__).parent / "test.sqlite"
database = SqliteDatabase(db_path)

class BaseModel(Model):
    class Meta:
        database = database

class Unit(BaseModel):
    name = TextField(unique=True)
    some_field = TextField(null=True)
    created_at = DateTimeField(default=datetime.now)

class Campaign(BaseModel):
    id_ = AutoField()
    created_at = DateTimeField(default=datetime.now)

class Task(BaseModel):
    id_ = AutoField()
    status = IntegerField(default=Status.CREATED)
    unit = ForeignKeyField(Unit, backref="tasks")
    campaign = ForeignKeyField(Campaign, backref="tasks")

Next code create units, campaign and tasks:

def fill_units(count):
    units = []
    with database.atomic():
        for i in range(count):
            units.append(Unit.create(name=f"unit{i}"))
    return units

def init_campaign(count):
    units = Unit.select().limit(count)
    with database.atomic():
        campaign = Campaign.create()
        for unit in units:
            Task.create(unit=unit, campaign=campaign)
    return campaign

The problem appears when I'm trying to add more units into existing campaign. I need to select units which haven't been used in this campaign. In SQL I can do this using next query:

SELECT * FROM unit WHERE id NOT IN (SELECT unit_id FROM task WHERE campaign_id = 1) LIMIT 10

But how to do this using peewee?

The only way I've found yet is:

def get_new_units_for_campaign(campaign, count):
    unit_names = [task.unit.name for task in campaign.tasks]
    units = Unit.select().where(Unit.name.not_in(unit_names)).limit(count)
    return units

It's somehow works but I'm 100% sure that it's the dumbest way to implement this. Could you show me the proper way to implement this?


Solution

  • Finally I found this:

    Unit.select().where(Unit.id.not_in(campaign.tasks.select(Task.unit))).limit(10)
    

    Which produces

    SELECT "t1"."id", "t1"."name", "t1"."some_field", "t1"."created_at" FROM "unit" AS "t1" WHERE ("t1"."id" NOT IN (SELECT "t2"."unit_id" FROM "task" AS "t2" WHERE ("t2"."campaign_id" = 1))) LIMIT 10
    

    Which matches with SQL query I've provided in my question.

    P.S. I've done some research and it seems to be a proper implementation, but I'd appreciate if somebody correct me and show the better way (if exist).