Skip to content

Commit 6369b16

Browse files
Thiyagu55Thiyagu55
authored
Add compatability for ; in sqlalchemy,psycopg2 & django (#123)
This PR fixes a zero day issue where comments are added after the end of SQL statement i.e. after `;`. This will create no or wrong association of the comment wrt the SQL statement. We previously fixed it for psycopg2 only, but this CL reactors the code and fixes the issue for sqlalchemy and as a precaution was added to django also.
1 parent 3a89de1 commit 6369b16

6 files changed

Lines changed: 71 additions & 30 deletions

File tree

‎python/sqlcommenter-python/google/cloud/sqlcommenter/__init__.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,12 +36,23 @@ def generate_sql_comment(**meta):
3636

3737
# Sort the keywords to ensure that caching works and that testing is
3838
# deterministic. It eases visual inspection as well.
39+
3940
return ' /*' + KEY_VALUE_DELIMITER.join(
4041
'{}={!r}'.format(url_quote(key), url_quote(value)) for key, value in sorted(meta.items())
4142
if value is not None
4243
) + '*/'
4344

4445

46+
def add_sql_comment(sql, **meta):
47+
comment = generate_sql_comment(**meta)
48+
sql = sql.rstrip()
49+
if sql[-1] == ';':
50+
sql = sql[:-1] + comment + ';'
51+
else:
52+
sql = sql + comment
53+
return sql
54+
55+
4556
def url_quote(s):
4657
if not isinstance(s, (str, bytes)):
4758
return s

‎python/sqlcommenter-python/google/cloud/sqlcommenter/django/middleware.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import django
2020
from django.db import connection
2121
from django.db.backends.utils import CursorDebugWrapper
22-
from google.cloud.sqlcommenter import generate_sql_comment
22+
from google.cloud.sqlcommenter import add_sql_comment
2323
from google.cloud.sqlcommenter.opencensus import get_opencensus_values
2424
from google.cloud.sqlcommenter.opentelemetry import get_opentelemetry_values
2525

@@ -62,7 +62,8 @@ def __call__(self, execute, sql, params, many, context):
6262
db_driver = context['connection'].settings_dict.get('ENGINE', '')
6363
resolver_match = self.request.resolver_match
6464

65-
sql_comment = generate_sql_comment(
65+
sql = add_sql_comment(
66+
sql,
6667
# Information about the controller.
6768
controller=resolver_match.view_name if resolver_match and with_controller else None,
6869
# route is the pattern that matched a request with a controller i.e. the regex
@@ -85,7 +86,6 @@ def __call__(self, execute, sql, params, many, context):
8586
# See:
8687
# * https://github.com/basecamp/marginalia/issues/61
8788
# * https://github.com/basecamp/marginalia/pull/80
88-
sql += sql_comment
8989

9090
# Add the query to the query log if debugging.
9191
if context['cursor'].__class__ is CursorDebugWrapper:

‎python/sqlcommenter-python/google/cloud/sqlcommenter/psycopg2/extension.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import psycopg2
2020
import psycopg2.extensions
21-
from google.cloud.sqlcommenter import generate_sql_comment
21+
from google.cloud.sqlcommenter import add_sql_comment
2222
from google.cloud.sqlcommenter.flask import get_flask_info
2323
from google.cloud.sqlcommenter.opencensus import get_opencensus_values
2424
from google.cloud.sqlcommenter.opentelemetry import get_opentelemetry_values
@@ -80,12 +80,7 @@ def execute(self, sql, args=None):
8080
if with_opentelemetry:
8181
data.update(get_opentelemetry_values())
8282

83-
if len(data) != 0:
84-
sql = sql.rstrip()
85-
if sql[-1] == ';':
86-
sql = sql[:-1] + generate_sql_comment(**data) + ';'
87-
else:
88-
sql = sql + generate_sql_comment(**data)
83+
sql = add_sql_comment(sql, **data)
8984

9085
return psycopg2.extensions.cursor.execute(self, sql, args)
9186

‎python/sqlcommenter-python/google/cloud/sqlcommenter/sqlalchemy/executor.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
import logging
2222

2323
import sqlalchemy
24-
from google.cloud.sqlcommenter import generate_sql_comment
24+
from google.cloud.sqlcommenter import add_sql_comment
2525
from google.cloud.sqlcommenter.fastapi import get_fastapi_info
2626
from google.cloud.sqlcommenter.flask import get_flask_info
2727
from google.cloud.sqlcommenter.opencensus import get_opencensus_values
@@ -76,15 +76,14 @@ def before_cursor_execute(conn, cursor, sql, parameters, context, executemany):
7676
if with_opentelemetry:
7777
data.update(get_opentelemetry_values())
7878

79-
sql_comment = generate_sql_comment(**data)
79+
sql = add_sql_comment(sql, **data)
8080

8181
# TODO: Check if the database type is MySQL and figure out
8282
# if we should prepend comments because MySQL server truncates
8383
# logs greater than 1kB.
8484
# See:
8585
# * https://github.com/basecamp/marginalia/issues/61
8686
# * https://github.com/basecamp/marginalia/pull/80
87-
sql += sql_comment
8887

8988
return sql, parameters
9089

‎python/sqlcommenter-python/tests/generic/test_add_sql_comment.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
#!/usr/bin/python
2+
#
3+
# Copyright 2022 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
from unittest import TestCase
18+
19+
from google.cloud.sqlcommenter import add_sql_comment
20+
21+
22+
class AddSqlCommentTests(TestCase):
23+
def test_no_args(self):
24+
self.assertEqual(add_sql_comment('Select 1'), 'Select 1')
25+
26+
def test_end_comment_escaping(self):
27+
self.assertEqual("Select 1 /*a='%%2A/abc'*/", add_sql_comment("Select 1", a='*/abc'))
28+
29+
def test_bytes(self):
30+
self.assertIn("Select 1 /*a='%%2A/abc'*/", add_sql_comment("Select 1",a=b'*/abc'))
31+
32+
def test_url_quoting(self):
33+
self.assertIn("Select 1 /*foo='bar%%2Cbaz'*/", add_sql_comment("Select 1",foo='bar,baz'))
34+
35+
def test_sql_with_semicolon(self):
36+
self.assertIn("Select 1 /*foo='bar%%2Cbaz'*/;", add_sql_comment("Select 1;",foo='bar,baz'))

‎python/sqlcommenter-python/tests/sqlalchemy/tests.py‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -52,30 +52,30 @@ def test_no_args(self):
5252

5353
def test_db_driver(self):
5454
self.assertSQL(
55-
"SELECT 1; /*db_driver='driver'*/",
55+
"SELECT 1 /*db_driver='driver'*/;",
5656
with_db_driver=True,
5757
)
5858

5959
def test_db_framework(self):
6060
self.assertSQL(
61-
"SELECT 1; /*db_framework='sqlalchemy%%3A{}'*/".format(sqlalchemy.__version__),
61+
"SELECT 1 /*db_framework='sqlalchemy%%3A{}'*/;".format(sqlalchemy.__version__),
6262
with_db_framework=True,
6363
)
6464

6565
def test_opencensus(self):
6666
with mock_opencensus_tracer():
6767
self.assertSQL(
68-
"SELECT 1; /*traceparent='00-trace%%20id-span%%20id-00',"
69-
"tracestate='congo%%3Dt61rcWkgMzE%%2Crojo%%3D00f067aa0ba902b7'*/",
68+
"SELECT 1 /*traceparent='00-trace%%20id-span%%20id-00',"
69+
"tracestate='congo%%3Dt61rcWkgMzE%%2Crojo%%3D00f067aa0ba902b7'*/;",
7070
with_opencensus=True,
7171
)
7272

7373
@skipIfPy2
7474
def test_opentelemetry(self):
7575
with mock_opentelemetry_context():
7676
self.assertSQL(
77-
"SELECT 1; /*traceparent='00-000000000000000000000000deadbeef-000000000000beef-00',"
78-
"tracestate='some_key%%3Dsome_value'*/",
77+
"SELECT 1 /*traceparent='00-000000000000000000000000deadbeef-000000000000beef-00',"
78+
"tracestate='some_key%%3Dsome_value'*/;",
7979
with_opentelemetry=True,
8080
)
8181

@@ -85,8 +85,8 @@ def test_both_opentelemetry_and_opencensus_warn(self):
8585
"google.cloud.sqlcommenter.sqlalchemy.executor.logger"
8686
) as logger_mock, mock_opencensus_tracer(), mock_opentelemetry_context():
8787
self.assertSQL(
88-
"SELECT 1; /*traceparent='00-000000000000000000000000deadbeef-000000000000beef-00',"
89-
"tracestate='some_key%%3Dsome_value'*/",
88+
"SELECT 1 /*traceparent='00-000000000000000000000000deadbeef-000000000000beef-00',"
89+
"tracestate='some_key%%3Dsome_value'*/;",
9090
with_opentelemetry=True,
9191
with_opencensus=True,
9292
)
@@ -103,27 +103,27 @@ class FlaskTests(SQLAlchemyTestCase):
103103
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_flask_info', return_value=flask_info)
104104
def test_all_data(self, get_info):
105105
self.assertSQL(
106-
"SELECT 1; /*controller='c',framework='flask',route='/'*/",
106+
"SELECT 1 /*controller='c',framework='flask',route='/'*/;",
107107
)
108108

109109
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_flask_info', return_value=flask_info)
110110
def test_framework_disabled(self, get_info):
111111
self.assertSQL(
112-
"SELECT 1; /*controller='c',route='/'*/",
112+
"SELECT 1 /*controller='c',route='/'*/;",
113113
with_framework=False,
114114
)
115115

116116
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_flask_info', return_value=flask_info)
117117
def test_controller_disabled(self, get_info):
118118
self.assertSQL(
119-
"SELECT 1; /*framework='flask',route='/'*/",
119+
"SELECT 1 /*framework='flask',route='/'*/;",
120120
with_controller=False,
121121
)
122122

123123
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_flask_info', return_value=flask_info)
124124
def test_route_disabled(self, get_info):
125125
self.assertSQL(
126-
"SELECT 1; /*controller='c',framework='flask'*/",
126+
"SELECT 1 /*controller='c',framework='flask'*/;",
127127
with_route=False,
128128
)
129129

@@ -138,26 +138,26 @@ class FastAPITests(SQLAlchemyTestCase):
138138
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
139139
def test_all_data(self, get_info):
140140
self.assertSQL(
141-
"SELECT 1; /*controller='c',framework='fastapi',route='/'*/",
141+
"SELECT 1 /*controller='c',framework='fastapi',route='/'*/;",
142142
)
143143

144144
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
145145
def test_framework_disabled(self, get_info):
146146
self.assertSQL(
147-
"SELECT 1; /*controller='c',route='/'*/",
147+
"SELECT 1 /*controller='c',route='/'*/;",
148148
with_framework=False,
149149
)
150150

151151
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
152152
def test_controller_disabled(self, get_info):
153153
self.assertSQL(
154-
"SELECT 1; /*framework='fastapi',route='/'*/",
154+
"SELECT 1 /*framework='fastapi',route='/'*/;",
155155
with_controller=False,
156156
)
157157

158158
@mock.patch('google.cloud.sqlcommenter.sqlalchemy.executor.get_fastapi_info', return_value=fastapi_info)
159159
def test_route_disabled(self, get_info):
160160
self.assertSQL(
161-
"SELECT 1; /*controller='c',framework='fastapi'*/",
161+
"SELECT 1 /*controller='c',framework='fastapi'*/;",
162162
with_route=False,
163-
)
163+
)

0 commit comments

Comments
 (0)