trovesuite 1.0.9__py3-none-any.whl → 1.0.11__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- trovesuite/__init__.py +10 -5
- trovesuite/auth/__init__.py +3 -5
- trovesuite/auth/auth_base.py +1 -0
- trovesuite/auth/auth_controller.py +6 -5
- trovesuite/auth/auth_read_dto.py +4 -3
- trovesuite/auth/auth_service.py +61 -51
- trovesuite/auth/auth_write_dto.py +2 -1
- trovesuite/configs/database.py +130 -36
- trovesuite/configs/settings.py +36 -132
- trovesuite/entities/health.py +4 -4
- trovesuite/notification/__init__.py +14 -0
- trovesuite/notification/notification_base.py +13 -0
- trovesuite/notification/notification_controller.py +21 -0
- trovesuite/notification/notification_read_dto.py +21 -0
- trovesuite/notification/notification_service.py +73 -0
- trovesuite/notification/notification_write_dto.py +21 -0
- trovesuite/storage/__init__.py +42 -0
- trovesuite/storage/storage_base.py +63 -0
- trovesuite/storage/storage_controller.py +198 -0
- trovesuite/storage/storage_read_dto.py +74 -0
- trovesuite/storage/storage_service.py +529 -0
- trovesuite/storage/storage_write_dto.py +70 -0
- {trovesuite-1.0.9.dist-info → trovesuite-1.0.11.dist-info}/METADATA +188 -11
- trovesuite-1.0.11.dist-info/RECORD +33 -0
- trovesuite-1.0.9.dist-info/RECORD +0 -21
- {trovesuite-1.0.9.dist-info → trovesuite-1.0.11.dist-info}/WHEEL +0 -0
- {trovesuite-1.0.9.dist-info → trovesuite-1.0.11.dist-info}/licenses/LICENSE +0 -0
- {trovesuite-1.0.9.dist-info → trovesuite-1.0.11.dist-info}/top_level.txt +0 -0
trovesuite/__init__.py
CHANGED
|
@@ -1,16 +1,21 @@
|
|
|
1
1
|
"""
|
|
2
|
-
TroveSuite
|
|
2
|
+
TroveSuite Package
|
|
3
3
|
|
|
4
|
-
A comprehensive authentication and
|
|
5
|
-
Provides JWT token validation, user authorization,
|
|
4
|
+
A comprehensive authentication, authorization, notification, and storage service for ERP systems.
|
|
5
|
+
Provides JWT token validation, user authorization, permission checking, notification capabilities,
|
|
6
|
+
and Azure Storage blob management.
|
|
6
7
|
"""
|
|
7
8
|
|
|
8
9
|
from .auth import AuthService
|
|
10
|
+
from .notification import NotificationService
|
|
11
|
+
from .storage import StorageService
|
|
9
12
|
|
|
10
|
-
__version__ = "1.0.
|
|
13
|
+
__version__ = "1.0.7"
|
|
11
14
|
__author__ = "Bright Debrah Owusu"
|
|
12
15
|
__email__ = "owusu.debrah@deladetech.com"
|
|
13
16
|
|
|
14
17
|
__all__ = [
|
|
15
|
-
"AuthService"
|
|
18
|
+
"AuthService",
|
|
19
|
+
"NotificationService",
|
|
20
|
+
"StorageService"
|
|
16
21
|
]
|
trovesuite/auth/__init__.py
CHANGED
|
@@ -5,12 +5,10 @@ Authentication and authorization services for ERP systems.
|
|
|
5
5
|
"""
|
|
6
6
|
|
|
7
7
|
from .auth_service import AuthService
|
|
8
|
-
from .
|
|
9
|
-
from .auth_read_dto import AuthServiceReadDto, AuthControllerReadDto
|
|
8
|
+
from .auth_write_dto import AuthServiceWriteDto
|
|
10
9
|
|
|
11
10
|
__all__ = [
|
|
12
11
|
"AuthService",
|
|
13
|
-
"
|
|
14
|
-
"AuthServiceReadDto",
|
|
15
|
-
"AuthControllerReadDto"
|
|
12
|
+
"AuthServiceWriteDto"
|
|
16
13
|
]
|
|
14
|
+
|
trovesuite/auth/auth_base.py
CHANGED
|
@@ -1,10 +1,11 @@
|
|
|
1
1
|
from fastapi import APIRouter
|
|
2
|
-
from
|
|
3
|
-
from
|
|
4
|
-
from
|
|
2
|
+
from .auth_write_dto import AuthControllerWriteDto
|
|
3
|
+
from .auth_read_dto import AuthControllerReadDto
|
|
4
|
+
from .auth_service import AuthService
|
|
5
|
+
from ..entities.sh_response import Respons
|
|
5
6
|
|
|
6
|
-
auth_router = APIRouter()
|
|
7
|
+
auth_router = APIRouter(tags=["Auth"])
|
|
7
8
|
|
|
8
|
-
@auth_router.post("/auth", response_model=AuthControllerReadDto)
|
|
9
|
+
@auth_router.post("/auth", response_model=Respons[AuthControllerReadDto])
|
|
9
10
|
async def authorize(data: AuthControllerWriteDto):
|
|
10
11
|
return AuthService.authorize(data=data)
|
trovesuite/auth/auth_read_dto.py
CHANGED
|
@@ -3,16 +3,17 @@ from pydantic import BaseModel
|
|
|
3
3
|
|
|
4
4
|
class AuthControllerReadDto(BaseModel):
|
|
5
5
|
org_id: Optional[str] = None
|
|
6
|
-
bus_id: Optional[str] = None
|
|
7
|
-
app_id: Optional[str] = None
|
|
6
|
+
bus_id: Optional[str] = None
|
|
7
|
+
app_id: Optional[str] = None
|
|
8
8
|
shared_resource_id: Optional[str] = None
|
|
9
9
|
user_id: Optional[str] = None
|
|
10
10
|
group_id: Optional[str] = None
|
|
11
11
|
role_id: Optional[str] = None
|
|
12
12
|
tenant_id: Optional[str] = None
|
|
13
13
|
permissions: Optional[List[str]] = None
|
|
14
|
-
shared_resource_id: Optional[str] = None
|
|
15
14
|
resource_id: Optional[str] = None
|
|
15
|
+
resource_type: Optional[str] = None
|
|
16
16
|
|
|
17
17
|
class AuthServiceReadDto(AuthControllerReadDto):
|
|
18
18
|
pass
|
|
19
|
+
|
trovesuite/auth/auth_service.py
CHANGED
|
@@ -45,7 +45,7 @@ class AuthService:
|
|
|
45
45
|
|
|
46
46
|
@staticmethod
|
|
47
47
|
def authorize(data: AuthServiceWriteDto) -> Respons[AuthServiceReadDto]:
|
|
48
|
-
|
|
48
|
+
|
|
49
49
|
user_id: str = data.user_id
|
|
50
50
|
tenant_id: str = data.tenant_id
|
|
51
51
|
|
|
@@ -59,7 +59,7 @@ class AuthService:
|
|
|
59
59
|
status_code=400,
|
|
60
60
|
error="INVALID_USER_ID"
|
|
61
61
|
)
|
|
62
|
-
|
|
62
|
+
|
|
63
63
|
if not tenant_id or not isinstance(tenant_id, str):
|
|
64
64
|
return Respons[AuthServiceReadDto](
|
|
65
65
|
detail="Invalid tenant_id: must be a non-empty string",
|
|
@@ -68,7 +68,7 @@ class AuthService:
|
|
|
68
68
|
status_code=400,
|
|
69
69
|
error="INVALID_TENANT_ID"
|
|
70
70
|
)
|
|
71
|
-
|
|
71
|
+
|
|
72
72
|
try:
|
|
73
73
|
|
|
74
74
|
is_tenant_verified = DatabaseManager.execute_query(
|
|
@@ -85,7 +85,7 @@ class AuthService:
|
|
|
85
85
|
status_code=404,
|
|
86
86
|
error="TENANT_NOT_FOUND"
|
|
87
87
|
)
|
|
88
|
-
|
|
88
|
+
|
|
89
89
|
if not is_tenant_verified[0]['is_verified']:
|
|
90
90
|
logger.warning("Login failed - tenant not verified for user: %s, tenant: %s", user_id, tenant_id)
|
|
91
91
|
return Respons[AuthServiceReadDto](
|
|
@@ -96,14 +96,36 @@ class AuthService:
|
|
|
96
96
|
error="TENANT_NOT_VERIFIED"
|
|
97
97
|
)
|
|
98
98
|
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
WHERE (delete_status = 'NOT_DELETED' AND is_active = true ) AND user_id = %s""",
|
|
104
|
-
(user_id,),
|
|
99
|
+
# 1️⃣ Get all groups the user belongs to
|
|
100
|
+
user_groups = DatabaseManager.execute_query(
|
|
101
|
+
f"""SELECT group_id FROM "{tenant_id}".{db_settings.TENANT_USER_GROUPS_TABLE}
|
|
102
|
+
WHERE delete_status = 'NOT_DELETED' AND is_active = true AND user_id = %s""",(user_id,),
|
|
105
103
|
)
|
|
106
104
|
|
|
105
|
+
# 2️⃣ Prepare list of group_ids
|
|
106
|
+
group_ids = [g["group_id"] for g in user_groups] if user_groups else []
|
|
107
|
+
|
|
108
|
+
# 3️⃣ Get login settings - check user-level first, then group-level
|
|
109
|
+
if group_ids:
|
|
110
|
+
login_settings_details = DatabaseManager.execute_query(
|
|
111
|
+
f"""SELECT user_id, group_id, is_suspended, can_always_login,
|
|
112
|
+
is_multi_factor_enabled, is_login_before, working_days,
|
|
113
|
+
login_on, logout_on FROM "{tenant_id}".{db_settings.TENANT_LOGIN_SETTINGS_TABLE}
|
|
114
|
+
WHERE (delete_status = 'NOT_DELETED' AND is_active = true )
|
|
115
|
+
AND (user_id = %s OR group_id = ANY(%s))
|
|
116
|
+
ORDER BY user_id NULLS LAST
|
|
117
|
+
LIMIT 1""",
|
|
118
|
+
(user_id, group_ids),
|
|
119
|
+
)
|
|
120
|
+
else:
|
|
121
|
+
login_settings_details = DatabaseManager.execute_query(
|
|
122
|
+
f"""SELECT user_id, group_id, is_suspended, can_always_login,
|
|
123
|
+
is_multi_factor_enabled, is_login_before, working_days,
|
|
124
|
+
login_on, logout_on FROM "{tenant_id}".{db_settings.TENANT_LOGIN_SETTINGS_TABLE}
|
|
125
|
+
WHERE (delete_status = 'NOT_DELETED' AND is_active = true ) AND user_id = %s""",
|
|
126
|
+
(user_id,),
|
|
127
|
+
)
|
|
128
|
+
|
|
107
129
|
if not login_settings_details or len(login_settings_details) == 0:
|
|
108
130
|
logger.warning("Authorization failed - user not found: %s in tenant: %s", user_id, tenant_id)
|
|
109
131
|
return Respons[AuthServiceReadDto](
|
|
@@ -126,17 +148,17 @@ class AuthService:
|
|
|
126
148
|
|
|
127
149
|
if not login_settings_details[0]['can_always_login']:
|
|
128
150
|
current_day = datetime.now().strftime("%A").upper()
|
|
129
|
-
|
|
151
|
+
|
|
130
152
|
if current_day not in login_settings_details[0]['working_days']:
|
|
131
153
|
logger.warning("Authorization failed - outside working days for user: %s checking custom login period", user_id)
|
|
132
|
-
|
|
154
|
+
|
|
133
155
|
# Get current datetime (full date and time) with timezone
|
|
134
156
|
current_datetime = datetime.now(timezone.utc).replace(microsecond=0, second=0)
|
|
135
|
-
|
|
157
|
+
|
|
136
158
|
# Get from database (should already be datetime objects)
|
|
137
159
|
login_on = login_settings_details[0]['login_on']
|
|
138
160
|
logout_on = login_settings_details[0]['logout_on']
|
|
139
|
-
|
|
161
|
+
|
|
140
162
|
# Set defaults if None (with timezone awareness)
|
|
141
163
|
if not login_on:
|
|
142
164
|
login_on = datetime.min.replace(tzinfo=timezone.utc)
|
|
@@ -153,27 +175,19 @@ class AuthService:
|
|
|
153
175
|
status_code=403,
|
|
154
176
|
error="LOGIN_TIME_RESTRICTED"
|
|
155
177
|
)
|
|
156
|
-
|
|
157
|
-
# 1️⃣ Get all groups the user belongs to
|
|
158
|
-
user_groups = DatabaseManager.execute_query(
|
|
159
|
-
f"""SELECT group_id FROM "{tenant_id}".{db_settings.USER_GROUPS_TABLE}
|
|
160
|
-
WHERE delete_status = 'NOT_DELETED' AND is_active = true AND user_id = %s""",(user_id,),
|
|
161
|
-
)
|
|
162
|
-
|
|
163
|
-
# 2️⃣ Prepare list of group_ids
|
|
164
|
-
group_ids = [g["group_id"] for g in user_groups] if user_groups else []
|
|
165
178
|
|
|
166
|
-
#
|
|
179
|
+
# 4️⃣ Build query dynamically to include groups (if any) + user
|
|
180
|
+
# ⚠️ CHANGED: Simplified to new schema - only select user_id, group_id, role_id, resource_type
|
|
167
181
|
if group_ids:
|
|
168
182
|
get_user_roles = DatabaseManager.execute_query(
|
|
169
183
|
f"""
|
|
170
|
-
SELECT DISTINCT ON (
|
|
171
|
-
|
|
172
|
-
FROM "{tenant_id}".{db_settings.
|
|
184
|
+
SELECT DISTINCT ON (group_id, user_id, role_id)
|
|
185
|
+
group_id, user_id, role_id, resource_type
|
|
186
|
+
FROM "{tenant_id}".{db_settings.TENANT_ASSIGN_ROLES_TABLE}
|
|
173
187
|
WHERE delete_status = 'NOT_DELETED'
|
|
174
188
|
AND is_active = true
|
|
175
189
|
AND (user_id = %s OR group_id = ANY(%s))
|
|
176
|
-
ORDER BY
|
|
190
|
+
ORDER BY group_id, user_id, role_id;
|
|
177
191
|
""",
|
|
178
192
|
(user_id, group_ids),
|
|
179
193
|
)
|
|
@@ -181,13 +195,13 @@ class AuthService:
|
|
|
181
195
|
# No groups, just check roles for user
|
|
182
196
|
get_user_roles = DatabaseManager.execute_query(
|
|
183
197
|
f"""
|
|
184
|
-
SELECT DISTINCT ON (
|
|
185
|
-
|
|
186
|
-
FROM "{tenant_id}".{db_settings.
|
|
198
|
+
SELECT DISTINCT ON (user_id, role_id)
|
|
199
|
+
user_id, role_id, resource_type
|
|
200
|
+
FROM "{tenant_id}".{db_settings.TENANT_ASSIGN_ROLES_TABLE}
|
|
187
201
|
WHERE delete_status = 'NOT_DELETED'
|
|
188
202
|
AND is_active = true
|
|
189
203
|
AND user_id = %s
|
|
190
|
-
ORDER BY
|
|
204
|
+
ORDER BY user_id, role_id;
|
|
191
205
|
""",
|
|
192
206
|
(user_id,),
|
|
193
207
|
)
|
|
@@ -196,7 +210,7 @@ class AuthService:
|
|
|
196
210
|
get_user_roles_with_tenant_and_permissions = []
|
|
197
211
|
for role in get_user_roles:
|
|
198
212
|
permissions = DatabaseManager.execute_query(
|
|
199
|
-
f"""SELECT permission_id FROM {db_settings.
|
|
213
|
+
f"""SELECT permission_id FROM {db_settings.MAIN_ROLE_PERMISSIONS_TABLE} WHERE role_id = %s""",
|
|
200
214
|
params=(role["role_id"],),)
|
|
201
215
|
|
|
202
216
|
role_dict = {**role, "tenant_id": tenant_id, "permissions": [p['permission_id'] for p in permissions]}
|
|
@@ -224,31 +238,27 @@ class AuthService:
|
|
|
224
238
|
status_code=500,
|
|
225
239
|
error="Authorization check failed due to an internal error"
|
|
226
240
|
)
|
|
227
|
-
|
|
241
|
+
|
|
228
242
|
@staticmethod
|
|
229
|
-
def check_permission(users_data: list, action=None,
|
|
230
|
-
resource_id=None, shared_resource_id=None) -> bool:
|
|
243
|
+
def check_permission(users_data: list, action=None, resource_type=None) -> bool:
|
|
231
244
|
"""
|
|
232
|
-
Check if user has a given permission (action)
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
245
|
+
Check if user has a given permission (action).
|
|
246
|
+
|
|
247
|
+
Args:
|
|
248
|
+
users_data: List of user authorization data containing roles and permissions
|
|
249
|
+
action: The permission/action to check for
|
|
250
|
+
resource_type: Optional resource type filter (e.g., 'rt-user', 'rt-group')
|
|
251
|
+
|
|
252
|
+
Returns:
|
|
253
|
+
bool: True if user has the permission, False otherwise
|
|
236
254
|
"""
|
|
237
255
|
for user_data in users_data:
|
|
238
|
-
# Check
|
|
239
|
-
if user_data.
|
|
240
|
-
continue
|
|
241
|
-
if user_data.bus_id not in (None, bus_id):
|
|
242
|
-
continue
|
|
243
|
-
if user_data.app_id not in (None, app_id):
|
|
244
|
-
continue
|
|
245
|
-
if user_data.resource_id not in (None, resource_id):
|
|
246
|
-
continue
|
|
247
|
-
if user_data.shared_resource_id not in (None, shared_resource_id):
|
|
256
|
+
# Check resource_type if specified
|
|
257
|
+
if resource_type and user_data.resource_type and user_data.resource_type != resource_type:
|
|
248
258
|
continue
|
|
249
259
|
|
|
250
260
|
# Check if the permission exists
|
|
251
|
-
if action in user_data.permissions:
|
|
261
|
+
if action and action in user_data.permissions:
|
|
252
262
|
return True
|
|
253
263
|
|
|
254
264
|
return False
|
trovesuite/configs/database.py
CHANGED
|
@@ -13,24 +13,32 @@ logger = get_logger("database")
|
|
|
13
13
|
|
|
14
14
|
# Database connection pool
|
|
15
15
|
_connection_pool: Optional[psycopg2.pool.ThreadedConnectionPool] = None
|
|
16
|
-
_sqlmodel_engine = None
|
|
17
16
|
|
|
18
17
|
|
|
19
18
|
class DatabaseConfig:
|
|
20
19
|
"""Database configuration and connection management"""
|
|
21
|
-
|
|
20
|
+
|
|
22
21
|
def __init__(self):
|
|
23
22
|
self.settings = db_settings
|
|
23
|
+
self.database_url = self.settings.database_url
|
|
24
24
|
self.pool_size = 5
|
|
25
25
|
self.max_overflow = 10
|
|
26
|
-
|
|
27
|
-
@property
|
|
28
|
-
def database_url(self):
|
|
29
|
-
"""Get database URL (lazy evaluation)"""
|
|
30
|
-
return self.settings.database_url
|
|
31
|
-
|
|
26
|
+
|
|
32
27
|
def get_connection_params(self) -> dict:
|
|
33
28
|
"""Get database connection parameters"""
|
|
29
|
+
if self.settings.DATABASE_URL:
|
|
30
|
+
# Use full DATABASE_URL if available
|
|
31
|
+
return {
|
|
32
|
+
"dsn": self.settings.DATABASE_URL,
|
|
33
|
+
"cursor_factory": RealDictCursor,
|
|
34
|
+
"keepalives": 1,
|
|
35
|
+
"keepalives_idle": 30,
|
|
36
|
+
"keepalives_interval": 10,
|
|
37
|
+
"keepalives_count": 5,
|
|
38
|
+
"connect_timeout": 10
|
|
39
|
+
}
|
|
40
|
+
|
|
41
|
+
# fallback to individual DB_* variables
|
|
34
42
|
return {
|
|
35
43
|
"host": self.settings.DB_HOST,
|
|
36
44
|
"port": self.settings.DB_PORT,
|
|
@@ -38,9 +46,14 @@ class DatabaseConfig:
|
|
|
38
46
|
"user": self.settings.DB_USER,
|
|
39
47
|
"password": self.settings.DB_PASSWORD,
|
|
40
48
|
"cursor_factory": RealDictCursor,
|
|
41
|
-
"application_name": f"{self.settings.APP_NAME}_{self.settings.ENVIRONMENT}"
|
|
49
|
+
"application_name": f"{self.settings.APP_NAME}_{self.settings.ENVIRONMENT}",
|
|
50
|
+
"keepalives": 1,
|
|
51
|
+
"keepalives_idle": 30,
|
|
52
|
+
"keepalives_interval": 10,
|
|
53
|
+
"keepalives_count": 5,
|
|
54
|
+
"connect_timeout": 10
|
|
42
55
|
}
|
|
43
|
-
|
|
56
|
+
|
|
44
57
|
def create_connection_pool(self) -> psycopg2.pool.ThreadedConnectionPool:
|
|
45
58
|
"""Create a connection pool for psycopg2"""
|
|
46
59
|
try:
|
|
@@ -54,7 +67,7 @@ class DatabaseConfig:
|
|
|
54
67
|
except Exception as e:
|
|
55
68
|
logger.error(f"Failed to create database connection pool: {str(e)}")
|
|
56
69
|
raise
|
|
57
|
-
|
|
70
|
+
|
|
58
71
|
def test_connection(self) -> bool:
|
|
59
72
|
"""Test database connection"""
|
|
60
73
|
try:
|
|
@@ -78,17 +91,17 @@ db_config = DatabaseConfig()
|
|
|
78
91
|
def initialize_database():
|
|
79
92
|
"""Initialize database connections and pool"""
|
|
80
93
|
global _connection_pool
|
|
81
|
-
|
|
94
|
+
|
|
82
95
|
try:
|
|
83
96
|
# Test connection first
|
|
84
97
|
if not db_config.test_connection():
|
|
85
98
|
raise Exception("Database connection test failed")
|
|
86
|
-
|
|
99
|
+
|
|
87
100
|
# Create connection pool
|
|
88
101
|
_connection_pool = db_config.create_connection_pool()
|
|
89
|
-
|
|
102
|
+
|
|
90
103
|
logger.info("Database initialization completed successfully")
|
|
91
|
-
|
|
104
|
+
|
|
92
105
|
except Exception as e:
|
|
93
106
|
logger.error(f"Database initialization failed: {str(e)}")
|
|
94
107
|
raise
|
|
@@ -98,16 +111,28 @@ def get_connection_pool() -> psycopg2.pool.ThreadedConnectionPool:
|
|
|
98
111
|
"""Get the database connection pool"""
|
|
99
112
|
global _connection_pool
|
|
100
113
|
if _connection_pool is None:
|
|
101
|
-
|
|
114
|
+
error_msg = (
|
|
115
|
+
"Database not initialized. This usually means:\n"
|
|
116
|
+
"1. Missing or incorrect .env file in app/ directory\n"
|
|
117
|
+
"2. Database credentials are wrong\n"
|
|
118
|
+
"3. Database container is not running\n"
|
|
119
|
+
"4. Database initialization failed during startup\n"
|
|
120
|
+
"Please check the startup logs for more details."
|
|
121
|
+
)
|
|
122
|
+
logger.error(error_msg)
|
|
123
|
+
raise Exception(error_msg)
|
|
102
124
|
return _connection_pool
|
|
103
125
|
|
|
104
126
|
|
|
105
|
-
def
|
|
106
|
-
"""
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
127
|
+
def _validate_connection(conn) -> bool:
|
|
128
|
+
"""Validate if a connection is still alive"""
|
|
129
|
+
try:
|
|
130
|
+
# Test if connection is alive with a simple query
|
|
131
|
+
with conn.cursor() as cursor:
|
|
132
|
+
cursor.execute("SELECT 1")
|
|
133
|
+
return True
|
|
134
|
+
except (psycopg2.OperationalError, psycopg2.InterfaceError):
|
|
135
|
+
return False
|
|
111
136
|
|
|
112
137
|
|
|
113
138
|
@contextmanager
|
|
@@ -117,52 +142,79 @@ def get_db_connection():
|
|
|
117
142
|
conn = None
|
|
118
143
|
try:
|
|
119
144
|
conn = pool.getconn()
|
|
145
|
+
|
|
146
|
+
# Validate connection before using it
|
|
147
|
+
if not _validate_connection(conn):
|
|
148
|
+
logger.warning("Stale connection detected, getting new connection")
|
|
149
|
+
pool.putconn(conn, close=True)
|
|
150
|
+
conn = pool.getconn()
|
|
151
|
+
|
|
120
152
|
logger.debug("Database connection acquired from pool")
|
|
121
153
|
yield conn
|
|
122
154
|
except Exception as e:
|
|
123
155
|
logger.error(f"Database connection error: {str(e)}")
|
|
124
156
|
if conn:
|
|
125
|
-
|
|
157
|
+
try:
|
|
158
|
+
# Only rollback if connection is still open
|
|
159
|
+
if not conn.closed:
|
|
160
|
+
conn.rollback()
|
|
161
|
+
except (psycopg2.OperationalError, psycopg2.InterfaceError) as rollback_error:
|
|
162
|
+
logger.warning(f"Could not rollback closed connection: {str(rollback_error)}")
|
|
126
163
|
raise
|
|
127
164
|
finally:
|
|
128
165
|
if conn:
|
|
129
|
-
|
|
130
|
-
|
|
166
|
+
try:
|
|
167
|
+
# If connection is broken, close it instead of returning to pool
|
|
168
|
+
if conn.closed:
|
|
169
|
+
pool.putconn(conn, close=True)
|
|
170
|
+
else:
|
|
171
|
+
pool.putconn(conn)
|
|
172
|
+
logger.debug("Database connection returned to pool")
|
|
173
|
+
except Exception as put_error:
|
|
174
|
+
logger.error(f"Error returning connection to pool: {str(put_error)}")
|
|
131
175
|
|
|
132
176
|
|
|
133
177
|
@contextmanager
|
|
134
178
|
def get_db_cursor():
|
|
135
179
|
"""Get a database cursor (context manager)"""
|
|
136
180
|
with get_db_connection() as conn:
|
|
137
|
-
cursor = conn.cursor()
|
|
181
|
+
cursor = conn.cursor(cursor_factory=RealDictCursor)
|
|
138
182
|
try:
|
|
139
183
|
yield cursor
|
|
140
|
-
conn.
|
|
184
|
+
if not conn.closed:
|
|
185
|
+
conn.commit()
|
|
141
186
|
except Exception as e:
|
|
142
|
-
conn.
|
|
187
|
+
if not conn.closed:
|
|
188
|
+
try:
|
|
189
|
+
conn.rollback()
|
|
190
|
+
except (psycopg2.OperationalError, psycopg2.InterfaceError) as rollback_error:
|
|
191
|
+
logger.warning(f"Could not rollback transaction on closed connection: {str(rollback_error)}")
|
|
143
192
|
logger.error(f"Database cursor error: {str(e)}")
|
|
144
193
|
raise
|
|
145
194
|
finally:
|
|
146
|
-
|
|
195
|
+
try:
|
|
196
|
+
cursor.close()
|
|
197
|
+
except Exception as close_error:
|
|
198
|
+
logger.warning(f"Error closing cursor: {str(close_error)}")
|
|
147
199
|
|
|
148
200
|
|
|
149
201
|
class DatabaseManager:
|
|
150
202
|
"""Database manager for common operations"""
|
|
151
|
-
|
|
203
|
+
|
|
152
204
|
@staticmethod
|
|
153
205
|
def execute_query(query: str, params: tuple = None) -> list:
|
|
154
206
|
"""Execute a SELECT query and return results"""
|
|
155
207
|
with get_db_cursor() as cursor:
|
|
156
208
|
cursor.execute(query, params)
|
|
157
209
|
return cursor.fetchall()
|
|
158
|
-
|
|
210
|
+
|
|
159
211
|
@staticmethod
|
|
160
212
|
def execute_update(query: str, params: tuple = None) -> int:
|
|
161
213
|
"""Execute an INSERT/UPDATE/DELETE query and return affected rows"""
|
|
162
214
|
with get_db_cursor() as cursor:
|
|
163
215
|
cursor.execute(query, params)
|
|
164
216
|
return cursor.rowcount
|
|
165
|
-
|
|
217
|
+
|
|
166
218
|
@staticmethod
|
|
167
219
|
def execute_scalar(query: str, params: tuple = None):
|
|
168
220
|
"""Execute a query and return a single value"""
|
|
@@ -178,7 +230,41 @@ class DatabaseManager:
|
|
|
178
230
|
# Handle tuple result
|
|
179
231
|
return result[0] if len(result) > 0 else None
|
|
180
232
|
return None
|
|
181
|
-
|
|
233
|
+
|
|
234
|
+
@staticmethod
|
|
235
|
+
@contextmanager
|
|
236
|
+
def transaction():
|
|
237
|
+
"""
|
|
238
|
+
Context manager for database transactions.
|
|
239
|
+
Wraps multiple operations in a single transaction.
|
|
240
|
+
|
|
241
|
+
Usage:
|
|
242
|
+
with DatabaseManager.transaction() as cursor:
|
|
243
|
+
cursor.execute("INSERT INTO table1 ...")
|
|
244
|
+
cursor.execute("INSERT INTO table2 ...")
|
|
245
|
+
# Auto-commits on success, auto-rollbacks on exception
|
|
246
|
+
"""
|
|
247
|
+
with get_db_connection() as conn:
|
|
248
|
+
cursor = conn.cursor(cursor_factory=RealDictCursor)
|
|
249
|
+
try:
|
|
250
|
+
yield cursor
|
|
251
|
+
if not conn.closed:
|
|
252
|
+
conn.commit()
|
|
253
|
+
logger.debug("Transaction committed successfully")
|
|
254
|
+
except Exception as e:
|
|
255
|
+
if not conn.closed:
|
|
256
|
+
try:
|
|
257
|
+
conn.rollback()
|
|
258
|
+
logger.warning(f"Transaction rolled back due to error: {str(e)}")
|
|
259
|
+
except (psycopg2.OperationalError, psycopg2.InterfaceError) as rollback_error:
|
|
260
|
+
logger.error(f"Could not rollback transaction: {str(rollback_error)}")
|
|
261
|
+
raise
|
|
262
|
+
finally:
|
|
263
|
+
try:
|
|
264
|
+
cursor.close()
|
|
265
|
+
except Exception as close_error:
|
|
266
|
+
logger.warning(f"Error closing transaction cursor: {str(close_error)}")
|
|
267
|
+
|
|
182
268
|
@staticmethod
|
|
183
269
|
def health_check() -> dict:
|
|
184
270
|
"""Perform database health check"""
|
|
@@ -186,7 +272,7 @@ class DatabaseManager:
|
|
|
186
272
|
with get_db_cursor() as cursor:
|
|
187
273
|
cursor.execute("SELECT version(), current_database(), current_user")
|
|
188
274
|
result = cursor.fetchone()
|
|
189
|
-
|
|
275
|
+
|
|
190
276
|
if result:
|
|
191
277
|
# Handle RealDictRow (dictionary-like) result
|
|
192
278
|
if hasattr(result, 'get'):
|
|
@@ -217,5 +303,13 @@ class DatabaseManager:
|
|
|
217
303
|
}
|
|
218
304
|
|
|
219
305
|
|
|
220
|
-
# Database initialization
|
|
221
|
-
|
|
306
|
+
# Database initialization on module import
|
|
307
|
+
try:
|
|
308
|
+
initialize_database()
|
|
309
|
+
except Exception as e:
|
|
310
|
+
logger.error(f"Failed to initialize database on startup: {str(e)}")
|
|
311
|
+
logger.error("⚠️ CRITICAL: Application started without database connection!")
|
|
312
|
+
logger.error("⚠️ Please check your .env file and database configuration")
|
|
313
|
+
logger.error("⚠️ The application will not function properly without a database connection")
|
|
314
|
+
# Don't raise here to allow the app to start even if DB is unavailable
|
|
315
|
+
# But log it clearly so it's obvious what's wrong
|