File size: 4,465 Bytes
a2ec7b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
PostgreSQL Login Helper for MCPMark
====================================
Handles PostgreSQL authentication and connection validation.
"""
import json
import psycopg2
from pathlib import Path
from typing import Optional, Dict, Any
from src.base.login_helper import BaseLoginHelper
from src.logger import get_logger
logger = get_logger(__name__)
class PostgresLoginHelper(BaseLoginHelper):
"""Handles PostgreSQL authentication and connection validation."""
def __init__(
self,
host: str = "localhost",
port: int = 5432,
database: str = "postgres",
username: str = "postgres",
password: str = None,
state_path: Optional[Path] = None,
):
"""Initialize PostgreSQL login helper.
Args:
host: Database host
port: Database port
database: Database name
username: Database username
password: Database password
state_path: Path to save connection state
"""
super().__init__()
self.host = host
self.port = port
self.database = database
self.username = username
self.password = password
self.state_path = state_path or Path.home() / ".mcpbench" / "postgres_auth.json"
# Ensure state directory exists
self.state_path.parent.mkdir(parents=True, exist_ok=True)
def login(self, **kwargs) -> bool:
"""Test PostgreSQL connection and save state.
Returns:
bool: True if connection successful
"""
try:
# Test connection
conn = psycopg2.connect(
host=self.host,
port=self.port,
database=self.database,
user=self.username,
password=self.password,
connect_timeout=10,
)
# Execute test query
with conn.cursor() as cur:
cur.execute("SELECT version()")
version = cur.fetchone()[0]
logger.info(f"PostgreSQL connection successful: {version}")
# Check permissions
cur.execute(
"""
SELECT has_database_privilege(%s, 'CREATE')
""",
(self.database,),
)
can_create = cur.fetchone()[0]
if not can_create:
logger.warning("User does not have CREATE privilege on database")
conn.close()
# Save connection state
self._save_connection_state(
{
"host": self.host,
"port": self.port,
"database": self.database,
"username": self.username,
"version": version,
"can_create": can_create,
"authenticated_at": self._get_current_timestamp(),
}
)
return True
except psycopg2.Error as e:
logger.error(f"PostgreSQL connection failed: {e}")
return False
except Exception as e:
logger.error(f"Unexpected error during PostgreSQL login: {e}")
return False
def _save_connection_state(self, state: Dict[str, Any]):
"""Save connection state to file."""
try:
# Don't save password
safe_state = {k: v for k, v in state.items() if k != "password"}
with open(self.state_path, "w") as f:
json.dump(safe_state, f, indent=2)
# Set restrictive permissions
self.state_path.chmod(0o600)
logger.info(f"Connection state saved to: {self.state_path}")
except Exception as e:
logger.error(f"Failed to save connection state: {e}")
def _get_current_timestamp(self) -> str:
"""Get current timestamp in ISO format."""
from datetime import datetime, timezone
return datetime.now(timezone.utc).isoformat()
def is_connected(self) -> bool:
"""Check if we can connect to PostgreSQL."""
return self.login()
def get_connection_params(self) -> Dict[str, Any]:
"""Get connection parameters (without password)."""
return {
"host": self.host,
"port": self.port,
"database": self.database,
"user": self.username,
}
|