Setting up SQLModel

At this stage, our application has a basic CRUD API. However, the data we are working with comes from a JSON file, rather than a database. In typical web APIs, the data usually comes from a database - this can be a SQL database such as SQLite, PostgreSQL or MySQL, or a NoSQL database such as DynamoDB or MongoDB.

We want to change the setup of our application to use a relational database. To keep things simple, we'll use SQLite, a file-based database.

SQLModel is a package that allows us to interact with relational, SQL-based databases. It was created by Sebastián Ramírez, who also created (among other things) the FastAPI package. In this post, we are going to introduce SQLModel and use it to create a Track model class that will be translated to a database table, with column names corresponding to field names, and types and constraints that are defined by Python type-hints.

The video for this content can be found below:


Objectives

In this post, we will demonstrate how to:

  • Install the new SQLModel package
  • Create SQLModel classes that map to database tables
  • Bootstrap our database tables with a command
  • Load in data from a JSON file using a FastAPI app startup-event - and learn about the Session and select constructs from SQLModel.

Adding SQLModel to our Application

Let's start by installing the SQLModel package. In your environment, run the following command.

pip install sqlmodel

Once installed, create a new file at the root of the project called database.py. This will be the home for database setup code and database model definitions. Add the following code.

from datetime import datetime
from typing import Optional

from sqlmodel import Field, SQLModel, create_engine

# There should be one engine for the entire application
DB_FILE = 'db.sqlite3'
engine = create_engine(f"sqlite:///{DB_FILE}", echo=True)

class TrackModel(SQLModel, table=True):
    # Note: Optional fields are marked as Nullable in the database
    __tablename__ = 'track'
    id: Optional[int] = Field(default=None, primary_key=True)
    title: str
    artist: str
    duration: float
    last_play: datetime


def create_tables():
    """Create the tables registered with SQLModel.metadata (i.e classes with table=True).
    More info: https://sqlmodel.tiangolo.com/tutorial/create-db-and-table/#sqlmodel-metadata
    """
    SQLModel.metadata.create_all(engine)

if __name__ == '__main__':
    # creates the table if this file is run independently, as a script
    create_tables()

On lines 7-8, we set up the database engine, using the create_engine function imported from SQLModel. We specify a SQLite connection string, and a database file called db.sqlite3.

On line 10, we define a TrackModel class which inherits from the SQLModel class. This class's fields (lines 13-17) will map to columns in a table with the name track.

The type-hints indicate what the column types should be, and additional constraints can be specified too - for example, an Optional type marks the column as NULLABLE, and we can set default values and primary key constraints using the sqlmodel.Field class (line 13).

Line 20-24 defines a simple function that wraps code that will create the table. This works by looking at all SQLModel classes that have the table=True parameter, and then creating a table based off the classes and attributes. This is wrapped in a function so that it can be called directly, rather than executed each time the module is loaded. The function is called when the script is executed as an independent Python script, which explains the code on lines 26-28.

The key feature of SQLModel is the ability to define type-hinted classes that inherit from the SQLModel class and map to tables in your database. Now, let's actually create the database and the track table. We need to run the database.py script using Python, such as below.

python database.py

This should create the SQLite database file, and create a table called track with the correct schema. You should now see a file called db.sqlite3 in the application folder, and you can use a tool such as VSCode's vscode-sqlite extension - right click on your new database file and select "Open Database" if using that extension, and a SQlite Explorer panel should open on the Explorer. Other tools such as DB Browser for SQLite can also be used to inspect your database.

Now that we have the model set up, and have created the database, we can start using these in our API.

Modifying main.py

Previously, we used FastAPI's startup-event to load a JSON file of music tracks into a list in the main.py module. We now want to modify this - instead of populating a list, we want to populate the SQLite database that we have created!

This is a more realistic real-life setup, because we want our data to be in a persistent data store, such as a database, rather than just loaded into memory.

We're going to modify our start-up function to interact with the database we have just set up, and load in the JSON data to the new track table, if the data does not already exist in the table. We'll see some new SQLModel features - including how to create a Session object that uses the database engine created in our database.py module to interact with the underlying database, how to query objects from a particular table, and how to do inserts to create new records in the table.

The new code for the main.py module will look like this (API routes are omitted).

import json
import pathlib
from typing import List, Union

from fastapi import FastAPI, Response
from sqlmodel import Session, select

from models import Track
from database import TrackModel, engine

# instantiate the FastAPI app
app = FastAPI()


# define app start-up event
@app.on_event("startup")
async def startup_event():
    DATAFILE = pathlib.Path() / 'data' / 'tracks.json'

    # create a Session scoped to the startup event
    # Note: we can also use a context manager
    session = Session(engine)

    # check if the database is already populated
    stmt = select(TrackModel)
    result = session.exec(stmt).first()

    # Load data if there's no results
    if result is None:
        with open(DATAFILE, 'r') as f:
            tracks = json.load(f)
            for track in tracks:
                session.add(TrackModel(**track))
    
        session.commit()
    session.close()

We import the Session object and the select function from SQLModel on line 6, and import our TrackModel and database engine on line 9. We then redefine the startup_event() function, creating a Session object on line 22 and creating a simple select(TrackModel) statement on line 25 that will select all of the columns from the table that corresponds to the TrackModel class.

We can then execute the statement on line 26, using the Session object's .exec() method, and getting the first result using the .first() method. If we retrieve a result, then we know the database table is populated - however if we don't, we load in the JSON file and create TrackModel objects for each of the JSON data records (line 29-33). These are added to the session and committed after the loop is finished (line 35).

Finally, we close the session on line 36 to release its resources.

The purpose of this startup-event function is to populate our database, if and only if the records don't already exist. This is why we execute the select statement and check if any records do exist.

You can run the uvicorn server to start the app, and therefore run this code, using the command below.

uvicorn main:app

If you are running this for the first time, you should see an INSERT INTO SQL statement on the shell, and some additional output. This corresponds to the records being inserted into the table.

If you stop and run the server again, you should not see the INSERT INTO statements, because the database has already been bootstrapped.

You can now inspect the table with your SQLite browsing tool to verify the data has been entered. For example, vscode-sqlite has a "Show Table" option when you right-click the table in the SQLite Explorer - this can be used to show a tabular view of the table.

If you only have Python available to you, here's a short script that will print out the first 5 records from the database - you can use this to verify that the data has been loaded.

import sqlite3

try:
    conn = sqlite3.connect('db.sqlite3')
except sqlite3.OperationalError:
    print("You need to create the database - python database.py")
else:
    cursor = conn.cursor()
    count = cursor.execute("SELECT COUNT(*) FROM track").fetchone()[0]
    print(f"There are {count} records in the Track table")
finally:
    conn.close()

Summary

In this post, we've introduced SQLModel and demonstrated model classes, the database engine, sessions and the select statement.

The next step is to modify our API endpoints to actually use the database. We will cover that in the next post!

If you enjoyed this post, please subscribe to our YouTube channel and follow us on Twitter to keep up with our new content!

Please also consider buying us a coffee, to encourage us to create more posts and videos!

;