From 48b035528df9252ef24e1baa669653da03cca6c7 Mon Sep 17 00:00:00 2001
From: Neha Nene <nehanene@google.com>
Date: Fri, 3 Nov 2023 12:03:08 -0400
Subject: [PATCH] feat: support query parameter for MSSQL connection (#1026)

---
 data_validation/cli_tools.py            | 2 ++
 docs/connections.md                     | 3 +++
 third_party/ibis/ibis_mssql/__init__.py | 9 ++++++++-
 third_party/ibis/ibis_mssql/api.py      | 2 ++
 4 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/data_validation/cli_tools.py b/data_validation/cli_tools.py
index 8e20db1d4..05a901aff 100644
--- a/data_validation/cli_tools.py
+++ b/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)"],
diff --git a/docs/connections.md b/docs/connections.md
index 759928569..c3c436e0d 100644
--- a/docs/connections.md
+++ b/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
diff --git a/third_party/ibis/ibis_mssql/__init__.py b/third_party/ibis/ibis_mssql/__init__.py
index 089ec2725..4b9fedbb6 100644
--- a/third_party/ibis/ibis_mssql/__init__.py
+++ b/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,6 +43,7 @@ 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":
@@ -49,6 +51,11 @@ def do_connect(
                     "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,
@@ -56,7 +63,7 @@ def do_connect(
                 username=user,
                 password=password,
                 database=database,
-                query={"driver": odbc_driver},
+                query=query,
             )
         else:
             alchemy_url = sa.engine.url.make_url(url)
diff --git a/third_party/ibis/ibis_mssql/api.py b/third_party/ibis/ibis_mssql/api.py
index 1bd963eba..eb128605d 100644
--- a/third_party/ibis/ibis_mssql/api.py
+++ b/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