1
1
from bot .database_service .auth import get_db_client
2
2
from datetime import datetime , timedelta
3
3
import pytz
4
+ from bot .error_handler import DatabaseError
4
5
5
6
6
7
class FirestoreService :
@@ -14,25 +15,37 @@ def __init__(self, collection_name="users"):
14
15
15
16
# New user setup
16
17
def new_user_setup (self , telegram_id , sheet_id , telegram_username ):
17
- user_ref = self .db .collection (self .collection_name ).document (str (telegram_id ))
18
- timestamp = datetime .now (pytz .timezone ("Asia/Singapore" ))
19
- user_ref .set (
20
- {
21
- "sheet_id" : sheet_id ,
22
- "datetime_created" : timestamp ,
23
- "username" : telegram_username ,
24
- "usage_count" : 0 ,
25
- "last_accessed" : timestamp ,
26
- "hourly_accessed" : timestamp ,
27
- "overusage_count" : 0 ,
28
- }
29
- )
18
+ try :
19
+ user_ref = self .db .collection (self .collection_name ).document (
20
+ str (telegram_id )
21
+ )
22
+ timestamp = datetime .now (pytz .timezone ("Asia/Singapore" ))
23
+ user_ref .set (
24
+ {
25
+ "sheet_id" : sheet_id ,
26
+ "datetime_created" : timestamp ,
27
+ "username" : telegram_username ,
28
+ "usage_count" : 0 ,
29
+ "last_accessed" : timestamp ,
30
+ "hourly_accessed" : timestamp ,
31
+ "overusage_count" : 0 ,
32
+ }
33
+ )
34
+ except Exception as e :
35
+ raise DatabaseError (message = "Error setting up new user" , extra_info = str (e ))
30
36
31
37
# Check if user exists
32
38
def check_if_user_exists (self , telegram_id ):
33
- user_ref = self .db .collection (self .collection_name ).document (str (telegram_id ))
34
- user_doc = user_ref .get ()
35
- return user_doc .exists
39
+ try :
40
+ user_ref = self .db .collection (self .collection_name ).document (
41
+ str (telegram_id )
42
+ )
43
+ user_doc = user_ref .get ()
44
+ return user_doc .exists
45
+ except Exception as e :
46
+ raise DatabaseError (
47
+ message = "Error checking if user exists" , extra_info = str (e )
48
+ )
36
49
37
50
# Get user sheet id
38
51
def get_user_sheet_id (self , telegram_id , telegram_username ):
@@ -76,19 +89,18 @@ def get_user_sheet_id(self, telegram_id, telegram_username):
76
89
)
77
90
return user_doc .get ("sheet_id" )
78
91
except Exception as e :
79
- raise e
80
- return None
92
+ raise DatabaseError (
93
+ message = "Error getting user sheet id" , extra_info = str (e )
94
+ )
95
+ raise DatabaseError (
96
+ message = "User does not exist" , extra_info = "User does not exist"
97
+ )
81
98
82
99
# Get all user IDs
83
100
def get_all_user_id (self ):
84
- users_ref = self .db .collection (self .collection_name )
85
- user_ids = [int (user .id ) for user in users_ref .stream ()]
86
- return user_ids
87
-
88
- # Get all sheet IDs
89
- def get_all_sheet_id (self ):
90
- users_ref = self .db .collection (self .collection_name )
91
- sheet_ids = []
92
- for user in users_ref .stream ():
93
- sheet_ids .append (user .get ("sheet_id" ))
94
- return sheet_ids
101
+ try :
102
+ users_ref = self .db .collection (self .collection_name )
103
+ user_ids = [int (user .id ) for user in users_ref .stream ()]
104
+ return user_ids
105
+ except Exception as e :
106
+ raise DatabaseError (message = "Error getting all user ids" , extra_info = str (e ))
0 commit comments