Skip to content
On this page

Database connection in Python

Simple connection

py
import sqlite3

# Connect to database
conn = sqlite3.connect("example.db")
cursor = conn.cursor()

cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)")
cursor.execute("INSERT INTO users (name) VALUES ('Alice')")
cursor.execute("SELECT * FROM users")
print(cursor.fetchall())

conn.commit()
conn.close()

Using context manager

py
import sqlite3
from contextlib import contextmanager

# `open_db` context manager
@contextmanager
def open_db(db_name: str):
    conn = sqlite3.connect(db_name)
    try:
        cursor = conn.cursor()
        yield cursor
    finally:
        conn.commit()
        conn.close()

# # Class based context manager
# class OpenDB:
#     def __init__(self, file_name: str):
#         self.file_name = file_name
#         self.conn = sqlite3.connect(self.file_name)

#     def __enter__(self):
#         return self.conn.cursor()

#     def __exit__(self, exc_type, exc_val, exc_tb):
#         self.conn.commit()
#         self.conn.close()

# Connect to database
with open_db("example.db") as cursor:
    cursor.execute("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT)")
    cursor.execute("INSERT INTO users (name) VALUES ('Alice')")
    cursor.execute("SELECT * FROM users")
    print(cursor.fetchall())

SQLAlchemy & Context Manager

py
from sqlalchemy import create_engine, String
from sqlalchemy.orm import sessionmaker, MappedAsDataclass, DeclarativeBase, Mapped, mapped_column

class Base(MappedAsDataclass, DeclarativeBase):
    pass

class User(Base):
    __tablename__ = "users"
    id: Mapped[int] = mapped_column(primary_key=True, init=False)
    name: Mapped[str] = mapped_column(String(30))

# Connect to database
engine = create_engine("sqlite:///example.db")
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)

# Execute queries using context manager
with Session() as session:
    session.add(User(name="Alice"))
    session.commit()
    users = session.query(User).all()
    print(users)

Last updated: