python - 使用 peewee 在 SQLite 上创建两个链接表的 bulk_create()：主键未更新

Question

我有两个具有 1-n 关系的表，我想使用bulk_create().

让我们以用户和推文为例。

from peewee import *

db = SqliteDatabase('my_app.db')

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

class User(BaseModel):
    username = CharField(unique=True)

class Tweet(BaseModel):
    user = ForeignKeyField(User, backref='tweets')
    message = TextField()

我想创建 和 的未保存实例User并Tweet用bulk_create(). 一个天真的解决方案是：

db.create_tables([User, Tweet])

john = User(username='John')
mehdi = User(username='Mehdi')
users = [john, mehdi]

tweets = [Tweet(user=john, message='Hello twitter world!'),
          Tweet(user=mehdi, message='This is my first message.'),
          Tweet(user=mehdi, message='This is my second message.')]

User.bulk_create(users)
Tweet.bulk_create(tweets)

不幸的是，这不起作用，因为User实例主键没有更新（如文档中所述，Postgres 数据库除外）。

由于似乎无法更新User实例的主键（即使可以，从数据库中读取主键可能效率很低），我能看到的唯一解决方案是使用我自己的主键并设置创建实例时。这意味着不使用 peewee 非常方便的自动递增主键系统，我也想知道在此之前是否有任何替代方法。

标签： pythonsqlitepeewee

Answer 1

这就是我出来的。我宁愿避免弄乱 peewee 的内部结构，但它工作正常。

from peewee import *

db = SqliteDatabase('my_app.db')


class BaseModel(Model):
    id_counter = 0
    id = IntegerField(primary_key=True, constraints=[SQL('AUTOINCREMENT')])

    def _set_id(self):
        if self.id is None:
            self.id = self.__class__.id_counter
            self.__class__.id_counter += 1

    def save(self, *args, **kwargs):
        return super(BaseModel, self).save(*args, **kwargs)

    @classmethod
    def bulk_create(cls, model_list, batch_size=None):
        max_id = cls.select(fn.MAX(cls.id)).scalar() or 0
        cls.id_counter = max_id + 1
        for model in model_list:
            model._set_id()
            model._update_fks()
        return super(BaseModel, cls).bulk_create(model_list=model_list, batch_size=batch_size)

    def _update_fks(self):
        for field_name, field in self._meta.fields.items():
            if isinstance(field, ForeignKeyField):
                fk_field_name = field_name + '_id'
                fk_id = getattr(self, field_name).id
                setattr(self, fk_field_name, fk_id)

    class Meta:
        database = db


class User(BaseModel):
    username = CharField(unique=True)


class Tweet(BaseModel):
    user = ForeignKeyField(User, backref='tweets', field='id')
    message = TextField()


db.create_tables([User, Tweet])

# Creating users and tweets one by one
sarah = User.create(username='Sarah')
Tweet.create(user=sarah, message='First tweet in history')

# Bulk user and tweet insertion
john = User(username='John')
mehdi = User(username='Mehdi')
users = [john, mehdi]

tweets = [Tweet(user=john, message='Hello twitter world!'),
          Tweet(user=mehdi, message='This is my first message.'),
          Tweet(user=mehdi, message='This is my second message.')]

User.bulk_create(users)
Tweet.bulk_create(tweets)

# Creating users and tweets one by one after bulk insertion
miranda = User.create(username='Miranda')

Tweet.create(user=miranda, message='The last tweet')