Skip to content

Commit 6b2f16f

Browse files
refactor: Remove unused functions and streamline permission granting in create_postgres_tables.py
1 parent 93b730a commit 6b2f16f

File tree

1 file changed

+8
-103
lines changed

1 file changed

+8
-103
lines changed

scripts/data_scripts/create_postgres_tables.py

Lines changed: 8 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -1,111 +1,13 @@
11
from azure_credential_utils import get_azure_credential
22
import psycopg2
33
from psycopg2 import sql
4-
import requests
54

65
principalId = "userPrincipalId"
76
user = "managedIdentityName"
87
host = "serverName"
98
dbname = "postgres"
109

1110

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-
10911
# Acquire the access token
11012
cred = get_azure_credential()
11113
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):
11719
conn = psycopg2.connect(conn_string)
11820
cursor = conn.cursor()
11921

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+
12030
# Drop and recreate the conversations table
12131
cursor.execute("DROP TABLE IF EXISTS conversations")
12232
conn.commit()
@@ -181,11 +91,6 @@ def grant_permissions(cursor, dbname, schema_name, principal_id):
18191
)
18292
conn.commit()
18393

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()
18994

19095
cursor.execute("ALTER TABLE public.conversations OWNER TO azure_pg_admin;")
19196
cursor.execute("ALTER TABLE public.messages OWNER TO azure_pg_admin;")

0 commit comments

Comments
 (0)