Skip to main content

Using Extensions with SQLAlchemy

PostgreSQL extensions are add-on modules that extend the functionality of the database by providing new data types, operators, functions, procedural languages, and more.

This guide explains how to define a schema field that uses the Geometry data type provided by the PostGIS extension, and configure the schema migration to manage both PostgreSQL extension installations and the SQLAlchemy models as a single migration unit using Atlas. This explanation can be extrapolated to other extensions and libraries that link with SQLAlchemy.

Atlas Pro Feature

Atlas support for Extensions is available exclusively to Pro users. To use this feature, run:

atlas login

Getting started with Atlas and SQLAlchemy​

Before following this guide, make sure you are properly set up with Atlas for SQLAlchemy integration.

Composite Schema​

In order to expand our PostgreSQL schema migration to include both extensions and our SQLAlchemy models, we configure Atlas to read the state of the schema from a composite schema data source. Follow the steps below to configure this for your project:

  1. Create a SQL or HCL file that defines the necessary extensions used by your database.
schema.sql
CREATE EXTENSION postgis;
  1. In your SQLAlchemy model, define a field that uses the data type provided by the extension.
from sqlalchemy import Column, Integer, String, ForeignKey, DateTime
from sqlalchemy.orm import relationship, DeclarativeBase
from geoalchemy2 import Geometry

class Base(DeclarativeBase):
    pass

class Venue(Base):
    __tablename__ = "venues"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    address = Column(Geometry('POINT', srid=4326), nullable=False)
  1. In your atlas.hcl config file, add a composite_schema that includes your extensions defined in schema.sql and your SQLAlchemy model:
atlas.hcl
data "external_schema" "sqlalchemy" {
    program = [
        "atlas-provider-sqlalchemy",
        "--path", "./models.py",   // path to your SQLAlchemy models directory or file
        "--dialect", "postgresql"
    ]
}

data "composite_schema" "app" {
  # Load custom types first to define extensions
  schema {
    url = "file://schema.sql"
  }
  # Then, load SQLAlchemy models
  schema "public" {
    url = data.external_schema.sqlalchemy.url
  }
}

env "local" {
  src = data.composite_schema.app.url   // link to current schema state
  dev = "docker://postgis/latest/dev"   // dev database
}

Usage​

After setting up our composite schema, we can get its representation using the atlas schema inspect command, generate schema migrations for it, apply them to a database, and more. Below are a few commands to get you started with Atlas:

Inspect the Schema​

The atlas schema inspect command is commonly used to inspect databases. However, we can also use it to inspect our composite_schema and print the SQL representation of it:

atlas schema inspect \
  --env local \
  --url env://src \
  --format '{{ sql . }}'

The command above prints the following SQL.

-- Add new schema named "public"
CREATE SCHEMA IF NOT EXISTS "public";
-- Set comment to schema: "public"
COMMENT ON SCHEMA "public" IS 'standard public schema';
-- Create extension "postgis"
CREATE EXTENSION "postgis" WITH SCHEMA "public" VERSION "3.5.2";
-- Create "venues" table
CREATE TABLE "public"."venues" ("id" serial NOT NULL, "name" character varying NOT NULL, "address" 
public.geometry(Point,4326) NOT NULL, PRIMARY KEY ("id"));
-- Create index "idx_venues_loc" to table: "venues"
CREATE INDEX "idx_venues_loc" ON "public"."venues" USING gist ("loc");
-- Create index "ix_venues_id" to table: "venues"
CREATE INDEX "ix_venues_id" ON "public"."venues" ("id");

Generate Migrations For the Schema​

To generate a migration for the schema, run the following command:

atlas migrate diff --env local

Note that a new migration file is created with the following content:

migrations/20250206130414.sql
-- Create extension "postgis"
CREATE EXTENSION "postgis" WITH SCHEMA "public" VERSION "3.5.2";
-- Create "venues" table
CREATE TABLE "public"."venues" ("id" serial NOT NULL, "name" character varying NOT NULL, 
"address" public.geometry(Point,4326) NOT NULL, PRIMARY KEY ("id"));
-- Create index "idx_venues_loc" to table: "venues"
CREATE INDEX "idx_venues_loc" ON "public"."venues" USING gist ("loc");
-- Create index "ix_venues_id" to table: "venues"
CREATE INDEX "ix_venues_id" ON "public"."venues" ("id");

Apply the Migrations​

To apply the migration generated above to a running database, run the following command:

atlas migrate apply \
  --env local \
  --url "postgres://postgres:pass@localhost:5432/database?sslmode=disable"
note

Extensions are database-level objects, so they are installed at the database level. Therefore, when working with extensions, the scope of the migration should be set to the database, where objects are qualified with the schema name. Hence, the search_path is dropped from the dev-database URL in the atlas.hcl file.