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

Question

我有两个具有 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()

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

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

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