|
8 | 8 | _get_version, |
9 | 9 | _format_auth |
10 | 10 | ) |
11 | | -import sys, subprocess, json, os, asyncio, functools, psycopg2 |
| 11 | +import sys, subprocess, json, os, asyncio, functools |
12 | 12 | from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor |
13 | | -from psycopg2.extras import RealDictCursor |
14 | 13 | import pandas as pd |
15 | 14 |
|
16 | 15 | from io import StringIO |
@@ -257,12 +256,20 @@ def __init__(self, |
257 | 256 | except OSError as e: |
258 | 257 | raise ValueError(f"Unable to create the log directory {log_dir}: {str(e)}") |
259 | 258 |
|
260 | | - if self.server_mode: |
261 | | - # server mode, connect to a server via the postgres wire protocol |
262 | | - self.server_address = server_address |
| 259 | + if self.server_mode: |
| 260 | + # server mode, connect to a server via the postgres wire protocol |
| 261 | + # Attempt to import psycopg2 only if server_mode is True |
| 262 | + global psycopg2, RealDictCursor |
| 263 | + try: |
| 264 | + import psycopg2 |
| 265 | + from psycopg2.extras import RealDictCursor |
| 266 | + except ImportError: |
| 267 | + raise ImportError("psycopg2 is required in server mode but is not installed. Please install psycopg2 and try again.") |
| 268 | + |
| 269 | + self.server_address = server_address |
263 | 270 | self.server_port = server_port |
264 | | - # establish the connection |
265 | | - self._conn = self._connect_to_server() |
| 271 | + # establish the connection |
| 272 | + self._conn = self._connect_to_server() |
266 | 273 | else: |
267 | 274 | # local mode, executes the binary locally |
268 | 275 | self.params = [] |
|
0 commit comments