Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix make_connection_configuration for MSSQL (adding a driver) #113

Merged
merged 6 commits into from
Aug 17, 2023
Merged

Fix make_connection_configuration for MSSQL (adding a driver) #113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f4a996d
add driver to MSSQL connection
TJaniF Jul 5, 2023
af870f6
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 5, 2023
041dc78
fix lint
TJaniF Jul 5, 2023
1d0563a
adjust unit test
TJaniF Jul 5, 2023
a62a3c6
make mssql connection explicit add else statement
TJaniF Jul 6, 2023
a578e5b
adding default name for MS SQL server db because it does not accept t…
TJaniF Jul 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions great_expectations_provider/operators/great_expectations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nitpick: not yet tested myself, but I think the following should work 🤔

Suggested change
ms_driver = self.conn.extra_dejson.get("driver") or "ODBC Driver 17 for SQL Server"
ms_driver = self.conn.extra_dejson.get("driver", "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()
Expand Down
3 changes: 2 additions & 1 deletion tests/operators/test_great_expectations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down