Coverage for pgsql_upserter / schema_inspector.py: 92.63%
75 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-26 18:15 -0300
1"""PostgreSQL table schema introspection utilities."""
3import logging
4import psycopg2
6from dataclasses import dataclass
7from psycopg2.extras import RealDictCursor
9from .exceptions import TableNotFoundError, SchemaIntrospectionError
11logger = logging.getLogger(__name__)
14@dataclass
15class ColumnInfo:
16 """Information about a table column."""
17 name: str
18 data_type: str
19 is_nullable: bool
20 default_value: str | None
21 is_auto_generated: bool
22 ordinal_position: int
25@dataclass
26class UniqueConstraint:
27 """Information about a unique constraint."""
28 name: str
29 columns: list[str]
30 is_primary: bool
33@dataclass
34class TableSchema:
35 """Complete table schema information."""
36 table_name: str
37 schema_name: str
38 columns: list[ColumnInfo]
39 unique_constraints: list[UniqueConstraint]
40 primary_key: UniqueConstraint | None
42 @property
43 def valid_columns(self) -> list[str]:
44 """Non-auto-generated column names for data insertion."""
45 return [col.name for col in self.columns if not col.is_auto_generated]
48def _is_auto_generated_column(column_default: str, data_type: str) -> bool:
49 """Determine if a column is auto-generated.
51 Checks for:
52 - SERIAL/BIGSERIAL types
53 - nextval() sequences
54 - GENERATED columns
55 - DEFAULT CURRENT_TIMESTAMP
56 """
57 if not column_default:
58 return False
60 column_default = column_default.lower()
61 data_type = data_type.lower()
63 # SERIAL types
64 if data_type in ('serial', 'bigserial'): 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true
65 return True
67 # Sequence defaults
68 if 'nextval(' in column_default:
69 return True
71 # Generated columns
72 if column_default.startswith('generated'):
73 return True
75 # Timestamp defaults
76 if any(ts in column_default for ts in ['current_timestamp', 'now()', 'clock_timestamp()']): 76 ↛ 79line 76 didn't jump to line 79 because the condition on line 76 was always true
77 return True
79 return False
82def inspect_table_schema(
83 connection,
84 table_name: str,
85 schema: str = 'public'
86) -> TableSchema:
87 """Inspect PostgreSQL table schema and return structured information.
89 Args:
90 connection: Active PostgreSQL connection
91 table_name: Name of the table to inspect
92 schema: Schema name (default: 'public')
94 Returns:
95 TableSchema: Complete table schema information
97 Raises:
98 TableNotFoundError: If table doesn't exist
99 SchemaIntrospectionError: If schema cannot be introspected
100 """
101 try:
102 with connection.cursor(cursor_factory=RealDictCursor) as cursor:
103 # First, verify table exists
104 cursor.execute("""
105 SELECT 1 FROM information_schema.tables
106 WHERE table_schema = %s AND table_name = %s
107 """, (schema, table_name))
109 if not cursor.fetchone():
110 raise TableNotFoundError(f"Table '{schema}.{table_name}' not found")
112 # Get column information
113 columns = _get_columns_info(cursor, table_name, schema)
115 # Get unique constraints
116 unique_constraints = _get_unique_constraints(cursor, table_name, schema)
118 # Find primary key
119 primary_key = next((uc for uc in unique_constraints if uc.is_primary), None)
121 logger.debug(f"Successfully introspected table '{schema}.{table_name}' "
122 f"with {len(columns)} columns and {len(unique_constraints)} constraints")
124 return TableSchema(
125 table_name=table_name,
126 schema_name=schema,
127 columns=columns,
128 unique_constraints=unique_constraints,
129 primary_key=primary_key
130 )
132 except psycopg2.Error as e:
133 raise SchemaIntrospectionError(f"Failed to introspect table '{schema}.{table_name}': {e}")
136def _get_columns_info(cursor, table_name: str, schema: str) -> list[ColumnInfo]:
137 """Get detailed column information from information_schema."""
138 cursor.execute("""
139 SELECT
140 c.column_name,
141 c.data_type,
142 c.is_nullable,
143 c.column_default,
144 c.ordinal_position,
145 c.is_generated,
146 c.generation_expression
147 FROM information_schema.columns c
148 WHERE c.table_schema = %s AND c.table_name = %s
149 ORDER BY c.ordinal_position
150 """, (schema, table_name))
152 columns = []
153 for row in cursor.fetchall():
154 # Check if column is auto-generated
155 is_auto_generated = (
156 # SERIAL types
157 row['data_type'] in ('serial', 'bigserial') or
158 # Sequence defaults
159 (row['column_default'] and 'nextval(' in row['column_default'].lower()) or
160 # Generated columns (PostgreSQL 12+)
161 row['is_generated'] == 'ALWAYS' or
162 # Timestamp defaults
163 (row['column_default'] and any(ts in row['column_default'].lower()
164 for ts in ['current_timestamp', 'now()', 'clock_timestamp()']))
165 )
167 columns.append(ColumnInfo(
168 name=row['column_name'],
169 data_type=row['data_type'],
170 is_nullable=row['is_nullable'] == 'YES',
171 default_value=row['column_default'],
172 is_auto_generated=is_auto_generated,
173 ordinal_position=row['ordinal_position']
174 ))
176 return columns
179def _get_unique_constraints(cursor, table_name: str, schema: str) -> list[UniqueConstraint]:
180 """Get unique constraints from information_schema."""
181 cursor.execute("""
182 SELECT
183 tc.constraint_name,
184 tc.constraint_type,
185 array_agg(kcu.column_name ORDER BY kcu.ordinal_position) as columns
186 FROM information_schema.table_constraints tc
187 JOIN information_schema.key_column_usage kcu
188 ON tc.constraint_name = kcu.constraint_name
189 AND tc.table_schema = kcu.table_schema
190 WHERE tc.table_schema = %s
191 AND tc.table_name = %s
192 AND tc.constraint_type IN ('PRIMARY KEY', 'UNIQUE')
193 GROUP BY tc.constraint_name, tc.constraint_type
194 ORDER BY tc.constraint_type DESC, tc.constraint_name
195 """, (schema, table_name))
197 constraints = []
198 for row in cursor.fetchall():
199 # Parse PostgreSQL array format: {col1,col2,col3} -> ['col1', 'col2', 'col3']
200 columns_raw = row['columns']
201 if isinstance(columns_raw, list):
202 # Already a list (some PostgreSQL drivers)
203 columns = columns_raw
204 elif isinstance(columns_raw, str): 204 ↛ 209line 204 didn't jump to line 209 because the condition on line 204 was always true
205 # Parse PostgreSQL array string format
206 columns = columns_raw.strip('{}').split(',')
207 else:
208 # Fallback
209 columns = [str(columns_raw)]
211 constraints.append(UniqueConstraint(
212 name=row['constraint_name'],
213 columns=columns,
214 is_primary=row['constraint_type'] == 'PRIMARY KEY'
215 ))
217 return constraints