Merge pull request #1 from singlestore-labs/users/asaini/allow-connection-override

Allow overriding `connection` while executing the sql magic.

Author: kesmit13

src/sql/command.py

@@ -22,7 +22,7 @@ class SQLCommand:
 
     """
 
-    def __init__(self, magic, user_ns, line, cell) -> None:
+    def __init__(self, magic, user_ns, line, cell, connection=None) -> None:
         self._line = line
         self._cell = cell
 
@@ -89,6 +89,9 @@ def __init__(self, magic, user_ns, line, cell) -> None:
         if add_alias:
             self.parsed["connection"] = self.args.line[0]
 
+        if connection is not None:
+            self.parsed["connection"] = connection
+
         if self.args.with_:
             self.args.with_ = [
                 Template(item).render(user_ns) for item in self.args.with_

src/sql/magic.py

@@ -335,7 +335,7 @@ def check_random_arguments(self, line="", cell=""):
         action="append",
         help="Interactive mode",
     )
-    def execute(self, line="", cell="", local_ns=None):
+    def execute(self, line="", cell="", local_ns=None, connection=None):
         """
         Runs SQL statement against a database, specified by
         SQLAlchemy connect string.
@@ -363,17 +363,26 @@ def execute(self, line="", cell="", local_ns=None):
 
         """
         return self._execute(
-            line=line, cell=cell, local_ns=local_ns, is_interactive_mode=False
+            line=line,
+            cell=cell,
+            local_ns=local_ns,
+            is_interactive_mode=False,
+            connection=connection,
         )
 
     @modify_exceptions
-    def _execute(self, line, cell, local_ns, is_interactive_mode=False):
+    def _execute(
+        self, line, cell, local_ns, is_interactive_mode=False, connection=None
+    ):
         """
         This function implements the cell logic; we create this private
         method so we can control how the function is called. Otherwise,
         decorating ``SqlMagic.execute`` will break when adding the ``@log_call``
         decorator with ``payload=True``
 
+        ``connection`` is any [DBAPI v2](https://peps.python.org/pep-0249/) compatible connection object.
+        If provided, it will override any other connection configuration.
+
         NOTE: telemetry has been removed, we can remove this function
         """
 
@@ -401,7 +410,7 @@ def interactive_execute_wrapper(**kwargs):
         user_ns = self.shell.user_ns.copy()
         user_ns.update(local_ns)
 
-        command = SQLCommand(self, user_ns, line, cell)
+        command = SQLCommand(self, user_ns, line, cell, connection)
         # args.line: contains the line after the magic with all options removed
 
         args = command.args