diff --git a/cid/__init__.py b/cid/__init__.py index 54b59c7..f87966b 100644 --- a/cid/__init__.py +++ b/cid/__init__.py @@ -1,4 +1,8 @@ -import pkg_resources +import sys +if sys.version_info >= (3, 8): + from importlib.metadata import version +else: + from importlib_metadata import version -__version__ = pkg_resources.get_distribution('django-cid').version +__version__ = version('django-cid') diff --git a/cid/cursor.py b/cid/cursor.py index 2531fd8..f7282c1 100644 --- a/cid/cursor.py +++ b/cid/cursor.py @@ -4,6 +4,7 @@ DEFAULT_CID_SQL_COMMENT_TEMPLATE = 'cid: {cid}' +DEFAULT_SQL_STATEMENT_TEMPLATE = '/* {cid} */\n{sql}' class CidCursorWrapper: @@ -31,12 +32,16 @@ def add_comment(self, sql): cid_sql_template = getattr( settings, 'CID_SQL_COMMENT_TEMPLATE', DEFAULT_CID_SQL_COMMENT_TEMPLATE ) + sql_statement_template = getattr( + settings, 'CID_SQL_STATEMENT_TEMPLATE', DEFAULT_SQL_STATEMENT_TEMPLATE + ) cid = get_cid() if not cid: return sql cid = cid.replace('/*', r'\/\*').replace('*/', r'\*\/') cid = cid_sql_template.format(cid=cid) - return f"/* {cid} */\n{sql}" + statement = sql_statement_template.format(cid=cid, sql=sql) + return statement # The following methods cannot be implemented in __getattr__, because the # code must run when the method is invoked, not just when it is accessed. diff --git a/setup.py b/setup.py index 33722ea..abed07e 100755 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ setup( name='django-cid', - version='2.4.dev0', + version='2.5.dev0', description="""Correlation IDs in Django for debugging requests""", long_description=readme, author='Snowball One', diff --git a/tests/test_cursor.py b/tests/test_cursor.py index 5d187e2..f1a912f 100644 --- a/tests/test_cursor.py +++ b/tests/test_cursor.py @@ -33,6 +33,16 @@ def test_adds_comment_setting_overriden(self, get_cid): self.cursor_wrapper.add_comment("SELECT 1;") ) + @override_settings(CID_SQL_STATEMENT_TEMPLATE='{sql}\n/* {cid} */') + @mock.patch('cid.cursor.get_cid') + def test_adds_comment_with_statement_template_setting_overriden(self, get_cid): + get_cid.return_value = 'testing-cursor-after-sql-statement' + expected = "SELECT 1;\n/* cid: testing-cursor-after-sql-statement */" + self.assertEqual( + expected, + self.cursor_wrapper.add_comment("SELECT 1;") + ) + @mock.patch('cid.cursor.get_cid') def test_no_comment_when_cid_is_none(self, get_cid): get_cid.return_value = None