Fix get_df with no related fields

This commit is contained in:
phil 2024-05-07 00:29:26 +02:00
parent 2ea672293d
commit 50920059b6

View file

@ -106,20 +106,21 @@ class BaseModel(SQLModel):
## Chamge column names to reflect the joined tables ## Chamge column names to reflect the joined tables
## Leave the first columns unchanged, as their names come straight ## Leave the first columns unchanged, as their names come straight
## from the model's fields ## from the model's fields
joined_columns = list(df.columns[len(cls.model_fields):]) if with_related:
renames: dict[str, str] = {} joined_columns = list(df.columns[len(cls.model_fields):])
# Match colum names with the joined tables renames: dict[str, str] = {}
# This uses the fact that orders of the joined tables # Match colum names with the joined tables
# and their columns is preserved by sqlalchemy query options (joinedlaod), # This uses the fact that orders of the joined tables
# and pandas' read_sql # and their columns is preserved by sqlalchemy query options (joinedlaod),
# Important: apparently, the order defined in cls.selectinload has to match the # and pandas' read_sql
# order of the relationship fields defined in the model class definition # Important: apparently, the order defined in cls.selectinload has to match the
for joined_table in joined_tables: # order of the relationship fields defined in the model class definition
target = joined_table.property.target # type: ignore for joined_table in joined_tables:
for col in target.columns: target = joined_table.property.target # type: ignore
## Pop the column from the colujmn list and make a new name for col in target.columns:
renames[joined_columns.pop(0)] = f'{target.schema}_{target.name}_{col.name}' ## Pop the column from the colujmn list and make a new name
df.rename(columns=renames, inplace=True) renames[joined_columns.pop(0)] = f'{target.schema}_{target.name}_{col.name}'
df.rename(columns=renames, inplace=True)
## Finally, set the index of the df as the index of cls ## Finally, set the index of the df as the index of cls
df.set_index([c.name for c in cls.__table__.primary_key.columns], # type: ignore df.set_index([c.name for c in cls.__table__.primary_key.columns], # type: ignore
inplace=True) inplace=True)