已开启
fix bug #200
zzh创建于 2025年7月11日
fix bug #200
已开启
共 3 个文件变更+11-1
| @@ -79,6 +79,12 @@ class Driver: | |||
| 79 | # Specify default schema is public. | 79 | # Specify default schema is public. |
| 80 | conn = psycopg2_connect(url) | 80 | conn = psycopg2_connect(url) |
| 81 | conn.cursor().execute('select 1;') | 81 | conn.cursor().execute('select 1;') |
| 82 | + cur = conn.cursor() | ||
| 83 | + cur.execute('select rolmonitoradmin from pg_catalog.pg_roles where rolname = pg_catalog.CURRENT_USER();') | ||
| 84 | + result = cur.fetchone() | ||
| 85 | + if not result[0]: | ||
| 86 | + raise PermissionError('Monitor admin role is required to access DBMind.') | ||
| 87 | + cur.close() | ||
| 82 | conn.close() | 88 | conn.close() |
| 83 | self._url = url | 89 | self._url = url |
| 84 | self.parsed_dsn = psycopg2.extensions.parse_dsn(url) | 90 | self.parsed_dsn = psycopg2.extensions.parse_dsn(url) |
| @@ -88,6 +94,8 @@ class Driver: | |||
| 88 | dsn=url, **_psycopg2_kwargs | 94 | dsn=url, **_psycopg2_kwargs |
| 89 | ) | 95 | ) |
| 90 | self.initialized = True | 96 | self.initialized = True |
| 97 | + except PermissionError as e: | ||
| 98 | + raise PermissionError(e) | ||
| 91 | except Exception as e: | 99 | except Exception as e: |
| 92 | raise ConnectionError(e) | 100 | raise ConnectionError(e) |
| 93 | 101 | ||
| @@ -80,7 +80,7 @@ class RPCServer: | |||
| 80 | ).json() | 80 | ).json() |
| 81 | 81 | ||
| 82 | # Validate credential. | 82 | # Validate credential. |
| 83 | - if not self.checker(req.username, req.pwd): | 83 | + if not self.checker(req.username, req.pwd)[0]: |
| 84 | return RPCResponse(req, success=False, | 84 | return RPCResponse(req, success=False, |
| 85 | exception='Failed to validate authorization.').json() | 85 | exception='Failed to validate authorization.').json() |
| 86 | 86 | ||
| @@ -63,5 +63,7 @@ def create_agent_rpc_service(): | |||
| 63 | return True | 63 | return True |
| 64 | except ConnectionError: | 64 | except ConnectionError: |
| 65 | return False | 65 | return False |
| 66 | + except PermissionError as e: | ||
| 67 | + return False, str(e) | ||
| 66 | 68 | ||
| 67 | return RPCServer(_rpc_register, checker) | 69 | return RPCServer(_rpc_register, checker) |