Coverage for pgsql_upserter / config.py: 100.00%

1"""Database connection and configuration utilities.""" 

2 

3import logging 

4import os 

5import psycopg2 

6 

7from dotenv import load_dotenv 

8from psycopg2.extras import RealDictCursor 

9 

10from .exceptions import ConnectionError, PermissionError 

11 

12load_dotenv() 

13 

14logger = logging.getLogger(__name__) 

15 

16 

17def create_connection_from_env(): 

18 """Create PostgreSQL connection from environment variables. 

19 

20 Expected environment variables: 

21 - PGHOST (default: localhost) 

22 - PGPORT (default: 5432) 

23 - PGDATABASE (required) 

24 - PGUSER (required) 

25 - PGPASSWORD (required) 

26 

27 Returns: 

28 psycopg2.connection: Active database connection 

29 

30 Raises: 

31 ConnectionError: If connection fails 

32 """ 

33 try: 

34 connection_params = { 

35 'host': os.getenv('PGHOST', 'localhost'), 

36 'port': int(os.getenv('PGPORT', '5432')), 

37 'database': os.getenv('PGDATABASE'), 

38 'user': os.getenv('PGUSER'), 

39 'password': os.getenv('PGPASSWORD'), 

40 } 

41 

42 # Validate required parameters 

43 if not all([connection_params['database'], connection_params['user'], connection_params['password']]): 

44 raise ConnectionError( 

45 "Missing required environment variables: PGDATABASE, PGUSER, PGPASSWORD") 

46 

47 logger.debug( 

48 f"Connecting to PostgreSQL: {connection_params['user']}@{connection_params['host']}:{connection_params['port']}/{connection_params['database']}") # noqa 

49 

50 connection = psycopg2.connect(**connection_params) 

51 connection.autocommit = False # Explicit transaction control 

52 

53 return connection 

54 

55 except psycopg2.Error as e: 

56 raise ConnectionError(f"Failed to connect to PostgreSQL: {e}") 

57 except Exception as e: 

58 raise ConnectionError(f"Unexpected error during connection: {e}") 

59 

60 

61def validate_permissions(connection) -> None: 

62 """Validate that user has required permissions for upsert operations. 

63 

64 Checks: 

65 - CREATE privilege on database (for temp tables) 

66 - General connection health 

67 

68 Args: 

69 connection: Active PostgreSQL connection 

70 

71 Raises: 

72 PermissionError: If user lacks required permissions 

73 """ 

74 try: 

75 with connection.cursor(cursor_factory=RealDictCursor) as cursor: 

76 # Test CREATE TEMP TABLE permission 

77 cursor.execute(""" 

78 CREATE TEMP TABLE pgsql_upserter_permission_test ( 

79 test_col INTEGER 

80 ); 

81 DROP TABLE pgsql_upserter_permission_test; 

82 """) 

83 connection.commit() 

84 

85 logger.debug("Permission validation successful") 

86 

87 except psycopg2.Error as e: 

88 connection.rollback() 

89 raise PermissionError(f"Insufficient permissions for upsert operations: {e}") 

90 

91 

92def test_connection(): 

93 """Test database connection and validate permissions. 

94 

95 Returns: 

96 psycopg2.connection: Validated connection ready for use 

97 

98 Raises: 

99 ConnectionError: If connection fails 

100 PermissionError: If user lacks required permissions 

101 """ 

102 connection = create_connection_from_env() 

103 validate_permissions(connection) 

104 

105 logger.debug("Database connection and permissions validated successfully") 

106 return connection