Skip to content

Commit 7a5f6e9

Browse files
feat(lint): add ruff lint to CI and pre-commit hook (#17)
* feat(lint): add ruff lint step to CI and pre-commit hook Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * chore(lint): apply ruff formatting and fix bare except clauses Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent fb7c0f4 commit 7a5f6e9

12 files changed

Lines changed: 216 additions & 147 deletions

File tree

.github/workflows/tests.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,20 @@ on:
77
- master
88

99
jobs:
10+
lint:
11+
runs-on: ubuntu-24.04
12+
13+
steps:
14+
- uses: actions/checkout@v4
15+
16+
- name: Set up Python
17+
uses: actions/setup-python@v5
18+
with:
19+
python-version: "3.13"
20+
21+
- name: Lint
22+
run: make lint CHECK=1
23+
1024
tests:
1125
runs-on: ubuntu-24.04
1226

.pre-commit-config.yaml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
repos:
2+
- repo: https://github.com/astral-sh/ruff-pre-commit
3+
rev: v0.15.12
4+
hooks:
5+
- id: ruff
6+
args: [--fix]
7+
- id: ruff-format

Makefile

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,14 @@ install-test: install ## Install test dependencies in local virtualenv
3131
install-lint: install ## Install lint dependencies in local virtualenv
3232
($(VENV_RUN); $(PIP_CMD) install -r $(LINT_REQS))
3333

34-
lint: install-lint ## Format code with ruff
35-
$(VENV_DIR)/bin/ruff format postgresql_proxy tests plugins
34+
install-pre-commit: install-lint ## Install and register the pre-commit hook
35+
$(VENV_DIR)/bin/pre-commit install
36+
37+
CHECK ?=
38+
39+
lint: install-lint ## Format code with ruff (use CHECK=1 to check without modifying)
40+
$(VENV_DIR)/bin/ruff format $(if $(CHECK),--check,) postgresql_proxy tests plugins
41+
$(VENV_DIR)/bin/ruff check $(if $(CHECK),,--fix) postgresql_proxy tests plugins
3642

3743
start-postgres: ## Start local PostgreSQL test container and wait until ready
3844
@set -euo pipefail; \

plugins/tableau_hll/__init__.py

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,19 @@
22
import re
33

44
# The field to replace
5-
field_pattern = re.compile('(?<=[^\w])count\(distinct (?:cast\()?("[^"]+")\.("[^"]+")(?: as text\))?\)', re.IGNORECASE)
5+
field_pattern = re.compile(
6+
'(?<=[^\w])count\(distinct (?:cast\()?("[^"]+")\.("[^"]+")(?: as text\))?\)',
7+
re.IGNORECASE,
8+
)
69
# Table name
7-
table_pattern = re.compile('from ([^\(\)]+?)\s*\)? (?:AS )?("[^"]+")', re.IGNORECASE | re.DOTALL)
10+
table_pattern = re.compile(
11+
'from ([^\(\)]+?)\s*\)? (?:AS )?("[^"]+")', re.IGNORECASE | re.DOTALL
12+
)
13+
814

915
def rewrite_query(query, context):
10-
original_table = ''
11-
table_alias = ''
16+
original_table = ""
17+
table_alias = ""
1218

1319
# cache only works on current query. Mainly because there's no way to tell if the table has been modified between
1420
# 2 different requests.
@@ -26,8 +32,8 @@ def replace(match):
2632
hll_column_candidate = match.group(2).strip()
2733

2834
# need to know which columns are hll
29-
if not hll_table.lower() in column_cache:
30-
db_conn_info = context['instance_config'].redirect
35+
if hll_table.lower() not in column_cache:
36+
db_conn_info = context["instance_config"].redirect
3137
conn = None
3238
try:
3339
conn = psycopg2.connect(
@@ -36,17 +42,17 @@ def replace(match):
3642
db_conn_info.host,
3743
db_conn_info.port,
3844
# Get auth information from the proxied request
39-
context['connect_params']['database'],
40-
context['connect_params']['user']
45+
context["connect_params"]["database"],
46+
context["connect_params"]["user"],
4147
)
4248
)
43-
49+
4450
hll_type_code = None
4551
cur = conn.cursor()
4652
try:
4753
cur.execute("SELECT oid FROM pg_type WHERE typname='hll';")
48-
hll_type_code, = cur.fetchone()
49-
except:
54+
(hll_type_code,) = cur.fetchone()
55+
except Exception:
5056
pass
5157
finally:
5258
cur.close()
@@ -65,7 +71,7 @@ def replace(match):
6571
hll_columns.add(desc.name.lower())
6672

6773
column_cache[hll_table.lower()] = hll_columns
68-
except:
74+
except Exception:
6975
pass
7076
finally:
7177
cur.close()
@@ -76,13 +82,17 @@ def replace(match):
7682
conn.close()
7783

7884
# Replace
79-
if hll_column_candidate.strip('"').lower() in column_cache[hll_table.lower()]:
80-
return ' hll_cardinality(hll_union_agg({}.{})) :: BIGINT '.format(match.group(1), match.group(2))
85+
if (
86+
hll_column_candidate.strip('"').lower()
87+
in column_cache[hll_table.lower()]
88+
):
89+
return " hll_cardinality(hll_union_agg({}.{})) :: BIGINT ".format(
90+
match.group(1), match.group(2)
91+
)
8192

8293
# Don't replace
8394
return match.group(0)
8495

85-
8696
# Matches this string. The 2 groups are `schema.table` and `"alias"`
8797
# FROM schema.table) "alias"
8898
table_result = table_pattern.search(query)

postgresql_proxy/config_schema.py

Lines changed: 47 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,18 @@
1+
"""This class is used to validate the config"""
2+
13
import logging
24

3-
''' This class is used to validate the config
4-
'''
5+
56
class Schema:
67
def _validate(self):
78
pass
89

910
def __hyphen_to_underscore(self, k):
10-
return k.replace('-', '_')
11+
return k.replace("-", "_")
1112

1213
def _populate(self, data, definition):
1314
try:
14-
for (k, v) in data.items():
15+
for k, v in data.items():
1516
k = self.__hyphen_to_underscore(k)
1617
if k in definition:
1718
vtype = definition[k]
@@ -30,59 +31,59 @@ def _populate(self, data, definition):
3031

3132
def _assert_non_empty(self, *attrs):
3233
for attr in attrs:
33-
assert len(getattr(self, attr)) > 0, "{}.{} must not be empty".format(type(self).__name__, attr)
34+
assert len(getattr(self, attr)) > 0, "{}.{} must not be empty".format(
35+
type(self).__name__, attr
36+
)
3437

3538
def _assert_non_null(self, *attrs):
3639
for attr in attrs:
37-
assert getattr(self, attr) is not None, "{}.{} must not be None".format(type(self).__name__, attr)
40+
assert getattr(self, attr) is not None, "{}.{} must not be None".format(
41+
type(self).__name__, attr
42+
)
3843

3944

4045
class InterceptQuerySettings(Schema):
4146
def __init__(self, data):
4247
self.plugin = None
4348
self.function = None
4449

45-
self._populate(data, {
46-
'plugin': str,
47-
'function': str
48-
})
50+
self._populate(data, {"plugin": str, "function": str})
4951

5052
def _validate(self):
51-
self._assert_non_null('plugin', 'function')
52-
self._assert_non_empty('plugin', 'function')
53+
self._assert_non_null("plugin", "function")
54+
self._assert_non_empty("plugin", "function")
5355

5456

5557
class InterceptCommandSettings(Schema):
5658
def __init__(self, data):
5759
self.queries = []
5860
self.connects = None
5961

60-
self._populate(data, {
61-
'queries': [InterceptQuerySettings],
62-
'connects': str
63-
})
62+
self._populate(data, {"queries": [InterceptQuerySettings], "connects": str})
6463

6564

6665
class InterceptResponseSettings(Schema):
6766
def __init__(self, data):
6867
self.parameter_responses = []
6968
self.connects = None
7069

71-
self._populate(data, {
72-
'parameter_status': [InterceptQuerySettings],
73-
'connects': str
74-
})
70+
self._populate(
71+
data, {"parameter_status": [InterceptQuerySettings], "connects": str}
72+
)
7573

7674

7775
class InterceptSettings(Schema):
7876
def __init__(self, data):
7977
self.commands = None
8078
self.responses = None
8179

82-
self._populate(data, {
83-
'commands': InterceptCommandSettings,
84-
'responses': InterceptResponseSettings,
85-
})
80+
self._populate(
81+
data,
82+
{
83+
"commands": InterceptCommandSettings,
84+
"responses": InterceptResponseSettings,
85+
},
86+
)
8687

8788

8889
class Connection(Schema):
@@ -91,15 +92,11 @@ def __init__(self, data):
9192
self.host = None
9293
self.port = None
9394

94-
self._populate(data, {
95-
'name': str,
96-
'host': str,
97-
'port': int
98-
})
95+
self._populate(data, {"name": str, "host": str, "port": int})
9996

10097
def _validate(self):
101-
self._assert_non_null('name', 'host', 'port')
102-
self._assert_non_empty('name')
98+
self._assert_non_null("name", "host", "port")
99+
self._assert_non_empty("name")
103100

104101

105102
class InstanceSettings(Schema):
@@ -108,15 +105,17 @@ def __init__(self, data):
108105
self.redirect = None
109106
self.intercept = None
110107

111-
self._populate(data, {
112-
'listen': Connection,
113-
'redirect': Connection,
114-
'intercept': InterceptSettings
115-
})
116-
108+
self._populate(
109+
data,
110+
{
111+
"listen": Connection,
112+
"redirect": Connection,
113+
"intercept": InterceptSettings,
114+
},
115+
)
117116

118117
def _validate(self):
119-
self._assert_non_null('listen', 'redirect')
118+
self._assert_non_null("listen", "redirect")
120119

121120

122121
class Settings(Schema):
@@ -125,15 +124,13 @@ def __init__(self, data):
125124
self.intercept_log = None
126125
self.general_log = None
127126

128-
self._populate(data, {
129-
'log_level': str,
130-
'intercept_log': str,
131-
'general_log': str
132-
})
127+
self._populate(
128+
data, {"log_level": str, "intercept_log": str, "general_log": str}
129+
)
133130

134131
def _validate(self):
135-
self._assert_non_null('log_level', 'intercept_log', 'general_log')
136-
self._assert_non_empty('log_level', 'intercept_log', 'general_log')
132+
self._assert_non_null("log_level", "intercept_log", "general_log")
133+
self._assert_non_empty("log_level", "intercept_log", "general_log")
137134

138135

139136
class Config(Schema):
@@ -142,11 +139,10 @@ def __init__(self, data):
142139
self.settings = None
143140
self.instances = []
144141

145-
self._populate(data, {
146-
'plugins' : [str],
147-
'settings' : Settings,
148-
'instances' : [InstanceSettings]
149-
})
142+
self._populate(
143+
data,
144+
{"plugins": [str], "settings": Settings, "instances": [InstanceSettings]},
145+
)
150146

151147
def _validate(self):
152-
self._assert_non_empty('instances')
148+
self._assert_non_empty("instances")

postgresql_proxy/connection.py

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,28 +13,28 @@ def __init__(self, sock, address, name, events, context):
1313
self.context = context
1414
self.interceptor = None
1515
self.redirect_conn: Optional[Connection] = None
16-
self.out_bytes = b''
17-
self.in_bytes = b''
16+
self.out_bytes = b""
17+
self.in_bytes = b""
1818
self.terminated = False
1919

2020
def parse_length(self, length_bytes):
21-
return int.from_bytes(length_bytes, 'big')
21+
return int.from_bytes(length_bytes, "big")
2222

2323
def encode_length(self, length):
24-
return length.to_bytes(4, byteorder='big')
24+
return length.to_bytes(4, byteorder="big")
2525

2626
def received(self, in_bytes):
2727
self.in_bytes += in_bytes
2828
# Read packet from byte array while there are enough bytes to make up a packet.
2929
# Otherwise wait for more bytes to be received (break and exit)
3030
while True:
3131
ptype = self.in_bytes[0:1]
32-
if ptype == b'\x00':
32+
if ptype == b"\x00":
3333
if len(self.in_bytes) < 4:
3434
break
3535
header_length = 4
3636
body_length = self.parse_length(self.in_bytes[0:4]) - 4
37-
elif ptype == b'N':
37+
elif ptype == b"N":
3838
header_length = 1
3939
body_length = 0
4040
else:
@@ -52,18 +52,22 @@ def received(self, in_bytes):
5252
self.in_bytes = self.in_bytes[length:]
5353

5454
def process_inbound_packet(self, header, body):
55-
if header != b'N':
55+
if header != b"N":
5656
packet_type = header[0:-4]
57-
_logger.info("intercepting packet of type '%s' from %s", packet_type, self.name)
57+
_logger.info(
58+
"intercepting packet of type '%s' from %s", packet_type, self.name
59+
)
5860
body = self.interceptor.intercept(packet_type, body)
5961
header = packet_type + self.encode_length(len(body) + 4)
60-
if packet_type == b'X':
62+
if packet_type == b"X":
6163
# this a termination packet, it will indicate that the proxied client wants to close the
6264
# postgres connection properly
6365
self.terminated = True
6466

6567
message = header + body
66-
_logger.debug("Received message. Relaying. Speaker: %s, message:\n%s", self.name, message)
68+
_logger.debug(
69+
"Received message. Relaying. Speaker: %s, message:\n%s", self.name, message
70+
)
6771

6872
if self.redirect_conn:
6973
# redirect_conn might not be set (anymore) at this stage

postgresql_proxy/constants.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
ALLOWED_CONNECTION_PARAMETERS = [
32
"host",
43
"hostaddr",

0 commit comments

Comments
 (0)