Coverage for gws-app/gws/plugin/postgres/storage_provider.py: 98%

1"""PostgreSQL storage provider.""" 

2 

3from typing import Optional 

4 

5import gws 

6import gws.config.util 

7import gws.lib.sa as sa 

8from sqlalchemy.dialects.postgresql import insert as pg_insert 

9 

10from . import provider 

11 

12 

13TABLE_DDL = """ 

14 CREATE TABLE IF NOT EXISTS {table_name} ( 

15 category TEXT NOT NULL, 

16 name TEXT NOT NULL, 

17 user_uid TEXT, 

18 data TEXT, 

19 created TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, 

20 updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, 

21 PRIMARY KEY (category, name) 

22 ) 

23""" 

24"""DDL of the storage table; ``{table_name}`` is the table name. The provider does not create the table.""" 

25 

26 

27@gws.ext.config.storageProvider('postgres') 

28class Config(gws.Config): 

29 """Storage provider that keeps saved user data in a PostgreSQL table.""" 

30 

31 dbUid: Optional[str] 

32 """UID of the database provider.""" 

33 tableName: str 

34 """Table for the stored records.""" 

35 

36 

37@gws.ext.object.storageProvider('postgres') 

38class Object(gws.StorageProvider): 

39 """Storage provider that keeps records in a PostgreSQL table.""" 

40 

41 db: provider.Object 

42 """Database provider.""" 

43 tableName: str 

44 """Table for the stored records.""" 

45 

46 def configure(self): 

47 self.configure_provider() 

48 self.configure_table() 

49 

50 def configure_table(self): 

51 """Set the table name from the configuration and check that the table exists. 

52 

53 The table must have the columns given in ``TABLE_DDL``. 

54 

55 Raises: 

56 ``gws.ConfigurationError``: If the table does not exist. 

57 """ 

58 self.tableName = self.cfg('tableName') or self.cfg('_defaultTableName') 

59 if not self.db.has_table(self.tableName): 

60 raise gws.ConfigurationError(f'table {self.tableName!r} not found') 

61 

62 def configure_provider(self): 

63 """Set the database provider from ``dbUid``, or the first ``postgres`` provider. 

64 

65 Returns: 

66 ``True`` if a provider was set. 

67 

68 Raises: 

69 ``gws.Error``: If no provider is found. 

70 """ 

71 return gws.config.util.configure_database_provider_for(self) 

72 

73 def list_names(self, category): 

74 with self.db.connect() as conn: 

75 tab = self._table() 

76 rs = conn.fetch_all(tab.select().where(tab.c.category == category).with_only_columns(tab.c.name)) 

77 return sorted(rec['name'] for rec in rs) 

78 

79 def read(self, category, name): 

80 with self.db.connect() as conn: 

81 tab = self._table() 

82 rec = conn.fetch_first(tab.select().where(tab.c.category == category, tab.c.name == name).limit(1)) 

83 if rec: 

84 return gws.StorageRecord(**rec) 

85 

86 def write(self, category, name, data, user_uid): 

87 with self.db.connect() as conn: 

88 tab = self._table() 

89 sql = ( 

90 pg_insert(tab) 

91 .values(category=category, name=name, user_uid=user_uid, data=data) 

92 .on_conflict_do_update( 

93 index_elements=['category', 'name'], 

94 set_=dict( 

95 user_uid=user_uid, 

96 data=data, 

97 updated=sa.func.now(), 

98 ) 

99 ) 

100 ) 

101 conn.exec_commit(sql) 

102 

103 def delete(self, category, name): 

104 with self.db.connect() as conn: 

105 tab = self._table() 

106 sql = tab.delete().where(tab.c.category == category, tab.c.name == name) 

107 conn.exec_commit(sql) 

108 

109 def _table(self): 

110 """Return the SQLAlchemy table object of the storage table.""" 

111 return self.db.table(self.tableName)