Repository navigation
Expand file tree
/
Copy pathserver.py
More file actions
164 lines (137 loc) · 5.76 KB
/
Copy pathserver.py
File metadata and controls
164 lines (137 loc) · 5.76 KB
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
# pystackql/core/server.py
"""
Server connection management for PyStackQL.
This module handles connections to a StackQL server using the Postgres wire protocol.
"""
class ServerConnection:
"""Manages connections to a StackQL server.
This class handles connecting to and querying a StackQL server
using the Postgres wire protocol.
"""
def __init__(self, server_address='127.0.0.1', server_port=5466):
"""Initialize the ServerConnection.
Args:
server_address (str, optional): Address of the server. Defaults to '127.0.0.1'.
server_port (int, optional): Port of the server. Defaults to 5466.
"""
self.server_address = server_address
self.server_port = server_port
self._conn = None
# Import psycopg on demand to avoid dependency issues
try:
global psycopg, dict_row
import psycopg
from psycopg.rows import dict_row
except ImportError:
raise ImportError("psycopg is required in server mode but is not installed. "
"Please install psycopg and try again.")
# Connect to the server
self._connect()
def _connect(self):
"""Connect to the StackQL server.
Returns:
bool: True if connection successful, False otherwise
"""
try:
self._conn = psycopg.connect(
dbname='stackql',
user='stackql',
host=self.server_address,
port=self.server_port,
autocommit=True,
row_factory=dict_row
)
return True
except psycopg.OperationalError as oe:
print(f"OperationalError while connecting to the server: {oe}")
except Exception as e:
print(f"Unexpected error while connecting to the server: {e}")
return False
def is_connected(self):
"""Check if the connection to the server is active.
Returns:
bool: True if connected, False otherwise
"""
return self._conn is not None and not self._conn.closed
def ensure_connected(self):
"""Ensure the connection to the server is active.
If the connection is closed, attempt to reconnect.
Returns:
bool: True if connected, False otherwise
"""
if not self.is_connected():
return self._connect()
return True
def execute_query(self, query, is_statement=False):
"""Execute a query on the server.
Args:
query (str): The query to execute
is_statement (bool, optional): Whether this is a statement (non-SELECT). Defaults to False.
Returns:
list: Results of the query as a list of dictionaries
Raises:
ConnectionError: If no active connection is available
"""
if not self.ensure_connected():
raise ConnectionError("No active connection to the server")
try:
with self._conn.cursor() as cur:
cur.execute(query)
if is_statement:
# Return status message for non-SELECT statements
result_msg = cur.statusmessage
return [{'message': result_msg}]
try:
# Fetch results for SELECT queries
rows = cur.fetchall()
return rows
except psycopg.ProgrammingError as e:
# Handle cases with no results
if "no results to fetch" in str(e):
return []
else:
raise
except psycopg.OperationalError as oe:
print(f"OperationalError during query execution: {oe}")
# Try to reconnect and retry once
if self._connect():
return self.execute_query(query, is_statement)
except Exception as e:
print(f"Unexpected error during query execution: {e}")
return []
def execute_query_with_new_connection(self, query):
"""Execute a query with a new connection.
This method creates a new connection to the server, executes the query,
and then closes the connection.
Args:
query (str): The query to execute
Returns:
list: Results of the query as a list of dictionaries
"""
try:
with psycopg.connect(
dbname='stackql',
user='stackql',
host=self.server_address,
port=self.server_port,
row_factory=dict_row
) as conn:
with conn.cursor() as cur:
cur.execute(query)
try:
rows = cur.fetchall()
except psycopg.ProgrammingError as e:
if str(e) == "no results to fetch":
rows = []
else:
raise
return rows
except psycopg.OperationalError as oe:
print(f"OperationalError while connecting to the server: {oe}")
except Exception as e:
print(f"Unexpected error while connecting to the server: {e}")
return []
def close(self):
"""Close the connection to the server."""
if self._conn and not self._conn.closed:
self._conn.close()