How to get separately column from sqlalchemy relationship using pydantic schema

Question

I have 4 tables: Hardware, SoftwareName, SoftwareVersion, and Software.

The Software table has an one-to-many relationship with SoftwareName table and SoftwareVersion table. Finally, the Hardware model has an one-to-many relationship with Software table.

I'm trying to get just a specific column from a model relationship using Pydantic Schema.

Now I'm getting this output:

[
  {
    "id": 1,
    "hostname": "hostname2",
    "softwares": [
      {
        "id": 1,
        "software_name": {
          "id": 1,
          "name": "nginx"
        },
        "software_version": {
          "id": 1,
          "version": "2.9"
        }
      },
      {
        "id": 2,
        "software_name": {
          "id": 2,
          "name": "vim"
        },
        "software_version": {
          "id": 2,
          "version": "0.3"
        }
      },
      {
        "id": 3,
        "software_name": {
          "id": 3,
          "name": "apache"
        },
        "software_version": {
          "id": 3,
          "version": "1.0"
        }
      }
    ]
  }
]

But what I expect is this output:


[
  {
    "id": 1,
    "hostname": "hostname2",
    "softwares": [
      {
        "id": 1,
        "name": "nginx",
        "version": "2.9"
      },
      {
        "id": 2,
        "name": "vim",
        "version": "0.3"
      },
      {
        "id": 3,
        "name": "apache",
        "version": "1.0"
      }
    ]
  }
]

I have the file main.py:

   
import uvicorn
from typing import Any, Iterator, List, Optional
from faker import Faker
from fastapi import Depends, FastAPI
from pydantic import BaseModel
from sqlalchemy import Column, ForeignKey, Integer, String, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, sessionmaker, relationship
from faker.providers import DynamicProvider

software_name = DynamicProvider(
     provider_name="software_name",
     elements=["bash", "vim", "vscode", "nginx", "apache"],
)


software_version = DynamicProvider(
     provider_name="software_version",
     elements=["1.0", "2.9", "1.1", "0.3", "2.0"],
)


hardware = DynamicProvider(
     provider_name="hardware",
     elements=["hostname1", "hostname2", "hostname3", "hostname4", "hostname5"],
)

fake = Faker()

# then add new provider to faker instance
fake.add_provider(software_name)
fake.add_provider(software_version)
fake.add_provider(hardware)


engine = create_engine("sqlite:///.db", connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=True, autoflush=True, bind=engine)

Base = declarative_base(bind=engine)


class Software(Base):
    __tablename__ = 'software'

    id = Column(Integer, primary_key=True)
    hardware_id = Column(Integer, ForeignKey('hardware.id'))
    name_id = Column(Integer, ForeignKey('software_name.id'))
    version_id = Column(Integer, ForeignKey('software_version.id'))

    software_name = relationship('SoftwareName', backref='software_name')
    software_version = relationship('SoftwareVersion',
                                    backref='software_version')


class SoftwareName(Base):
    __tablename__ = 'software_name'

    id = Column(Integer, primary_key=True)
    name = Column(String)


class SoftwareVersion(Base):
    __tablename__ = 'software_version'

    id = Column(Integer, primary_key=True)
    version = Column(String)


class Hardware(Base):
    __tablename__ = "hardware"

    id = Column(Integer, primary_key=True, autoincrement=True)
    hostname = Column(String, nullable=False)

    softwares = relationship(Software)


Base.metadata.drop_all()
Base.metadata.create_all()

class BaseSchema(BaseModel):
    id: int

    class Config:
        orm_mode = True


class SoftwareNameSchema(BaseSchema):
    name: str


class SoftwareVersionSchema(BaseSchema):
    version: str


class SoftwareSchema(BaseSchema):
    software_name: SoftwareNameSchema
    software_version:  SoftwareVersionSchema


class HardwareOut(BaseSchema):
    hostname: str
    softwares: List[SoftwareSchema]


app = FastAPI()


@app.on_event("startup")
def on_startup() -> None:
    session = SessionLocal()

    for _ in range(10):
        software_list = []
        for _ in range(3):
            sn = SoftwareName(name=fake.software_name())
            sv = SoftwareVersion(version=fake.software_version())
            s = Software(software_name=sn, software_version=sv)
            software_list.append(s)

        h = Hardware(hostname=fake.hardware(), softwares=software_list)
        session.add(h)
        session.flush()

    session.close()


def get_db() -> Iterator[Session]:
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


@app.get("/hardwares", response_model=List[HardwareOut])
def get_hardwares(db: Session = Depends(get_db)) -> Any:
    return [HardwareOut.from_orm(hardware) for hardware in db.query(Hardware).all()]

How can I change the HardwareOut Schema to return what I expect?

user_16441707 · Answer

I finally got the answer I wanted.

I added 2 changes to get it:

Use the Union type from typing lib for the attributes software_name e software_version like that:
Add a Pydantic validator for each field to change the returned value, like that:

from typing import Union
from pydantic import validator

...

class SoftwareSchema(BaseSchema):
    software_name: Union[str, SoftwareNameSchema]
    software_version:  Union[str, SoftwareVersionSchema]

    @validator('software_name')
    def name_to_str(cls, v, values, **kwargs):
        return v.name if not isinstance(v, str) else v

    @validator('software_version')
    def version_to_str(cls, v, values, **kwargs):
        return v.version if not isinstance(v, str) else v

...

And the answer was this:

[
  {
    "id": 1,
    "hostname": "hostname2",
    "softwares": [
      {
        "id": 1,
        "software_name": "nginx",
        "software_version": "2.9"
      },
      {
        "id": 2,
        "software_name": "vim",
        "software_version": "0.3"
      },
      {
        "id": 3,
        "software_name": "apache",
        "software_version": "1.0"
      }
    ]
  }
]

update:

As an improvement, I add an alias for each attribute for a better semantic response. So, I change software_name to name and software_version to version. Like this:

from typing import Union
from pydantic import validator

...

class SoftwareSchema(BaseSchema):
    software_name: Union[str, SoftwareNameSchema] = Field(None, alias="name")
    software_version:  Union[str, SoftwareVersionSchema] = Field(None, alias="version")

    @validator('software_name')
    def name_to_str(cls, v, values, **kwargs):
        return v.name if not isinstance(v, str) else v

    @validator('software_version')
    def version_to_str(cls, v, values, **kwargs):
        return v.version if not isinstance(v, str) else v

...

How to get separately column from sqlalchemy relationship using pydantic schema

1 Answers