Hiccups connecting to Azure SQL Server instances using Azure Active Directory

[inigohidalgo]

Hello,

I recently developed this small Kedro extension to load Ibis tables with: repo. It was a very clean solution for querying sql tables from kedro, but as I pointed out to @cpcloud a couple of months ago, I was having issues connecting to MSSQL servers running on Azure. I hadn't looked too far into so I didn't want to raise an issue.

Now I am trying to start including it in some parts of our production code but am having an issue correctly connecting to the server. I am not familiar with SQL backends at all, but the way I currently connect to SQL from pandas is as follows (slightly adapted from here):

# requirements
# pypi:
#  ibis-project[mssql]==7.0.0
#  pandas
#  sqlalchemy
#  pyodbc
#  azure-identity


import struct
from azure.identity import AzureCliCredential
import sqlalchemy as sa
import pandas as pd
import ibis


credential = AzureCliCredential()
databaseToken = credential.get_token("https://database.windows.net/")

# get bytes from token obtained
tokenb = bytes(databaseToken[0], "UTF-8")
exptoken = b""
for i in tokenb:
    exptoken += bytes({i})
    exptoken += bytes(1)
tokenstruct = struct.pack("=i", len(exptoken)) + exptoken
SQL_COPT_SS_ACCESS_TOKEN = 1256
connect_args = {
    "attrs_before": {
            SQL_COPT_SS_ACCESS_TOKEN: tokenstruct
    }
}

connection_string_right = f"Driver={{ODBC Driver 17 for SQL Server}};SERVER={server};DATABASE={database}"

dummy_query = "SELECT * FROM CLIENT.ADDRESS"

pd.read_sql(
    dummy_query,
    con=sa.create_engine(f"mssql://?odbc_connect={connection_string_right}", connect_args=connect_args)
)

ibis.connect(
    f"mssql://?odbc_connect={connection_string_right}", connect_args=connect_args
)
"""
backends/base/__init__.py:781, in BaseBackend.connect(self, *args, **kwargs)
    758 \"""Connect to the database.
    759 
    760 Parameters
   (...)
    778     An instance of the backend
    779 \"""
    780 new_backend = self.__class__(*args, **kwargs)
--> 781 new_backend.reconnect()
    782 return new_backend
...
    794 def reconnect(self) -> None:
    795     \"""Reconnect to the database already configured with connect.\"""
--> 796     self.do_connect(*self._con_args, **self._con_kwargs)
TypeError: Backend.do_connect() got an unexpected keyword argument 'odbc_connect'
"""


ibis.connect(
    f"mssql://{connection_string_right}", connect_args=connect_args
)
"""
backends/base/__init__.py:781, in BaseBackend.connect(self, *args, **kwargs)
    758 \"""Connect to the database.
    759 
    760 Parameters
   (...)
    778     An instance of the backend
    779 \"""
    780 new_backend = self.__class__(*args, **kwargs)
--> 781 new_backend.reconnect()
    782 return new_backend
...
    794 def reconnect(self) -> None:
    795     \"""Reconnect to the database already configured with connect.\"""
--> 796     self.do_connect(*self._con_args, **self._con_kwargs)
TypeError: Backend.do_connect() got an unexpected keyword argument 'connect_args'
"""

ibis.connect(
    f"mssql://{connection_string_right}", **connect_args
)
"""
backends/base/__init__.py:781, in BaseBackend.connect(self, *args, **kwargs)
    758 \"""Connect to the database.
    759 
    760 Parameters
   (...)
    778     An instance of the backend
    779 \"""
    780 new_backend = self.__class__(*args, **kwargs)
--> 781 new_backend.reconnect()
    782 return new_backend
...
    794 def reconnect(self) -> None:
    795     \"""Reconnect to the database already configured with connect.\"""
--> 796     self.do_connect(*self._con_args, **self._con_kwargs)
TypeError: Backend.do_connect() got an unexpected keyword argument 'attrs_before'
"""

Am I missing some place I can provide those "attrs_before"?

Thank you!

[inigohidalgo]

I tried to edit

ibis/ibis/backends/mssql/__init__.py

Line 51 in cb1dc65

engine = sa.create_engine(alchemy_url, poolclass=sa.pool.StaticPool)

to make it

engine = sa.create_engine(alchemy_url, poolclass=sa.pool.StaticPool, connect_args=kwargs)

in line with this change in mysql backend

This allowed me to instantiate the connection, but it seems to have just kicked the can down the road as trying to list_tables raises this error:

sqlalchemy/engine/default.py", line 615, in connect
    return self.loaded_dbapi.connect(*cargs, **cparams)
  File "src/pymssql/_pymssql.pyx", line 593, in pymssql._pymssql.connect
TypeError: connect() got an unexpected keyword argument 'connect_args'

This is where my understanding more or less runs out. I get the feeling it might be due to the fact that Ibis uses pymssql instead of the default pyodbc, as I have just found that my connection method is listed under SQLAlchemy's docs for MSSQL.

[inigohidalgo]

Related: #6640 (comment)

[cpcloud]

@inigohidalgo Thanks for opening a discussion!

I think we've been putting off PyODBC MSSQL support, but given the status of pymssql, and the fact that PyODBC seems to be the only way to connect Azure services we should likely move to that in the next release. It's likely some functionality will break since I think pymssql has slightly better support for some of the mssql specific types, but we can see how it goes.

Let's move to supporting connection with pyodbc in the next release to unblock @inigohidalgo's work!

[inigohidalgo]

Thanks for the quick answer @cpcloud.

I'd like to help get this into Ibis so let me know where I can help. I don't think I'll be very useful to you for the actual implementation but let me know if I can help you test any early builds against our infrastructure.

[cpcloud]

@inigohidalgo I've put up #7317 to move our MSSQL backend over to ODBC. Can you give it a try? Let's move the discussion to that PR as well!