使用 peewee 的 SQLite 上两个链接表的 bulk_create():主键未更新

bulk_create() of two linked tables on SQLite using peewee : primary key not updated

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

在这些表中插入数据

我们来 User and Tweet example.

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()

我想创建 UserTweet 的未保存实例并用 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 实例主键未更新 (as stated in the documentation except for Postgres databases)。

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

这就是我想出来的。我宁愿避免弄乱 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')