Python: How to dynamically choose the database and table based on input with SQLAlchemy?
If you want to dynamically choose the database URL and table based on input, you can use a more flexible approach by creating a factory function to handle this. This will involve dynamically creating the engine and session based on the input, and then dynamically loading the appropriate model.
Here's a step-by-step guide to achieve this:
1. Define a Model Registry
You can maintain a dictionary of model classes for each database. This way, you can dynamically choose the model based on the table name.
from sqlalchemy import create_engine, Column, Integer, String
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy import inspect
Base = declarative_base()
# Define models for different tables
class User(Base):
__tablename__ = 'users'
id = Column(Integer, primary_key=True)
name = Column(String(255), default="", nullable=False)
class Product(Base):
__tablename__ = 'products'
id = Column(Integer, primary_key=True)
product_name = Column(String(255), default="", nullable=False)
# Map tables to models
model_registry = {
'users': User,
'products': Product
}
# 2. Create a Function to Initialize Session and Model
# This function will create the engine and session based on the input database URL and table name.
def get_engine_and_session(db_url):
engine = create_engine(db_url)
Session = sessionmaker(bind=engine)
session = Session()
return engine, session
def get_model_for_table(table_name):
return model_registry.get(table_name)
def initialize_table(db_url, table_name):
# Create engine and session
engine, session = get_engine_and_session(db_url)
# Get the model for the given table name
model = get_model_for_table(table_name)
if not model:
raise ValueError(f"No model defined for table '{table_name}'")
# Check if the table exists
inspector = inspect(engine)
if not inspector.has_table(model.__tablename__):
print(f"Creating table '{model.__tablename__}'")
create_table_for_model(engine, model)
# Or Create all tables if they don't exist
# Base.metadata.create_all(engine)
return session, model
def create_table_for_model(engine, model):
""" Creates only the specific table for the given model."""
model.metadata.create_all(engine, tables=[model.__table__])
# 3. Here's how you can use the function to dynamically choose the database URL and table:
if __name__ == "__main__":
# input_db_url = "mysql+mysqlclient://username:password@localhost/dbname"
input_db_url = "mysql+pymysql://username:password@localhost/dbname" # if using PyMySQL.
input_table = "users" # or "products"
session, model = initialize_table(input_db_url, input_table)
# Example: Add a new record
if model:
new_entry = model(name='John Doe') if input_table == 'users' else model(product_name='Gadget')
session.add(new_entry)
session.commit()
# Query example
results = session.query(model).all()
# Convert to list of dict:
result_list = [{k: v for k, v in result.__dict__.items() if k != '_sa_instance_state'} for result in results]
print(result_list)
***Note: You need to install these packages::
-
pip install sqlalchemy pymysql - or:
pip install sqlalchemy mysqlclient
