Make primary and foreign_key information available in model_fields()

[1cadumagalhaes]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the SQLModel documentation, with the integrated search.
• I already searched in Google "How to X in SQLModel" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to SQLModel but to Pydantic.
• I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• I commit to help with one of those options 👆

Example Code

def generate_astro_table(cls, primary_keys: List[str], foreign_keys: List[str]):
        # Map Pydantic types to SQLAlchemy types
        type_map = {
            int: types.INTEGER,
            str: types.VARCHAR(255),
            bool: types.BOOLEAN,
            datetime: types.DateTime,
            date: types.Date,
        }

        # Get table name from model's __tablename__ attribute
        table_name = cls.__tablename__

        # Get schema from model's metadata
        schema = cls.metadata.schema

        # Generate columns
        columns = []
        print("*" * 40)
        for name, field in cls.model_fields.items():
            actual_type = next(
                t for t in get_args(field.annotation) if t is not type(None)
            )
            column_type = type_map[actual_type]
            is_primary_key = name in primary_keys
            foreign_key = next(
                (fk for fk in foreign_keys if fk.split(".")[1] == name), None
            )

            if foreign_key:
                column = Column(
                    name,
                    column_type,
                    ForeignKey(f"{name}.id"),
                    primary_key=is_primary_key,
                )
            else:
                column = Column(name, column_type, primary_key=is_primary_key)
            columns.append(column)

        return (table_name, schema, columns)

Description

I'm using astro sdk to make some pipelines, and I use SQLModel in my application.
In astro we describe Tables and sometimes we need to list it's columns, the same way we do in SQLAlchemy.

Since I already have my tables described as SQLModel classes, I wanted to make a function to get these "Column" informations by just passing the class to it.

I managed to get table name, schema, column names and types; but for now, to mark columns as primary or foreign keys I am passing arrays with these informations.

Is there anyway for me to access the SQLAlchemy model generated?
Or maybe any method to access these informations.

I also created an issue on their repo, cause I'll probably keep working on that
astronomer/astro-sdk#2150

Operating System

Linux

Operating System Details

No response

SQLModel Version

0.0.18

Python Version

3.10.12

Additional Context

No response

[imneedle]

I ran into the same sort of problem a while back. I had resorted to using the sqlalchemy module and its inspect function.

from sqlalchemy import inspect

mapper = inspect(model)  # where model is some SQLModel, with Table=True

for column in mapper.columns:
      # Boolean values
      primary_key = column.primary_key
      nullable = column.nullable
      unique = column.unique
      index = column.index

      # Gather information on the foreign keys
      for foreign_key in column.foreign_keys:
           foreign_table = foreign_key.column.table.name
           foreign_column = foreign_key.column.name

If for whatever reason you needed to get the constraints set on the table like UniqueConstraint, CheckConstraint...

from sqlalchemy.sql.schema import UniqueConstraint, CheckConstraint


mapper = inspect(model)  # where model is some SQLModel, with Table=True
table = mapper.tables[0]
for constraint in table.constraints:
    if isinstance(constraint, UniqueConstraint):
        ...
    if isinstance(constraint, CheckConstaint):
        ...

Note: I know that sqlmodel has its own inspect function. But I am not sure how different it is to sqlalchemy's inspect function.