Add foreign key validation in Collection

[gab23r]

Lets say you have this basic example:

class Country(dy.Schema):
    country_code = dy.String(primary_key=True)
    capital = dy.String()

class CountryPair(dy.Schema):
    a_country_code = dy.String(primary_key=True)
    b_country_code = dy.String(primary_key=True)
    distance = dy.Float64()

class MyCollection(dy.Collection):
    """Collection of all input tables for DMU balancing optimization."""
    country: dy.LazyFrame[Country]
    country_pair: dy.LazyFrame[CountryPair]

You might want to check/assert that a_country_code and b_country_code are present in Country.country_code.

Today you can use require_relationship_one_to_at_least_one but this will not raise a Expection but only filter.
And I would like to have something that is closer to my Collection definition. I tried to add in my Collection:

@dy.filter()
def at_least_one_diagnosis_per_invoice(self) -> pl.LazyFrame:
    return self.country.select(pl.col.country_code.name.prefix("a_")).join(
        self.country_pair,
        on="a_country_code",
        how="inner",
    )
# But it fails with: ImplementationError: Members of a collection must have an overlapping primary key but did not find any

But even if this worked, I am more interesting into an assertion rather than a filter.

Do you think we can add support for foreign_key in the same manner this is done by sqlmodel

Exmaple:

class Country(dy.Schema):
    country_code = dy.String(primary_key=True)
    capital = dy.String()

class CountryPair(dy.Schema):
    a_country_code = dy.String(primary_key=True, foreign_key="country.country_code") # reference Collection attribute
    b_country_code = dy.String(primary_key=True, foreign_key="country.country_code") # reference Collection attribute
    distance = dy.Float64()

class MyCollection(dy.Collection):
    """Collection of all input tables for DMU balancing optimization."""
    country: dy.LazyFrame[Country]
    country_pair: dy.LazyFrame[CountryPair]