def get_db_classes()

in genai-on-vertex-ai/gemini/evals_playbook/utils/evals_playbook.py [0:0]


def get_db_classes():
    # Define engine, metadata and session
    engine = create_engine(f'bigquery://{cfg.PROJECT_ID}')
    metadata = MetaData()
    # Auto populate metadata
    for table_class in BQ_TABLE_MAP:
        table_name, update_key_cols = get_db_object(table_class)
        Table(table_name, metadata, *update_key_cols, autoload_with=engine, schema=cfg.BQ_DATASET_ID)
    # create objects
    Base = automap_base(metadata=metadata)
    Base.prepare()
    return Base