Skip to content

Commit 5df6e02

Browse files
Prajwal-MicrosoftPrajwal-Microsoft
committed
fix: Code changes to Connect to SQL with Managed Identity using pyodbc
1 parent fd5dda9 commit 5df6e02

File tree

7 files changed

+191
-22
lines changed

7 files changed

+191
-22
lines changed

‎AzureFunctions/km-charts-function/Dockerfile‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22
# FROM mcr.microsoft.com/azure-functions/python:4-python3.11-appservice
33
FROM mcr.microsoft.com/azure-functions/python:4-python3.11
44

5+
# Install Microsoft ODBC Driver
6+
RUN apt-get update && \
7+
ACCEPT_EULA=Y apt-get install -y msodbcsql17
8+
59
ENV AzureWebJobsScriptRoot=/home/site/wwwroot \
610
AzureFunctionsJobHost__Logging__Console__IsEnabled=true
711

‎AzureFunctions/km-charts-function/function_app.py‎

Lines changed: 58 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,60 @@
11
from datetime import datetime
22
import azure.functions as func
3+
from azure.identity import DefaultAzureCredential
34
import logging
45
import json
56
import os
6-
import pymssql
7+
import pyodbc
78
import pandas as pd
9+
import struct
810
# from dotenv import load_dotenv
911
# load_dotenv()
1012

1113
app = func.FunctionApp(http_auth_level=func.AuthLevel.ANONYMOUS)
1214

15+
# get database connection
16+
def get_db_connection():
17+
driver = "{ODBC Driver 17 for SQL Server}"
18+
server = os.environ.get("SQLDB_SERVER")
19+
database = os.environ.get("SQLDB_DATABASE")
20+
username = os.environ.get("SQLDB_USERNAME")
21+
password = os.environ.get("SQLDB_PASSWORD")
22+
23+
# Attempt connection using Username & Password
24+
try:
25+
conn = pyodbc.connect(
26+
f"DRIVER={driver};SERVER={server};DATABASE={database};UID={username};PWD={password}",
27+
timeout=5
28+
)
29+
logging.info("Connected using Username & Password")
30+
return conn
31+
except pyodbc.Error as e:
32+
print(f"Failed with Username & Password: {str(e)}")
33+
34+
# If first attempt fails, try Azure Default Credential
35+
try:
36+
credential = DefaultAzureCredential()
37+
38+
token_bytes = credential.get_token(
39+
"https://database.windows.net/.default"
40+
).token.encode("utf-16-LE")
41+
token_struct = struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)
42+
SQL_COPT_SS_ACCESS_TOKEN = (
43+
1256 # This connection option is defined by microsoft in msodbcsql.h
44+
)
45+
46+
# Set up the connection
47+
connection_string = f"DRIVER={driver};SERVER={server};DATABASE={database};"
48+
conn = pyodbc.connect(
49+
connection_string, attrs_before={SQL_COPT_SS_ACCESS_TOKEN: token_struct}
50+
)
51+
52+
logging.info("Connected using Default Azure Credential")
53+
return conn
54+
except Exception as e:
55+
logging.error(f"Failed with Default Credential: {str(e)}")
56+
return None # Return None if both attempts fail
57+
1358
# add post methods - filters will come in the body (request.body), if body is not empty, update the where clause in the query
1459
@app.route(route="get_metrics", methods=["GET","POST"], auth_level=func.AuthLevel.ANONYMOUS)
1560
def get_metrics(req: func.HttpRequest) -> func.HttpResponse:
@@ -19,12 +64,7 @@ def get_metrics(req: func.HttpRequest) -> func.HttpResponse:
1964
if not data_type:
2065
data_type = 'filters'
2166

22-
server = os.environ.get("SQLDB_SERVER")
23-
database = os.environ.get("SQLDB_DATABASE")
24-
username = os.environ.get("SQLDB_USERNAME")
25-
password = os.environ.get("SQLDB_PASSWORD")
26-
27-
conn = pymssql.connect(server, username, password, database)
67+
conn = get_db_connection()
2868
cursor = conn.cursor()
2969

3070
# Adjust the dates to the current date
@@ -73,8 +113,12 @@ def get_metrics(req: func.HttpRequest) -> func.HttpResponse:
73113
) t'''
74114

75115
cursor.execute(sql_stmt)
76-
rows = cursor.fetchall()
116+
#rows = cursor.fetchall()
117+
118+
# Convert pyodbc.Row objects to tuples
119+
rows = [tuple(row) for row in cursor.fetchall()]
77120

121+
# Define column names
78122
column_names = [i[0] for i in cursor.description]
79123
df = pd.DataFrame(rows, columns=column_names)
80124
df.rename(columns={'key1':'key'}, inplace=True)
@@ -167,7 +211,8 @@ def get_metrics(req: func.HttpRequest) -> func.HttpResponse:
167211
#charts pt1
168212
cursor.execute(sql_stmt)
169213

170-
rows = cursor.fetchall()
214+
# rows = cursor.fetchall()
215+
rows = [tuple(row) for row in cursor.fetchall()]
171216

172217
column_names = [i[0] for i in cursor.description]
173218
df = pd.DataFrame(rows, columns=column_names)
@@ -208,7 +253,8 @@ def get_metrics(req: func.HttpRequest) -> func.HttpResponse:
208253

209254
cursor.execute(sql_stmt)
210255

211-
rows = cursor.fetchall()
256+
# rows = cursor.fetchall()
257+
rows = [tuple(row) for row in cursor.fetchall()]
212258

213259
column_names = [i[0] for i in cursor.description]
214260
df = pd.DataFrame(rows, columns=column_names)
@@ -286,7 +332,8 @@ def get_metrics(req: func.HttpRequest) -> func.HttpResponse:
286332

287333
cursor.execute(sql_stmt)
288334

289-
rows = cursor.fetchall()
335+
# rows = cursor.fetchall()
336+
rows = [tuple(row) for row in cursor.fetchall()]
290337

291338
column_names = [i[0] for i in cursor.description]
292339
df = pd.DataFrame(rows, columns=column_names)

‎AzureFunctions/km-charts-function/requirements.txt‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,5 +3,6 @@
33
# Manually managing azure-functions-worker may cause unexpected issues
44

55
azure-functions
6-
pymssql==2.3.0
7-
pandas
6+
pandas
7+
pyodbc==5.2.0
8+
azure.identity

‎AzureFunctions/km-rag-function/Dockerfile‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22
# FROM mcr.microsoft.com/azure-functions/python:4-python3.11-appservice
33
FROM mcr.microsoft.com/azure-functions/python:4-python3.11
44

5+
# Install Microsoft ODBC Driver
6+
RUN apt-get update && \
7+
ACCEPT_EULA=Y apt-get install -y msodbcsql17
8+
59
ENV AzureWebJobsScriptRoot=/home/site/wwwroot \
610
AzureFunctionsJobHost__Logging__Console__IsEnabled=true
711

‎AzureFunctions/km-rag-function/function_app.py‎

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import azure.functions as func
2+
from azure.identity import DefaultAzureCredential
3+
import logging
24
import openai
35
from azurefunctions.extensions.http.fastapi import Request, StreamingResponse
46
import asyncio
@@ -17,7 +19,8 @@
1719
from semantic_kernel.functions.kernel_arguments import KernelArguments
1820
from semantic_kernel.functions.kernel_function_decorator import kernel_function
1921
from semantic_kernel.kernel import Kernel
20-
import pymssql
22+
import pyodbc
23+
import struct
2124

2225
# from semantic_kernel import Kernel
2326
# from semantic_kernel.connectors.ai.open_ai import AzureChatCompletion
@@ -37,6 +40,49 @@
3740
# search_endpoint = os.environ.get("AZURE_AI_SEARCH_ENDPOINT")
3841
# search_key = os.environ.get("AZURE_AI_SEARCH_API_KEY")
3942

43+
# get database connection
44+
def get_db_connection():
45+
driver = "{ODBC Driver 17 for SQL Server}"
46+
server = os.environ.get("SQLDB_SERVER")
47+
database = os.environ.get("SQLDB_DATABASE")
48+
username = os.environ.get("SQLDB_USERNAME")
49+
password = os.environ.get("SQLDB_PASSWORD")
50+
51+
# Attempt connection using Username & Password
52+
try:
53+
conn = pyodbc.connect(
54+
f"DRIVER={driver};SERVER={server};DATABASE={database};UID={username};PWD={password}",
55+
timeout=5
56+
)
57+
logging.info("Connected using Username & Password")
58+
return conn
59+
except pyodbc.Error as e:
60+
print(f"Failed with Username & Password: {str(e)}")
61+
62+
# If first attempt fails, try Azure Default Credential
63+
try:
64+
credential = DefaultAzureCredential()
65+
66+
token_bytes = credential.get_token(
67+
"https://database.windows.net/.default"
68+
).token.encode("utf-16-LE")
69+
token_struct = struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)
70+
SQL_COPT_SS_ACCESS_TOKEN = (
71+
1256 # This connection option is defined by microsoft in msodbcsql.h
72+
)
73+
74+
# Set up the connection
75+
connection_string = f"DRIVER={driver};SERVER={server};DATABASE={database};"
76+
conn = pyodbc.connect(
77+
connection_string, attrs_before={SQL_COPT_SS_ACCESS_TOKEN: token_struct}
78+
)
79+
80+
logging.info("Connected using Default Azure Credential")
81+
return conn
82+
except Exception as e:
83+
logging.error(f"Failed with Default Credential: {str(e)}")
84+
return None # Return None if both attempts fail
85+
4086
class ChatWithDataPlugin:
4187
@kernel_function(name="Greeting", description="Respond to any greeting or general questions")
4288
def greeting(self, input: Annotated[str, "the question"]) -> Annotated[str, "The output is a string"]:
@@ -111,14 +157,8 @@ def get_SQL_Response(
111157
sql_query = completion.choices[0].message.content
112158
sql_query = sql_query.replace("```sql",'').replace("```",'')
113159
#print(sql_query)
114-
115-
# connectionString = os.environ.get("SQLDB_CONNECTION_STRING")
116-
server = os.environ.get("SQLDB_SERVER")
117-
database = os.environ.get("SQLDB_DATABASE")
118-
username = os.environ.get("SQLDB_USERNAME")
119-
password = os.environ.get("SQLDB_PASSWORD")
120160

121-
conn = pymssql.connect(server, username, password, database)
161+
conn = get_db_connection()
122162
# conn = pyodbc.connect(connectionString)
123163
cursor = conn.cursor()
124164
cursor.execute(sql_query)

‎AzureFunctions/km-rag-function/requirements.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ azure-functions
66
azurefunctions-extensions-http-fastapi==1.0.0b1
77
openai==1.57.0
88
semantic_kernel==1.0.4
9-
pymssql==2.3.0
109
azure-search-documents==11.6.0b3
10+
pyodbc==5.2.0
11+
azure.identity
1112
# httpx==0.27.2

‎Deployment/scripts/add_sql_users.ps1‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
param(
2+
[string] $resourceGroupName,
3+
[string] $serverName,
4+
[string] $databaseName,
5+
[string] $chartJsFuncAppName, # Function App Name of Chart JS
6+
[string] $ragFuncAppName # Function App Name of RAG
7+
)
8+
9+
try {
10+
Write-Host "Installing modules"
11+
Import-Module -Name SqlServer
12+
13+
# Import modules
14+
Import-Module SqlServer
15+
16+
Write-Host "Modules imported successfully."
17+
18+
# Authenticate using current logged in user
19+
az login
20+
$loggedInUser = (Get-AzADUser -SignedIn).UserPrincipalName
21+
22+
Set-AzSqlServerActiveDirectoryAdministrator -ResourceGroupName $resourceGroupName `
23+
-ServerName $serverName `
24+
-DisplayName $loggedInUser
25+
26+
Write-Host "Entra ID Admin set to: $loggedInUser"
27+
28+
# Define connection details
29+
$connectionString = "Server=tcp:${serverName}.database.windows.net,1433;Initial Catalog=$databaseName;Authentication=Active Directory Integrated;"
30+
31+
Write-Host "Database Connection String: $connectionString"
32+
33+
# Define the T-SQL script to add a new user and grant permissions
34+
$queryChartJs = @"
35+
IF NOT EXISTS (SELECT 1 FROM sys.database_principals WHERE name = '${chartJsFuncAppName}')
36+
BEGIN
37+
CREATE USER [${chartJsFuncAppName}] FROM EXTERNAL PROVIDER;
38+
ALTER ROLE db_datareader ADD MEMBER [${chartJsFuncAppName}]; -- Grant SELECT on all user tables and views.
39+
ALTER ROLE db_datawriter ADD MEMBER [${chartJsFuncAppName}]; -- Grant INSERT, UPDATE, and DELETE on all user tables and views.
40+
END
41+
"@
42+
43+
# Define the T-SQL script to add a new user and grant permissions
44+
$queryRag = @"
45+
IF NOT EXISTS (SELECT 1 FROM sys.database_principals WHERE name = '${ragFuncAppName}')
46+
BEGIN
47+
CREATE USER [${ragFuncAppName}] FROM EXTERNAL PROVIDER;
48+
ALTER ROLE db_datareader ADD MEMBER [${ragFuncAppName}]; -- Grant SELECT on all user tables and views.
49+
ALTER ROLE db_datawriter ADD MEMBER [${ragFuncAppName}]; -- Grant INSERT, UPDATE, and DELETE on all user tables and views.
50+
END
51+
"@
52+
53+
# ALTER ROLE db_ddladmin ADD MEMBER [${principalName}]; -- Grants CREATE, ALTER, and DROP on tables, views, functions, procedures, etc.
54+
55+
Write-Host "Executing SQL script for Chart JS Func..."
56+
Invoke-Sqlcmd -ConnectionString $connectionString -Query $queryChartJs
57+
Write-Host "SQL statement executed for Chart JS."
58+
59+
Write-Host "Executing SQL script for RAG.."
60+
Invoke-Sqlcmd -ConnectionString $connectionString -Query $queryRag
61+
62+
Write-Host "SQL statement executed for RAG."
63+
} catch {
64+
# Print the detailed error message
65+
Write-Host "An error occurred while querying the database:"
66+
Write-Host "Error Message: $($_.Exception.Message)"
67+
Write-Host "Error Type: $($_.Exception.GetType())"
68+
Write-Host "Stack Trace: $($_.Exception.StackTrace)"
69+
if ($_.Exception.InnerException) {
70+
Write-Host "Inner Exception: $($_.Exception.InnerException.Message)"
71+
}
72+
}

0 commit comments

Comments
 (0)