1
+ import re
2
+ from datetime import datetime
3
+ from enum import Enum
4
+ from typing import Any , List , Optional
5
+ from uuid import uuid4
6
+
7
+ from pydantic import BaseModel , Field , TypeAdapter , field_serializer
8
+ from pydantic_core import from_json
9
+ from redis .asyncio import Redis
10
+ from redis .commands .search .document import Document
11
+ from redis .commands .search .field import TextField
12
+ from redis .commands .search .query import Query
13
+ from redis .commands .search .indexDefinition import IndexDefinition , IndexType
14
+ from redis .exceptions import ResponseError
15
+
16
+ from app .logger import logger
17
+
18
+ TODOS_INDEX = "todos-idx"
19
+ TODOS_PREFIX = "todos:"
20
+
21
+
22
+ class TodoStatus (str , Enum ):
23
+ todo = 'todo'
24
+ in_progress = 'in progress'
25
+ complete = 'complete'
26
+
27
+
28
+ class Todo (BaseModel ):
29
+ name : str
30
+ status : TodoStatus
31
+ created_date : datetime = None
32
+ updated_date : datetime = None
33
+
34
+ @field_serializer ('created_date' )
35
+ def serialize_created_date (self , dt : datetime , _info ):
36
+ return dt .strftime ('%Y-%m-%dT%H:%M:%SZ' )
37
+
38
+ @field_serializer ('updated_date' )
39
+ def serialize_updated_date (self , dt : datetime , _info ):
40
+ return dt .strftime ('%Y-%m-%dT%H:%M:%SZ' )
41
+
42
+ class TodoDocument (BaseModel ):
43
+ id : str
44
+ value : Todo
45
+
46
+ # def __init__(self, doc: Any, test: Any):
47
+ # print("test")
48
+ # print(doc)
49
+ # super().__init__(id=doc.id, document=Todo(**from_json(doc.json, allow_partial=True)))
50
+
51
+ class Todos (BaseModel ):
52
+ total : int
53
+ documents : List [TodoDocument ]
54
+
55
+ class TodoStore :
56
+ """Stores todos"""
57
+
58
+ def __init__ (self , redis : Redis ):
59
+ self .redis = redis
60
+ self .INDEX = TODOS_INDEX
61
+ self .PREFIX = TODOS_PREFIX
62
+
63
+ async def initialize (self ) -> None :
64
+ await self .create_index_if_not_exists ()
65
+ return None
66
+
67
+ async def have_index (self ) -> bool :
68
+ try :
69
+ info = await self .redis .ft (self .INDEX ).info () # type: ignore
70
+ except ResponseError as e :
71
+ if "Unknown index name" in str (e ):
72
+ logger .info (f'Index { self .INDEX } does not exist' )
73
+ return False
74
+
75
+ logger .info (f"Index { self .INDEX } already exists" )
76
+ return True
77
+
78
+ async def create_index_if_not_exists (self ) -> None :
79
+ if await self .have_index ():
80
+ return None
81
+
82
+ logger .debug (f"Creating index { self .INDEX } " )
83
+
84
+ schema = (
85
+ TextField ("$.name" , as_name = "name" ),
86
+ TextField ("$.status" , as_name = "status" ),
87
+ )
88
+
89
+ try :
90
+ await self .redis .ft (self .INDEX ).create_index ( # type: ignore
91
+ schema ,
92
+ definition = IndexDefinition ( # type: ignore
93
+ prefix = [TODOS_PREFIX ], index_type = IndexType .JSON
94
+ ),
95
+ )
96
+ except Exception as e :
97
+ logger .error (f"Error setting up index { self .INDEX } : { e } " )
98
+ raise
99
+
100
+ logger .debug (f"Index { self .INDEX } created successfully" )
101
+
102
+ return None
103
+
104
+ async def drop_index (self ) -> None :
105
+ if not await self .have_index ():
106
+ return None
107
+
108
+ try :
109
+ await self .redis .ft (self .INDEX ).dropindex ()
110
+ except Exception as e :
111
+ logger .error (f"Error dropping index ${ self .INDEX } : { e } " )
112
+ raise
113
+
114
+ logger .debug (f"Index { self .INDEX } dropped successfully" )
115
+
116
+ return None
117
+
118
+ def format_id (self , id : str ) -> str :
119
+ if re .match (f'^{ self .PREFIX } ' , id ):
120
+ return id
121
+
122
+ return f'{ self .PREFIX } { id } '
123
+
124
+ def parse_todo_document (self , todo : Document ) -> TodoDocument :
125
+ return TodoDocument (id = todo .id , value = Todo (** from_json (todo .json , allow_partial = True )))
126
+
127
+ def parse_todo_documents (self , todos : list [Document ]) -> Todos :
128
+ todo_docs = [];
129
+
130
+ for doc in todos :
131
+ todo_docs .append (self .parse_todo_document (doc ))
132
+
133
+ return todo_docs
134
+
135
+ async def all (self ):
136
+ try :
137
+ result = await self .redis .ft (self .INDEX ).search ("*" )
138
+ return Todos (total = result .total , documents = self .parse_todo_documents (result .docs ))
139
+ except Exception as e :
140
+ logger .error (f"Error getting all todos: { e } " )
141
+ raise
142
+
143
+ async def one (self , id : str ) -> Todo :
144
+ id = self .format_id (id )
145
+
146
+ try :
147
+ json = await self .redis .json ().get (id )
148
+ except Exception as e :
149
+ logger .error (f"Error getting todo ${ id } : { e } " )
150
+ raise
151
+
152
+ return Todo (** json )
153
+
154
+ async def search (self , name : str | None , status : TodoStatus | None ) -> Todo :
155
+ searches = []
156
+
157
+ if name is not None and len (name ) > 0 :
158
+ searches .append (f'@name:({ name } )' )
159
+
160
+ if status is not None and len (status ) > 0 :
161
+ searches .append (f'@status:{ status .value } ' )
162
+
163
+ try :
164
+ result = await self .redis .ft (self .INDEX ).search (Query (' ' .join (searches )))
165
+ return Todos (total = result .total , documents = self .parse_todo_documents (result .docs ))
166
+ except Exception as e :
167
+ logger .error (f"Error getting todo { id } : { e } " )
168
+ raise
169
+
170
+ async def create (self , id : Optional [str ], name : Optional [str ]) -> TodoDocument :
171
+ dt = datetime .now ()
172
+
173
+ if name is None :
174
+ raise Exception ("Todo must have a name" )
175
+
176
+ if id is None :
177
+ id = str (uuid4 ())
178
+
179
+ todo = TodoDocument (** {
180
+ "id" : self .format_id (id ),
181
+ "value" : {
182
+ "name" : name ,
183
+ "status" : "todo" ,
184
+ "created_date" : dt ,
185
+ "updated_date" : dt ,
186
+ }
187
+ })
188
+
189
+ try :
190
+ result = await self .redis .json ().set (todo .id , "$" , todo .value .model_dump ())
191
+ except Exception as e :
192
+ logger .error (f'Error creating todo { todo } : { e } ' )
193
+ raise
194
+
195
+ if result != True :
196
+ raise Exception (f'Error creating todo { todo } ' )
197
+
198
+ return todo
199
+
200
+ async def update (self , id : str , status : str ) -> Todo :
201
+ dt = datetime .now ()
202
+
203
+ todo = await self .one (id )
204
+
205
+ todo .status = status
206
+ todo .updated_date = dt
207
+
208
+ try :
209
+ result = await self .redis .json ().set (self .format_id (id ), "$" , todo .model_dump ())
210
+ except Exception as e :
211
+ logger .error (f'Error updating todo { todo } : { e } ' )
212
+ raise
213
+
214
+ if result != True :
215
+ raise Exception (f'Error creating todo { todo } ' )
216
+
217
+ return todo
218
+
219
+ async def delete (self , id : str ) -> None :
220
+ try :
221
+ await self .redis .json ().delete (self .format_id (id ))
222
+ except Exception as e :
223
+ logger .error (f'Error deleting todo { id } : { e } ' )
224
+ raise
225
+
226
+ return None
227
+
228
+ async def delete_all (self ) -> None :
229
+ todos = await self .all ()
230
+
231
+ try :
232
+ for todo in todos .documents :
233
+ await self .redis .json ().delete (todo .id )
234
+ except Exception as e :
235
+ logger .error (f'Error deleting todos: { e } ' )
236
+ raise
237
+
238
+ return None
0 commit comments