From eb9f05afdb08aae58d33a53ba5e6c91b44728ffe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cau=C3=AA=20Garcia=20Polimanti?= <caue.polimanti@gmail.com>
Date: Wed, 20 Sep 2023 02:41:49 -0300
Subject: [PATCH] Adds a setting to define the correlation id placement
 injected in the SQL statement. (#68)

---
 HISTORY.rst           |  3 +++
 docs/installation.rst |  9 +++++++++
 src/cid/cursor.py     |  7 ++++++-
 tests/test_cursor.py  | 10 ++++++++++
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index 96b78f8..be664b5 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -21,6 +21,9 @@ those versions, you should pin ``django-cid < 3``.
 - |backward-incompatible| Drop support of Django 3.0 and prior
   versions. Only Django 3.1 and later are supported.
 
+- Add `CID_SQL_STATEMENT_TEMPLATE` setting to customize the position
+  of the correlation relative to the original SQL statement.
+  Contributed by Cauê Garcia Polimanti (@CaueP).
 
 2.3 (2022-06-13)
 ++++++++++++++++
diff --git a/docs/installation.rst b/docs/installation.rst
index fd31266..7811940 100644
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -244,6 +244,15 @@ a string with a ``cid`` format parameter:
     CID_SQL_COMMENT_TEMPLATE = 'correlation={cid}'
 
 
+Also, you may change the position of the correlation id injected in
+the statement by defining a ``CID_SQL_STATEMENT_TEMPLATE`` that is
+a string with a ``cid`` and a ``sql`` format parameter:
+
+.. code-block:: python
+
+    CID_SQL_STATEMENT_TEMPLATE = '/* {cid} */\n{sql}'
+
+
 Inclusion of the correlation id in templates
 --------------------------------------------
 
diff --git a/src/cid/cursor.py b/src/cid/cursor.py
index 2531fd8..f7282c1 100644
--- a/src/cid/cursor.py
+++ b/src/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/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