Fix make_connection_configuration for MSSQL (adding a driver) (#113)

astronomer · Aug 17, 2023 · 09dbdac · 09dbdac
1 parent e63ae36
commit 09dbdac
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/great_expectations_provider/operators/great_expectations.py b/great_expectations_provider/operators/great_expectations.py
@@ -239,6 +239,7 @@ def __init__(
     def make_connection_configuration(self) -> Dict[str, str]:
         """Builds connection strings based off existing Airflow connections. Only supports necessary extras."""
         uri_string = ""
+        driver = ""
         if not self.conn:
             raise ValueError(f"Connections does not exist in Airflow for conn_id: {self.conn_id}")
         self.schema = self.schema or self.conn.schema
@@ -247,11 +248,21 @@ def make_connection_configuration(self) -> Dict[str, str]:
             odbc_connector = ""
             if conn_type in ("redshift", "postgres"):
                 odbc_connector = "postgresql+psycopg2"
+                database_name = self.schema
             elif conn_type == "mysql":
                 odbc_connector = "mysql"
-            else:
+                database_name = self.schema
+            elif conn_type == "mssql":
                 odbc_connector = "mssql+pyodbc"
-            uri_string = f"{odbc_connector}://{self.conn.login}:{self.conn.password}@{self.conn.host}:{self.conn.port}/{self.schema}"  # noqa
+                ms_driver = self.conn.extra_dejson.get("driver") or "ODBC Driver 17 for SQL Server"
+                driver = f"?driver={ms_driver}"
+                database_name = self.conn.schema or "master"
+            else:
+                raise ValueError(f"Conn type: {conn_type} is not supported.")
+            uri_string = (
+                f"{odbc_connector}://{self.conn.login}:{self.conn.password}@"
+                f"{self.conn.host}:{self.conn.port}/{database_name}{driver}"
+            )
         elif conn_type == "snowflake":
             try:
                 return self.build_snowflake_connection_config_from_hook()

diff --git a/tests/operators/test_great_expectations.py b/tests/operators/test_great_expectations.py
@@ -793,7 +793,7 @@ def test_great_expectations_operator__make_connection_string_mysql():
 
 
 def test_great_expectations_operator__make_connection_string_mssql():
-    test_conn_conf = {"connection_string": "mssql+pyodbc://user:password@connection:5439/schema"}
+    test_conn_conf = {"connection_string": "mssql+pyodbc://user:password@connection:5439/schema?driver=driver"}
     operator = GreatExpectationsOperator(
         task_id="task_id",
         data_context_config=in_memory_data_context_config,
@@ -810,6 +810,7 @@ def test_great_expectations_operator__make_connection_string_mssql():
         password="password",
         schema="schema",
         port=5439,
+        extra='{"driver": "driver"}',
     )
     operator.conn_type = operator.conn.conn_type
     assert operator.make_connection_configuration() == test_conn_conf