How to manage a peewee database in a separate module?

Viewed 5256

I want to have my database implementation in a separate module or class. But I am struggling with a few details. A simple example:

from peewee import *

db = SqliteDatabase(':memory:')

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

class User(BaseModel):
    name = CharField()

db.connect()
db.create_tables([User,])
db.commit()

@db.atomic()
def add_user(name):
    User.create(name=name).save()

@db.atomic()
def get_user(name):
    return User.get(User.name == name)

So far this is working fine. I can implement my interface to the database here and import this as a module.

Now I want to be able to choose the database file at runtime. So I need a way to define the Model classes without defining SqliteDatabase('somefile') before. I tried to encapsulate everything in a new Database class, which I can later import and create an instance from:

from peewee import *

class Database:

    def __init__(self, dbfile):
        self.db = SqliteDatabase(dbfile)

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

        class User(BaseModel):
            name = CharField()

        self.User = User

        self.db.connect()
        self.db.create_tables([User,])
        self.db.commit()

    @self.db.atomic()    # Error as self is not known on this level
    def add_user(self, name):
        self.User.create(name=name).save()

    @self.db.atomic()    # Error as self is not known on this level
    def get_user(self, name):
        return self.User.get(self.User.name == name)

Now I can call for example database = Database('database.db') or choose any other file name. I can even use multiple database instance in the same program, each with its own file.

However, there are two problems with this approach:

  • I still need to specify the database driver (SqliteDatabase) before defining the Model classes. To solve this I define the Model classes within the __init__() method, and then create an alias to with self.User = User. I don't really like this approach (it just doesn't feel like neat code), but at least it works.
  • I cannot use the @db.atomic() decorator since self is not known at class level, I would an instance here.

So this class approach does not seem to work very well. Is there some better way to define the Model classes without having to choose where you want to store your database first?

3 Answers
Related