-
Notifications
You must be signed in to change notification settings - Fork 8
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
1 changed file
with
51 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
from typing import Optional | ||
|
||
from sqlalchemy import types | ||
from sqlalchemy.util import generic_repr | ||
from typeid import TypeID, from_uuid | ||
|
||
|
||
class TypeIDType(types.TypeDecorator): | ||
""" | ||
A SQLAlchemy TypeDecorator that allows storing TypeIDs in the database. | ||
The prefix will not be persisted, instead the database-native UUID field will be used. | ||
At retrieval time a TypeID will be constructed based on the configured prefix and the | ||
UUID value from the database. | ||
Usage: | ||
# will result in TypeIDs such as "user_01h45ytscbebyvny4gc8cr8ma2" | ||
id = mapped_column( | ||
TypeIDType("user"), | ||
primary_key=True, | ||
default=lambda: TypeID("user") | ||
) | ||
""" | ||
impl = types.Uuid | ||
|
||
cache_ok = True | ||
|
||
prefix: Optional[str] | ||
|
||
def __init__(self, prefix: Optional[str], *args, **kwargs): | ||
self.prefix = prefix | ||
super().__init__(*args, **kwargs) | ||
|
||
def __repr__(self) -> str: | ||
# Customize __repr__ to ensure that auto-generated code e.g. from alembic includes | ||
# the right __init__ params (otherwise by default prefix will be omitted because | ||
# uuid.__init__ does not have such an argument). | ||
return generic_repr( | ||
self, | ||
to_inspect=TypeID(self.prefix), | ||
) | ||
|
||
def process_bind_param(self, value: TypeID, dialect): | ||
if self.prefix is None: | ||
assert value.prefix is None | ||
else: | ||
assert value.prefix == self.prefix | ||
|
||
return value.uuid | ||
|
||
def process_result_value(self, value, dialect): | ||
return from_uuid(value, self.prefix) |