class Movie(Table):
name = Varchar(length=300)
rating = Real(help_text="The rating on IMDB.")
duration = Interval()
director = ForeignKey(references=Director)
oscar_nominations = Integer()
won_oscar = Boolean()
description = Text()
release_date = Timestamp()
box_office = Numeric(digits=(5, 1), help_text="In millions of US dollars.")
tags = Array(base_column=Varchar())
class Category(Table):
"""
An category table.
"""
name = Varchar(length=200)
slug = Varchar(length=200)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.name])
def test_subtract(self):
kwargs = {"class_name": "Manager", "tablename": "manager"}
name_column_1 = Varchar(unique=False)
name_column_1._meta.name = "name"
table_1 = DiffableTable(**kwargs, columns=[name_column_1])
name_column_2 = Varchar(unique=True)
name_column_2._meta.name = "name"
table_2 = DiffableTable(**kwargs, columns=[name_column_2])
delta = table_2 - table_1
self.assertEqual(delta.alter_columns[0].params, {"unique": True})
self.assertEqual(delta.alter_columns[0].old_params, {"unique": False})
def test_drop_column(self, get_migration_managers: MagicMock):
"""
Test dropping a column with MigrationManager.
"""
manager_1 = MigrationManager()
name_column = Varchar()
name_column._meta.name = "name"
manager_1.add_table(
class_name="Musician", tablename="musician", columns=[name_column]
)
asyncio.run(manager_1.run())
self.run_sync("INSERT INTO musician VALUES (default, 'Dave');")
response = self.run_sync("SELECT * FROM musician;")
self.assertEqual(response, [{"id": 1, "name": "Dave"}])
manager_2 = MigrationManager()
manager_2.drop_column(
table_class_name="Musician",
tablename="musician",
column_name="name",
)
asyncio.run(manager_2.run())
response = self.run_sync("SELECT * FROM musician;")
self.assertEqual(response, [{"id": 1}])
# Reverse
set_mock_return_value(get_migration_managers, [manager_1])
asyncio.run(manager_2.run_backwards())
response = self.run_sync("SELECT * FROM musician;")
self.assertEqual(response, [{"id": 1, "name": ""}])
class Movie(Table):
name = Varchar(length=100, required=True)
rating = Integer()
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.name])
def test_drop_table(self, get_migration_managers: MagicMock):
self.run_sync("DROP TABLE IF EXISTS musician;")
name_column = Varchar()
name_column._meta.name = "name"
manager_1 = MigrationManager(migration_id="1", app_name="music")
manager_1.add_table(
class_name="Musician", tablename="musician", columns=[name_column]
)
asyncio.run(manager_1.run())
manager_2 = MigrationManager(migration_id="2", app_name="music")
manager_2.drop_table(class_name="Musician", tablename="musician")
asyncio.run(manager_2.run())
set_mock_return_value(get_migration_managers, [manager_1])
self.assertTrue(not self.table_exists("musician"))
asyncio.run(manager_2.run_backwards())
get_migration_managers.assert_called_with(
app_name="music", max_migration_id="2", offset=-1
)
self.assertTrue(self.table_exists("musician"))
self.run_sync("DROP TABLE IF EXISTS musician;")
class Ad(Table):
title = Varchar(length=255)
slug = Varchar(length=255)
content = Text()
created = Timestamp()
view = Integer(default=0)
price = Integer()
room = Integer()
visitor = Integer()
address = Varchar(length=255)
city = Varchar(length=255)
ad_user = ForeignKey(references=User)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.title])
class Manager(Table):
name = Varchar(length=50)
touring = Boolean(default=False)
@classmethod
def get_readable(cls) -> Readable:
return Readable(template="%s", columns=[cls.name])
class Image(Table):
path = Varchar(length=255)
ad_image = ForeignKey(references=Ad)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.path])
class Task(Table):
"""
An example table.
"""
name = Varchar()
completed = Boolean(default=False)
class Venue(Table):
name = Varchar(length=100)
capacity = Integer(default=0, secret=True)
@classmethod
def get_readable(cls) -> Readable:
return Readable(template="%s", columns=[cls.name])
class Migration(Table):
name = Varchar(length=200)
app_name = Varchar(length=200)
ran_on = Timestamp(default=TimestampNow())
@classmethod
async def get_migrations_which_ran(cls,
app_name: t.Optional[str] = None
) -> t.List[str]:
"""
Returns the names of migrations which have already run, by inspecting
the database.
"""
query = cls.select(cls.name, cls.ran_on).order_by(cls.ran_on)
if app_name is not None:
query = query.where(cls.app_name == app_name)
return [i["name"] for i in await query.run()]
class Band(Table):
name = Varchar(length=50)
manager = ForeignKey(Manager, null=True)
popularity = Integer(default=0)
@classmethod
def get_readable(cls) -> Readable:
return Readable(template="%s", columns=[cls.name])
class Movie(Table, db=DB):
name = Varchar(length=300)
rating = Real()
duration = Integer()
director = ForeignKey(references=Director)
won_oscar = Boolean()
description = Text()
release_date = Timestamp()
box_office = Numeric(digits=(5, 1))
class Director(Table, help_text="The main director for a movie."):
class Gender(enum.Enum):
male = "m"
female = "f"
non_binary = "n"
name = Varchar(length=300, null=False)
years_nominated = Array(
base_column=Integer(),
help_text=(
"Which years this director was nominated for a best director "
"Oscar."),
)
gender = Varchar(length=1, choices=Gender)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.name])
class Question(Table):
"""
An question table.
"""
title = Varchar(length=200)
slug = Varchar(length=200)
description = Text()
created_at = Timestamp()
view = Integer(default=0)
question_like = Integer(default=0)
accepted_answer = Boolean(default=False)
user = ForeignKey(references=User)
category = ForeignKey(references=Category)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.title])
class Post(Table):
"""
A simple blog post.
"""
title = Varchar()
content = Text()
published = Boolean(default=False)
created_on = Timestamp()
class Shirt(Table):
"""
Used for testing columns with a choices attribute.
"""
class Size(str, Enum):
small = "s"
medium = "m"
large = "l"
size = Varchar(length=1, choices=Size, default=Size.large)
class Notification(Table):
message = Varchar(length=150)
created = Timestamp()
is_read = Boolean(default=False)
sender = ForeignKey(references=User)
recipient = ForeignKey(references=User)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.message])
class Category(Table):
"""
An Category table.
"""
name = Varchar()
description = Text()
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.name])
class Director(Table, help_text="The main director for a movie."):
name = Varchar(length=300, null=False)
years_nominated = Array(
base_column=Integer(),
help_text=(
"Which years this director was nominated for a best director "
"Oscar."
),
)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.name])
class Topic(Table):
"""
An Topic table.
"""
subject = Varchar()
created = Timestamp()
category = ForeignKey(references=Category)
topic_user = ForeignKey(references=BaseUser)
@classmethod
def get_readable(cls):
return Readable(template="%s", columns=[cls.subject])
class Movie(Table):
class Genre(int, enum.Enum):
fantasy = 1
sci_fi = 2
documentary = 3
horror = 4
action = 5
comedy = 6
romance = 7
musical = 8
name = Varchar(length=300)
rating = Real(help_text="The rating on IMDB.")
duration = Interval()
director = ForeignKey(references=Director)
oscar_nominations = Integer()
won_oscar = Boolean()
description = Text()
release_date = Timestamp()
box_office = Numeric(digits=(5, 1), help_text="In millions of US dollars.")
tags = Array(base_column=Varchar())
barcode = BigInt(default=0)
genre = SmallInt(choices=Genre, null=True)
def test_create_table_class(self):
"""
Make sure a basic `Table` can be created successfully.
"""
_Table = create_table_class(class_name="MyTable")
self.assertEqual(_Table._meta.tablename, "my_table")
_Table = create_table_class(class_name="MyTable",
class_kwargs={"tablename": "my_table_1"})
self.assertEqual(_Table._meta.tablename, "my_table_1")
column = Varchar()
_Table = create_table_class(class_name="MyTable",
class_members={"name": column})
self.assertIn(column, _Table._meta.columns)
def test_integer_to_varchar(self):
"""
Test converting an Integer column to Varchar.
"""
self.insert_row()
alter_query = Band.alter().set_column_type(old_column=Band.popularity,
new_column=Varchar())
alter_query.run_sync()
self.assertEqual(
self.get_postgres_column_type(tablename="band",
column_name="popularity"),
"CHARACTER VARYING",
)
popularity = (Band.select(
Band.popularity).first().run_sync()["popularity"])
self.assertEqual(popularity, "1000")
def test_add_table(self):
"""
Test adding a table to a MigrationManager.
"""
self.run_sync("DROP TABLE IF EXISTS musician;")
manager = MigrationManager()
name_column = Varchar()
name_column._meta.name = "name"
manager.add_table(
class_name="Musician", tablename="musician", columns=[name_column]
)
asyncio.run(manager.run())
self.run_sync("INSERT INTO musician VALUES (default, 'Bob Jones');")
response = self.run_sync("SELECT * FROM musician;")
self.assertEqual(response, [{"id": 1, "name": "Bob Jones"}])
# Reverse
asyncio.run(manager.run_backwards())
self.assertEqual(self.table_exists("musician"), False)
self.run_sync("DROP TABLE IF EXISTS musician;")
class BaseUser(Table, tablename="piccolo_user"):
"""
Provides a basic user, with authentication support.
"""
id: Serial
username = Varchar(length=100, unique=True)
password = Secret(length=255)
first_name = Varchar(null=True)
last_name = Varchar(null=True)
email = Varchar(length=255, unique=True)
active = Boolean(default=False)
admin = Boolean(
default=False, help_text="An admin can log into the Piccolo admin GUI."
)
superuser = Boolean(
default=False,
help_text=(
"If True, this user can manage other users's passwords in the "
"Piccolo admin GUI."
),
)
last_login = Timestamp(
null=True,
default=None,
required=False,
help_text="When this user last logged in.",
)
_min_password_length = 6
_max_password_length = 128
def __init__(self, **kwargs):
# Generating passwords upfront is expensive, so might need reworking.
password = kwargs.get("password", None)
if password:
if not password.startswith("pbkdf2_sha256"):
kwargs["password"] = self.__class__.hash_password(password)
super().__init__(**kwargs)
@classmethod
def get_salt(cls):
return secrets.token_hex(16)
@classmethod
def get_readable(cls) -> Readable:
"""
Used to get a readable string, representing a table row.
"""
return Readable(template="%s", columns=[cls.username])
###########################################################################
@classmethod
def update_password_sync(cls, user: t.Union[str, int], password: str):
"""
A sync equivalent of :meth:`update_password`.
"""
return run_sync(cls.update_password(user, password))
@classmethod
async def update_password(cls, user: t.Union[str, int], password: str):
"""
The password is the raw password string e.g. ``'password123'``.
The user can be a user ID, or a username.
"""
if isinstance(user, str):
clause = cls.username == user
elif isinstance(user, int):
clause = cls.id == user
else:
raise ValueError(
"The `user` arg must be a user id, or a username."
)
password = cls.hash_password(password)
await cls.update({cls.password: password}).where(clause).run()
###########################################################################
@classmethod
def hash_password(
cls, password: str, salt: str = "", iterations: int = 10000
) -> str:
"""
Hashes the password, ready for storage, and for comparing during
login.
:raises ValueError:
If an excessively long password is provided.
"""
if len(password) > cls._max_password_length:
logger.warning("Excessively long password provided.")
raise ValueError("The password is too long.")
if salt == "":
salt = cls.get_salt()
hashed = hashlib.pbkdf2_hmac(
"sha256",
bytes(password, encoding="utf-8"),
bytes(salt, encoding="utf-8"),
iterations,
).hex()
return f"pbkdf2_sha256${iterations}${salt}${hashed}"
def __setattr__(self, name: str, value: t.Any):
"""
Make sure that if the password is set, it's stored in a hashed form.
"""
if name == "password" and not value.startswith("pbkdf2_sha256"):
value = self.__class__.hash_password(value)
super().__setattr__(name, value)
@classmethod
def split_stored_password(cls, password: str) -> t.List[str]:
elements = password.split("$")
if len(elements) != 4:
raise ValueError("Unable to split hashed password")
return elements
###########################################################################
@classmethod
def login_sync(cls, username: str, password: str) -> t.Optional[int]:
"""
A sync equivalent of :meth:`login`.
"""
return run_sync(cls.login(username, password))
@classmethod
async def login(cls, username: str, password: str) -> t.Optional[int]:
"""
Make sure the user exists and the password is valid. If so, the
``last_login`` value is updated in the database.
:returns:
The id of the user if a match is found, otherwise ``None``.
"""
if len(username) > cls.username.length:
logger.warning("Excessively long username provided.")
return None
if len(password) > cls._max_password_length:
logger.warning("Excessively long password provided.")
return None
response = (
await cls.select(cls._meta.primary_key, cls.password)
.where(cls.username == username)
.first()
.run()
)
if not response:
# No match found
return None
stored_password = response["password"]
algorithm, iterations, salt, hashed = cls.split_stored_password(
stored_password
)
if (
cls.hash_password(password, salt, int(iterations))
== stored_password
):
await cls.update({cls.last_login: datetime.datetime.now()}).where(
cls.username == username
)
return response["id"]
else:
return None
###########################################################################
@classmethod
def create_user_sync(
cls, username: str, password: str, **extra_params
) -> BaseUser:
"""
A sync equivalent of :meth:`create_user`.
"""
return run_sync(
cls.create_user(
username=username, password=password, **extra_params
)
)
@classmethod
async def create_user(
cls, username: str, password: str, **extra_params
) -> BaseUser:
"""
Creates a new user, and saves it in the database. It is recommended to
use this rather than instantiating and saving ``BaseUser`` directly, as
we add extra validation.
:raises ValueError:
If the username or password is invalid.
:returns:
The created ``BaseUser`` instance.
"""
if not username:
raise ValueError("A username must be provided.")
if not password:
raise ValueError("A password must be provided.")
if len(password) < cls._min_password_length:
raise ValueError("The password is too short.")
if len(password) > cls._max_password_length:
raise ValueError("The password is too long.")
if password.startswith("pbkdf2_sha256"):
logger.warning(
"Tried to create a user with an already hashed password."
)
raise ValueError("Do not pass a hashed password.")
user = cls(username=username, password=password, **extra_params)
await user.save()
return user
class BaseUser(Table, tablename="piccolo_user"):
"""
Provides a basic user, with authentication support.
"""
username = Varchar(length=100, unique=True)
password = Secret(length=255)
first_name = Varchar(null=True)
last_name = Varchar(null=True)
email = Varchar(length=255, unique=True)
active = Boolean(default=False)
admin = Boolean(default=False)
def __init__(self, **kwargs):
"""
Generating passwords upfront is expensive, so might need reworking.
"""
password = kwargs.get("password", None)
if password:
kwargs["password"] = self.__class__.hash_password(password)
super().__init__(**kwargs)
@classmethod
def get_salt(cls):
return secrets.token_hex(16)
@classmethod
def get_readable(cls) -> Readable:
"""
Used to get a readable string, representing a table row.
"""
return Readable(template="%s", columns=[cls.username])
###########################################################################
@classmethod
def update_password_sync(cls, user: t.Union[str, int], password: str):
return run_sync(cls.update_password(user, password))
@classmethod
async def update_password(cls, user: t.Union[str, int], password: str):
"""
The password is the raw password string e.g. password123.
The user can be a user ID, or a username.
"""
if isinstance(user, str):
clause = cls.username == user
elif isinstance(user, int):
clause = cls.id == user # type: ignore
else:
raise ValueError(
"The `user` arg must be a user id, or a username."
)
password = cls.hash_password(password)
await cls.update().values({cls.password: password}).where(clause).run()
###########################################################################
@classmethod
def hash_password(
cls, password: str, salt: str = "", iterations: int = 10000
) -> str:
"""
Hashes the password, ready for storage, and for comparing during
login.
"""
if salt == "":
salt = cls.get_salt()
hashed = hashlib.pbkdf2_hmac(
"sha256",
bytes(password, encoding="utf-8"),
bytes(salt, encoding="utf-8"),
iterations,
).hex()
return f"pbkdf2_sha256${iterations}${salt}${hashed}"
def __setattr__(self, name: str, value: t.Any):
"""
Make sure that if the password is set, it's stored in a hashed form.
"""
if name == "password":
if not value.startswith("pbkdf2_sha256"):
value = self.__class__.hash_password(value)
super().__setattr__(name, value)
@classmethod
def split_stored_password(cls, password: str) -> t.List[str]:
elements = password.split("$")
if len(elements) != 4:
raise ValueError("Unable to split hashed password")
return elements
@classmethod
def login_sync(cls, username: str, password: str) -> t.Optional[int]:
"""
Returns the user_id if a match is found.
"""
return run_sync(cls.login(username, password))
@classmethod
async def login(cls, username: str, password: str) -> t.Optional[int]:
"""
Returns the user_id if a match is found.
"""
query = (
cls.select()
.columns(cls.id, cls.password)
.where((cls.username == username))
.first()
)
response = await query.run()
if not response:
# No match found
return None
stored_password = response["password"]
algorithm, iterations, salt, hashed = cls.split_stored_password(
stored_password
)
if (
cls.hash_password(password, salt, int(iterations))
== stored_password
):
return response["id"]
else:
return None
class SessionsBase(Table, tablename="sessions"):
"""
Use this table, or inherit from it, to create a session store.
"""
#: Stores the session token.
token: Varchar = Varchar(length=100, null=False)
#: Stores the user ID.
user_id: Integer = Integer(null=False)
#: Stores the expiry date for this session.
expiry_date: Timestamp = Timestamp(default=TimestampOffset(hours=1),
null=False)
#: We set a hard limit on the expiry date - it can keep on getting extended
#: up until this value, after which it's best to invalidate it, and either
#: require login again, or just create a new session token.
max_expiry_date: Timestamp = Timestamp(default=TimestampOffset(days=7),
null=False)
@classmethod
async def create_session(
cls,
user_id: int,
expiry_date: t.Optional[datetime] = None,
max_expiry_date: t.Optional[datetime] = None,
) -> SessionsBase:
"""
Creates a session in the database.
"""
while True:
token = secrets.token_urlsafe(nbytes=32)
if not await cls.exists().where(cls.token == token).run():
break
session = cls(token=token, user_id=user_id)
if expiry_date:
session.expiry_date = expiry_date
if max_expiry_date:
session.max_expiry_date = max_expiry_date
await session.save().run()
return session
@classmethod
def create_session_sync(
cls,
user_id: int,
expiry_date: t.Optional[datetime] = None) -> SessionsBase:
"""
A sync equivalent of :meth:`create_session`.
"""
return run_sync(cls.create_session(user_id, expiry_date))
@classmethod
async def get_user_id(
cls,
token: str,
increase_expiry: t.Optional[timedelta] = None) -> t.Optional[int]:
"""
Returns the ``user_id`` if the given token is valid, otherwise
``None``.
:param increase_expiry:
If set, the ``expiry_date`` will be increased by the given amount
if it's close to expiring. If it has already expired, nothing
happens. The ``max_expiry_date`` remains the same, so there's a
hard limit on how long a session can be used for.
"""
session: SessionsBase = (await cls.objects().where(cls.token == token
).first().run())
if not session:
return None
now = datetime.now()
if (session.expiry_date > now) and (session.max_expiry_date > now):
if increase_expiry and (t.cast(datetime, session.expiry_date) - now
< increase_expiry):
session.expiry_date = (t.cast(datetime, session.expiry_date) +
increase_expiry)
await session.save().run()
return t.cast(t.Optional[int], session.user_id)
else:
return None
@classmethod
def get_user_id_sync(cls, token: str) -> t.Optional[int]:
"""
A sync wrapper around :meth:`get_user_id`.
"""
return run_sync(cls.get_user_id(token))
@classmethod
async def remove_session(cls, token: str):
"""
Deletes a matching session from the database.
"""
await cls.delete().where(cls.token == token).run()
@classmethod
def remove_session_sync(cls, token: str):
"""
A sync wrapper around :meth:`remove_session`.
"""
return run_sync(cls.remove_session(token))
class BandMember(Table):
name = Varchar(length=50, index=True)