BUG#34467201: Add init_command connection option

oscpache · oscpache · commit 7cdff89a111b · 2022-10-26T08:35:15.000-05:00
With this patch, support for the "inti_command" option is added as
part of the connection settings. This new option allows the user to
specify a command to be executed right after the connection is
established, that's to say, as part of the connection initialization.

Additionally, adding this option inside the Django database
configuration when connection is established now works,
previously this was ignored.

Thanks Sander van de Graaf for the contribution.

Change-Id: I24069cb071cdf388ebb543500c22db6d2c141d40
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -19,6 +19,7 @@ v8.0.32
 - BUG#34675508: Character set 'utf8' unsupported in python mysql connector when using MariaDB
 - BUG#34556157: Kerberos authorization fails when using SSPI as security interface
 - BUG#34499578: MySQLCursor.executemany() fails to correctly identify BULK data loading ops
+- BUG#34467201: Add init_command connection option
 - BUG#32625155: Tests fail against group replication cluster
 - BUG#30089671: Fix decoding VARBINARY columns when using a prepared cursor
 - BUG#28020811: Fix multiple reference leaks in the C extension
diff --git a/lib/mysql/connector/abstracts.py b/lib/mysql/connector/abstracts.py
@@ -198,6 +198,7 @@ def __init__(self) -> None:
         self._compress: bool = False
 
         self._consume_results: bool = False
+        self._init_command: Optional[str] = None
 
     def __enter__(self) -> MySQLConnectionAbstract:
         return self
@@ -611,6 +612,11 @@ def config(self, **kwargs: Any) -> None:
         if "ssl_disabled" in config:
             self._ssl_disabled = config.pop("ssl_disabled")
 
+        # If an init_command is set, keep it, so we can execute it in _post_connection
+        if "init_command" in config:
+            self._init_command = config["init_command"]
+            del config["init_command"]
+
         # Other configuration
         set_ssl_flag = False
         for key, value in config.items():
@@ -1149,6 +1155,8 @@ def _post_connection(self) -> None:
             self.time_zone = self._time_zone
         if self._sql_mode:
             self.sql_mode = self._sql_mode
+        if self._init_command:
+            self._execute_query(self._init_command)
 
     @abstractmethod
     def disconnect(self) -> Any:
diff --git a/lib/mysql/connector/constants.py b/lib/mysql/connector/constants.py
@@ -92,6 +92,7 @@
     "oci_config_file": None,
     "fido_callback": None,
     "kerberos_auth_mode": None,
+    "init_command": None,
 }
 
 CNX_POOL_ARGS: Tuple[str, str, str] = ("pool_name", "pool_size", "pool_reset_session")
diff --git a/lib/mysql/connector/django/base.py b/lib/mysql/connector/django/base.py
@@ -360,6 +360,8 @@ def get_connection_params(self) -> Dict[str, Any]:
             kwargs["host"] = settings_dict["HOST"]
         if settings_dict["PORT"]:
             kwargs["port"] = int(settings_dict["PORT"])
+        if settings_dict.get("OPTIONS", {}).get("init_command"):
+            kwargs["init_command"] = settings_dict["OPTIONS"]["init_command"]
 
         # Raise exceptions for database warnings if DEBUG is on
         kwargs["raise_on_warnings"] = settings.DEBUG
diff --git a/lib/mysql/connector/django/client.py b/lib/mysql/connector/django/client.py
@@ -62,8 +62,9 @@ def settings_to_cmd_args_env(
         if defaults_file:
             args.append(f"--defaults-file={defaults_file}")
 
-        # We force SQL_MODE to TRADITIONAL
-        args.append("--init-command=SET @@session.SQL_MODE=TRADITIONAL")
+        # Load any custom init_commands. We always force SQL_MODE to TRADITIONAL
+        init_command = settings_dict["OPTIONS"].get("init_command", "")
+        args.append(f"--init-command=SET @@session.SQL_MODE=TRADITIONAL;{init_command}")
 
         if user:
             args.append(f"--user={user}")
diff --git a/tests/test_bugs.py b/tests/test_bugs.py
@@ -7300,3 +7300,33 @@ def test_batch_insert_happens_in_executemany(self):
                 ),
                 "Batch insert failed!",
             )
+
+
+class BugOra34467201(tests.MySQLConnectorTests):
+    """BUG#34467201: Add init_command connection option
+
+    With this patch, support for the inti_command option is added as
+    part of the connection settings. This new option allows the user to
+    specify a command to be executed right after the connection is
+    established, that's to say, as part of the connection initialization.
+    """
+
+    var_name = "test_var"
+    var_value = "BugOra34467201"
+
+    def setUp(self):
+        config = tests.get_mysql_config()
+        config["init_command"] = f"SET @{self.var_name}='{self.var_value}'"
+        self.cnx = mysql.connector.connect(**config)
+
+    def tearDown(self):
+        with self.cnx.cursor() as cur:
+            cur.execute(f"SET @{self.var_name} = NULL")
+        self.cnx.close()
+
+    @foreach_cnx
+    def test_init_command(self):
+        with self.cnx.cursor() as cur:
+            cur.execute(f"SELECT @{self.var_name}")
+            res = cur.fetchall()
+            self.assertEqual((self.var_value,), res[0])
diff --git a/tests/test_django.py b/tests/test_django.py
@@ -575,3 +575,36 @@ def test_safe_string(self):
         )
         self.cur.execute("SELECT col1, col2 FROM {0}".format(self.tbl), ())
         self.assertEqual(self.cur.fetchall(), [(safe_text, safe_bytes.encode())])
+
+
+class BugOra34467201(tests.MySQLConnectorTests):
+    """BUG#34467201: Add init_command connection option
+
+    With this patch, support for the inti_command option is added as
+    part of the connection settings. This new option allows the user to
+    specify a command to be executed right after the connection is
+    established, that's to say, as part of the connection initialization.
+    """
+
+    var_name = "test_var"
+    var_value = "BugOra34467201"
+
+    def setUp(self):
+        settings.DATABASES["default"]["OPTIONS"] = {
+            "init_command": f"SET @{self.var_name}='{self.var_value}'"
+        }
+        dbconfig = tests.get_mysql_config()
+        self.conn = mysql.connector.connect(**dbconfig)
+        self.cnx = DatabaseWrapper(settings.DATABASES["default"])
+
+    def tearDown(self):
+        with self.cnx.cursor() as cur:
+            cur.execute(f"SET @{self.var_name} = NULL")
+        self.cnx.close()
+        del settings.DATABASES["default"]["OPTIONS"]
+
+    def test_init_command(self):
+        with self.cnx.cursor() as cur:
+            cur.execute(f"SELECT @{self.var_name}")
+            res = cur.fetchall()
+            self.assertEqual((self.var_value,), res[0])

Original file line number	Diff line number	Diff line change
`@@ -92,6 +92,7 @@`
`92`	`92`	`"oci_config_file": None,`
`93`	`93`	`"fido_callback": None,`
`94`	`94`	`"kerberos_auth_mode": None,`
	`95`	`+ "init_command": None,`
`95`	`96`	`}`
`96`	`97`
`97`	`98`	`CNX_POOL_ARGS: Tuple[str, str, str] = ("pool_name", "pool_size", "pool_reset_session")`