Python: from SQL to ORM (SQLModel)

I had to update a small Python script that I wrote a long time ago, and this was the occasion to try SQLModel.

The script reads the data from an Excel file and inserts them into a (remote) PostgreSQL database (I use Aiven, the free plan, more enough for my simple needs. They do not sponsor me; mentioning them is my way to say thanks for the free option).

The first step is data-wrangling to ensure the input data are consistent. I use Pandas for this. The main advantage of using a DataFrame is that each column must contain the same data type, which helps validate data.

Then, I used a function to iterate through the DataFrame and insert each row into the database using psycopg2 and pure SQL statements. I always use type hints in my code so the linter can immediately report any issue. Still, we know that Python is dynamically typed, so data at runtime may be of the wrong type, and we won’t be able to catch it without type-checking (in this particular case, using the DataFrame is already a sort of control).

SQLModel is built upon SQLAlchemy and Pydantic; the latter of which is an excellent tool for data validation. To create the models, I used SQLAcodegen, which reads the Database and automatically writes the corresponding classes into a Python file. SQLAcodegen can be used via the command line or inside a Python script. I chose to use Python, and I generated the models for just some tables in the database that were needed for the purpose.

"""Script to generate SQLModel classes from database schema. sqlacodegen 3.0.0"""

from dotenv import dotenv_values
from sqlacodegen.generators import SQLModelGenerator
from sqlalchemy import MetaData, create_engine

# Load configuration
config_dict = dotenv_values(".env")

# Database connection
DATABASE_URL = f"postgresql://{config_dict['PGUSER']}:{config_dict['PGPASSWORD']}"
DATABASE_URL += f"@{config_dict['PGHOST']}:{db_port}/{config_dict['PGDATABASE']}"
DATABASE_URL += "?sslmode=require" # Aiven requires SSL

engine = create_engine(DATABASE_URL)

# Define the tables you want to generate models for
selected_tables = [
    "table_1",
    "table_2",
]

# Reflect metadata for selected tables only
metadata = MetaData()
metadata.reflect(bind=engine, only=selected_tables)

# Reflect only selected tables
for table_name in selected_tables:
    metadata.reflect(bind=engine, only=[table_name])

# Define SQLModelGenerator options (as a list)
generator_options = [
    "use_inflect",  # use the inflect library when naming classes and relationships
    "nojoined",  # don’t try to detect joined-class inheritance 
]
# Generate SQLModel output
generator = SQLModelGenerator(metadata, bind=engine, options=generator_options)  # Use SQLModel format
generated_code = generator.generate()

# Save the generated models to a file
with open("models.py", "w") as f:
    f.write(generated_code)

Since my code already used a parsing function to return each row adequately formatted, it was just a matter of replacing the return value from a Dictionary to a model object. In doing so, the linter suggested several adjustments for better matching with Pydantic fields.

Get Andrea Dalseno’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

Even if I’m comfortable with SQL, using the ORM simplified the code, making it more readable and easier to maintain. Having understood the differences between pure SQL and using an ORM, the effort was minimal for this simple script and undoubtedly worthwhile.

To understand the syntax differences between pure SQL and the ORM, let’s look at this sample script:

import datetime

from sqlmodel import Field, SQLModel, func, select

# Create a simple model 
class Sample(SQLModel, table=True):
    id: int = Field(default=None, primary_key=True)
    field1: int
    field2: float
    field3: float
    field4: datetime.date
    field5: str


max_date = datetime.date(2025, 1, 1)


statement = select(
    func.count().label("count"),
    func.sum(Sample.field1).label("sum1"),
    func.sum(Sample.field2).label("sum2"),
    func.sum(Sample.field3).label("sum3"),
).where(Sample.field4 > max_date)
print(statement)

The above snippet will produce:

SELECT count(*) AS count, sum(sample.field1) AS sum1, sum(sample.field2) AS sum2, sum(sample.field3) AS sum3 
FROM sample 
WHERE sample.field4 > :field4_1

and

statement = select(Sample.id).where(Sample.field5.like("AB%"))  # type: ignore
print(statement)

will give (note the type: ignore since the linter complains about the like, which is instead correct):

SELECT sample.id 
FROM sample 
WHERE sample.field5 LIKE :field5_1