Coverage for slidge/db/meta.py: 91%

1from __future__ import annotations 

2 

3import json 

4from typing import Union 

5 

6import sqlalchemy as sa 

7from slixmpp import JID 

8 

9 

10class JIDType(sa.TypeDecorator[JID]): 

11 """ 

12 Custom SQLAlchemy type for JIDs 

13 """ 

14 

15 impl = sa.types.TEXT 

16 cache_ok = True 

17 

18 def process_bind_param(self, value: JID | None, dialect: sa.Dialect) -> str | None: 

19 if value is None: 

20 return value 

21 return str(value) 

22 

23 def process_result_value( 

24 self, value: str | None, dialect: sa.Dialect 

25 ) -> JID | None: 

26 if value is None: 

27 return value 

28 return JID(value) 

29 

30 

31class JSONEncodedDict(sa.TypeDecorator): 

32 """ 

33 Custom SQLAlchemy type for dictionaries stored as JSON 

34 

35 Note that mutations of the dictionary are not detected by SQLAlchemy, 

36 which is why use ``attributes.flag_modified()`` in ``UserStore.update()`` 

37 """ 

38 

39 impl = sa.VARCHAR 

40 

41 cache_ok = True 

42 

43 def process_bind_param(self, value, dialect): 

44 if value is not None: 

45 value = json.dumps(value) 

46 

47 return value 

48 

49 def process_result_value(self, value, dialect): 

50 if value is not None: 

51 value = json.loads(value) 

52 return value 

53 

54 

55JSONSerializableTypes = Union[str, float, None, "JSONSerializable"] 

56JSONSerializable = dict[str, JSONSerializableTypes] 

57 

58 

59class Base(sa.orm.DeclarativeBase): 

60 type_annotation_map = {JSONSerializable: JSONEncodedDict, JID: JIDType} 

61 naming_convention = { 

62 "ix": "ix_%(column_0_label)s", 

63 "uq": "uq_%(table_name)s_%(column_0_name)s", 

64 "ck": "ck_%(table_name)s_`%(constraint_name)s`", 

65 "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", 

66 "pk": "pk_%(table_name)s", 

67 } 

68 

69 

70def get_engine(path: str) -> sa.Engine: 

71 engine = sa.create_engine(path) 

72 return engine