def test_foreignkey(self):
ForeignKey(references=Table(), default=None, null=True)
ForeignKey(references=Table(), default=1)
with self.assertRaises(ValueError):
ForeignKey(references=Table, default="hello world")
with self.assertRaises(ValueError):
ForeignKey(references=Table, default=None, null=False)
def test_foreignkey(self):
class MyTable(Table):
pass
ForeignKey(references=MyTable, default=None, null=True)
ForeignKey(references=MyTable, default=1)
with self.assertRaises(ValueError):
ForeignKey(references=MyTable, default="hello world")
class MegaTable(Table):
"""
A table containing all of the column types, and different column kwargs.
"""
bigint_col = BigInt()
boolean_col = Boolean()
bytea_col = Bytea()
date_col = Date()
foreignkey_col = ForeignKey(SmallTable)
integer_col = Integer()
interval_col = Interval()
json_col = JSON()
jsonb_col = JSONB()
numeric_col = Numeric(digits=(5, 2))
real_col = Real()
double_precision_col = DoublePrecision()
smallint_col = SmallInt()
text_col = Text()
timestamp_col = Timestamp()
timestamptz_col = Timestamptz()
uuid_col = UUID()
varchar_col = Varchar()
unique_col = Varchar(unique=True)
null_col = Varchar(null=True)
not_null_col = Varchar(null=False)
class TokenAuth(Table):
"""
Holds randomly generated tokens.
Useful for mobile authentication, IOT etc. Session auth is recommended for
web usage.
"""
token = Varchar(default=generate_token)
user = ForeignKey(references=BaseUser)
@classmethod
async def create_token(cls,
user_id: int,
one_per_user: bool = True) -> str:
"""
Create a new token.
:param one_per_user:
If True, a ValueError is raised if a token already exists for that
user.
"""
if (await cls.exists().where(cls.user.id == user_id).run()
and one_per_user):
raise ValueError(f"User {user_id} already has a token.")
token_auth = cls(user=user_id)
await token_auth.save().run()
return t.cast(str, token_auth.token)
@classmethod
def create_token_sync(cls, user_id: int) -> str:
return run_sync(cls.create_token(user_id))
@classmethod
async def authenticate(cls, token: str) -> Select:
return cls.select(cls.user.id).where(cls.token == token).first()
@classmethod
async def authenticate_sync(cls, token: str) -> t.Optional[int]:
return run_sync(cls.authenticate(token))
@classmethod
async def get_user_id(cls, token: str) -> t.Optional[int]:
"""
Returns the user_id if the given token is valid, otherwise None.
"""
data = (await cls.select(cls.user).where(cls.token == token
).output(as_list=True
).first().run())
return data.get("user", None) if data else None
def test_self_referencing_fk(self):
"""
Make sure self-referencing foreign keys are handled correctly.
"""
MegaTable.alter().add_column("self_referencing_fk",
ForeignKey("self")).run_sync()
output_schema: OutputSchema = run_sync(get_output_schema())
# Make sure the 'references' value of the generated column is "self".
for table in output_schema.tables:
if table.__name__ == "MegaTable":
column: ForeignKey = output_schema.tables[
1].self_referencing_fk
self.assertEqual(
column._foreign_key_meta.references._meta.tablename,
MegaTable._meta.tablename,
)
self.assertEqual(column._meta.params["references"], "self")
class Concert(Table):
band = ForeignKey(Band)
venue = ForeignKey(Venue)
class Band(Table):
name = Varchar()
manager = ForeignKey(Manager)
assistant_manager = ForeignKey(Manager)
class SmallToMega(Table):
small = ForeignKey(MegaTable)
mega = ForeignKey(SmallTable)
class Movie(Table):
title = Varchar()
director = ForeignKey(Director)
def test_foreign_key(self):
self._test_add_column(
column=ForeignKey(references=Manager),
column_name="assistant_manager",
expected_value=None,
)
class TableC(Table):
table_a = ForeignKey(TableA, target_column=TableA.name)
class GenreToBand(Table):
band = ForeignKey(Band)
genre = ForeignKey(Genre)
class BandA(Table):
manager = ForeignKey(Manager, target_column=Manager.name)
class Book(Table):
name = Varchar(length=50)
writer = ForeignKey(Writer, null=True)
popularity = Integer(default=0)
class Writer(Table, tablename="schema1.writer"):
name = Varchar(length=50)
publication = ForeignKey(Publication, null=True)
class TableC(Table):
table_b = ForeignKey(TableB)
class Musician(Table):
name = Varchar()
user = ForeignKey(BaseUser)
class TableA(Table):
table_a = ForeignKey("self")
class BandB(Table):
manager = ForeignKey(Manager, target_column="name")
class TableB(Table):
table_a = ForeignKey(TableA)
class BandC(Table):
manager = ForeignKey(Manager)
class TableB(Table):
table_a = ForeignKey(TableA, target_column="name")
class Band(Table):
pk = UUID(primary=True, key=True)
name = Varchar()
manager = ForeignKey(Manager)
class TableB(Table):
table_a = ForeignKey(references=TableA)
class Band(Table):
name = Varchar(length=10)
manager = ForeignKey(Manager)
class CustomerToConcert(Table):
customer = ForeignKey(Customer)
concert = ForeignKey(Concert)
class Manager(Table):
name = Varchar()
country = ForeignKey(Country)
class Band(Table):
name = Varchar()
manager = ForeignKey(Manager)
popularity = Integer()
class GenreToBand(Table):
band = ForeignKey(Band)
genre = ForeignKey(Genre)
reason = Text(help_text="For testing additional columns on join tables.")
