1
1
from azure_credential_utils import get_azure_credential
2
2
import psycopg2
3
3
from psycopg2 import sql
4
- import requests
5
4
6
5
principalId = "userPrincipalId"
7
6
user = "managedIdentityName"
8
7
host = "serverName"
9
8
dbname = "postgres"
10
9
11
10
12
- def get_user_principal_name (principal_id , credential ):
13
- """
14
- Get user principal name (UPN/email) from a principal ID using Microsoft Graph API.
15
-
16
- Parameters:
17
- - principal_id: The Azure AD object ID of the user or service principal
18
- - credential: Azure credential object for authentication
19
-
20
- Returns:
21
- - UPN/email of the user or the principal_id if not found
22
- """
23
- try :
24
- # Get access token for Microsoft Graph
25
- token = credential .get_token ("https://graph.microsoft.com/.default" )
26
-
27
- headers = {
28
- "Authorization" : f"Bearer { token .token } " ,
29
- "Content-Type" : "application/json"
30
- }
31
-
32
- # Try to get user details
33
- user_url = f"https://graph.microsoft.com/v1.0/users/{ principal_id } "
34
- response = requests .get (user_url , headers = headers , timeout = 30 )
35
-
36
- if response .status_code == 200 :
37
- user_data = response .json ()
38
- # Return userPrincipalName if available, otherwise return mail
39
- return user_data .get ("userPrincipalName" , user_data .get ("mail" , principal_id ))
40
-
41
- # If not found as user, try as service principal
42
- sp_url = f"https://graph.microsoft.com/v1.0/servicePrincipals/{ principal_id } "
43
- response = requests .get (sp_url , headers = headers , timeout = 30 )
44
-
45
- if response .status_code == 200 :
46
- sp_data = response .json ()
47
- # For service principals, use appId or displayName as identifier
48
- return sp_data .get ("displayName" , principal_id )
49
-
50
- # If neither found, return the ID itself
51
- print (f"Could not find UPN for principal ID { principal_id } . Using ID as fallback." )
52
- return principal_id
53
-
54
- except Exception as e :
55
- print (f"Error retrieving UPN: { str (e )} " )
56
- return principal_id
57
-
58
-
59
- def grant_permissions (cursor , dbname , schema_name , principal_id ):
60
- """
61
- Grants database and schema-level permissions to a specified principal.
62
-
63
- Parameters:
64
- - cursor: psycopg2 cursor object for database operations.
65
- - dbname: Name of the database to grant CONNECT permission.
66
- - schema_name: Name of the schema to grant table-level permissions.
67
- - principal_id: ID of the principal (role or user) to grant permissions.
68
- """
69
-
70
- print (f"Granting permissions to principal '{ principal_id } '" )
71
- # Check if the principal exists in the database
72
- cursor .execute (
73
- sql .SQL ("SELECT 1 FROM pg_roles WHERE rolname = {principal}" ).format (
74
- principal = sql .Literal (principal_id )
75
- )
76
- )
77
- if cursor .fetchone () is None :
78
- add_principal_user_query = sql .SQL (
79
- "SELECT * FROM pgaadauth_create_principal({principal}, false, false)"
80
- )
81
- cursor .execute (
82
- add_principal_user_query .format (
83
- principal = sql .Literal (principal_id ),
84
- )
85
- )
86
-
87
- # Grant CONNECT on database
88
- grant_connect_query = sql .SQL ("GRANT CONNECT ON DATABASE {database} TO {principal}" )
89
- cursor .execute (
90
- grant_connect_query .format (
91
- database = sql .Identifier (dbname ),
92
- principal = sql .Identifier (principal_id ),
93
- )
94
- )
95
- print (f"Granted CONNECT on database '{ dbname } ' to '{ principal_id } '" )
96
-
97
- # Grant SELECT, INSERT, UPDATE, DELETE on schema tables
98
- grant_permissions_query = sql .SQL (
99
- "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {schema} TO {principal}"
100
- )
101
- cursor .execute (
102
- grant_permissions_query .format (
103
- schema = sql .Identifier (schema_name ),
104
- principal = sql .Identifier (principal_id ),
105
- )
106
- )
107
-
108
-
109
11
# Acquire the access token
110
12
cred = get_azure_credential ()
111
13
access_token = cred .get_token ("https://ossrdbms-aad.database.windows.net/.default" )
@@ -117,6 +19,14 @@ def grant_permissions(cursor, dbname, schema_name, principal_id):
117
19
conn = psycopg2 .connect (conn_string )
118
20
cursor = conn .cursor ()
119
21
22
+ # Grant admin privileges to the principal
23
+ cursor .execute (
24
+ sql .SQL ("ALTER ROLE {} WITH SUPERUSER" ).format (
25
+ sql .Identifier (principalId )
26
+ )
27
+ )
28
+ conn .commit ()
29
+
120
30
# Drop and recreate the conversations table
121
31
cursor .execute ("DROP TABLE IF EXISTS conversations" )
122
32
conn .commit ()
@@ -181,11 +91,6 @@ def grant_permissions(cursor, dbname, schema_name, principal_id):
181
91
)
182
92
conn .commit ()
183
93
184
- if principalId and principalId .strip ():
185
- identifier_to_use = get_user_principal_name (principalId , cred )
186
- print (f"Using identifier '{ identifier_to_use } ' for principal ID '{ principalId } '" )
187
- grant_permissions (cursor , dbname , "public" , identifier_to_use )
188
- conn .commit ()
189
94
190
95
cursor .execute ("ALTER TABLE public.conversations OWNER TO azure_pg_admin;" )
191
96
cursor .execute ("ALTER TABLE public.messages OWNER TO azure_pg_admin;" )
0 commit comments