How to connect to CockroachDB using SQLModel, creating the database automatically if it doesn't exist

#!/usr/bin/env python3
from sqlmodel import Field, Session, SQLModel, create_engine
from sqlalchemy_utils import database_exists, create_database

from sqlalchemy.exc import ProgrammingError
    
def init_db():
    db = "mydb"
    db_url = f"cockroachdb+psycopg2://root@localhost:26257/{db}"
    engine = create_engine(db_url)
    # Create database if it doesnt exist
    try:
        SQLModel.metadata.create_all(engine)
    except ProgrammingError as ex:  # Catch the specific exception
        if f'database "{db}" does not exist' in str(ex):
            create_database(engine.url)
            # Now try again
            engine.dispose()
            engine = create_engine(db_url)
            SQLModel.metadata.create_all(engine)
        else: # Other exception
            raise ex
    return engine