feat: support query parameter for MSSQL connection (#1026)

data_validation/cli_tools.py

@@ -81,6 +81,8 @@
         ["user", "User used to connect"],
         ["password", "Password for supplied user"],
         ["database", "Database to connect to (default master)"],
+        ["query", "Connection query parameters"],
+        ["url", "SQL Server SQLAlchemy connection URL"],
     ],
     "MySQL": [
         ["host", "Desired MySQL host (default localhost)"],

docs/connections.md

@@ -151,6 +151,7 @@ data-validation connections add
 
 ## MSSQL Server
 MSSQL Server connections require [pyodbc](https://pypi.org/project/pyodbc/) as the driver: `pip install pyodbc`.
+For connection query parameter options, see https://docs.sqlalchemy.org/en/20/dialects/mssql.html#hostname-connections.
 
 ```
 data-validation connections add 
@@ -162,6 +163,8 @@ data-validation connections add
     --user USER                                         MSSQL user
     --password PASSWORD                                 MSSQL password
     --database DATABASE                                 MSSQL database
+    [--url URL]                                         SQLAlchemy connection URL
+    [--query QUERY]                                     Connection query parameters i.e. '{"TrustServerCertificate": "yes"}'
 ```
 
 ## Postgres

third_party/ibis/ibis_mssql/__init__.py

@@ -21,6 +21,7 @@
 from ibis.backends.mssql.datatypes import _type_from_result_set_info
 
 import third_party.ibis.ibis_mssql.datatypes
+import json
 
 
 # The MSSQL backend uses the Ibis MSSQL compiler, but overrides
@@ -42,21 +43,27 @@ def do_connect(
         url: str = None,
         driver: Literal["pyodbc"] = "pyodbc",
         odbc_driver: str = "ODBC Driver 17 for SQL Server",
+        query: str = None,
     ) -> None:
         if url is None:
             if driver != "pyodbc":
                 raise NotImplementedError(
                     "pyodbc is currently the only supported driver"
                 )
 
+            if query:
+                query = json.loads(query)
+            else:
+                query = {"driver": odbc_driver}
+
             alchemy_url = sa.engine.url.URL.create(
                 f"mssql+{driver}",
                 host=host,
                 port=port,
                 username=user,
                 password=password,
                 database=database,
-                query={"driver": odbc_driver},
+                query=query,
             )
         else:
             alchemy_url = sa.engine.url.make_url(url)

third_party/ibis/ibis_mssql/api.py

@@ -25,6 +25,7 @@ def mssql_connect(
     url: str = None,
     driver: Literal["pyodbc"] = "pyodbc",
     odbc_driver: str = "ODBC Driver 17 for SQL Server",
+    query: str = None,
 ):
     backend = MsSqlBackend()
     backend.do_connect(
@@ -36,5 +37,6 @@ def mssql_connect(
         url=url,
         driver=driver,
         odbc_driver=odbc_driver,
+        query=query,
     )
     return backend