|
import argparse |
|
import logging |
|
|
|
import psycopg2 |
|
from azure.identity import DefaultAzureCredential |
|
|
|
logger = logging.getLogger("scripts") |
|
|
|
|
|
def assign_role_for_webapp(postgres_host, postgres_username, app_identity_name): |
|
if not postgres_host.endswith(".database.azure.com"): |
|
logger.info("This script is intended to be used with Azure Database for PostgreSQL.") |
|
logger.info("Please set the environment variable DBHOST to the Azure Database for PostgreSQL server hostname.") |
|
return |
|
|
|
logger.info("Authenticating to Azure Database for PostgreSQL using Azure Identity...") |
|
azure_credential = DefaultAzureCredential() |
|
token = azure_credential.get_token("https://ossrdbms-aad.database.windows.net/.default") |
|
conn = psycopg2.connect( |
|
database="postgres", |
|
user=postgres_username, |
|
password=token.token, |
|
host=postgres_host, |
|
sslmode="require", |
|
) |
|
|
|
conn.autocommit = True |
|
cur = conn.cursor() |
|
|
|
cur.execute(f"select * from pgaadauth_list_principals(false) WHERE rolname = '{app_identity_name}'") |
|
identities = cur.fetchall() |
|
if len(identities) > 0: |
|
logger.info(f"Found an existing PostgreSQL role for identity {app_identity_name}") |
|
else: |
|
logger.info(f"Creating a PostgreSQL role for identity {app_identity_name}") |
|
cur.execute(f"SELECT * FROM pgaadauth_create_principal('{app_identity_name}', false, false)") |
|
|
|
logger.info(f"Granting permissions to {app_identity_name}") |
|
|
|
cur.execute(f'GRANT USAGE ON SCHEMA public TO "{app_identity_name}"') |
|
cur.execute(f'GRANT CREATE ON SCHEMA public TO "{app_identity_name}"') |
|
cur.execute(f'GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO "{app_identity_name}"') |
|
cur.execute( |
|
f"ALTER DEFAULT PRIVILEGES IN SCHEMA public " |
|
f'GRANT SELECT, UPDATE, INSERT, DELETE ON TABLES TO "{app_identity_name}"' |
|
) |
|
|
|
cur.close() |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
logging.basicConfig(level=logging.WARNING) |
|
logger.setLevel(logging.INFO) |
|
parser = argparse.ArgumentParser(description="Create database schema") |
|
parser.add_argument("--host", type=str, help="Postgres host") |
|
parser.add_argument("--username", type=str, help="Postgres username") |
|
parser.add_argument("--app-identity-name", type=str, help="Azure App Service identity name") |
|
|
|
args = parser.parse_args() |
|
if not args.host.endswith(".database.azure.com"): |
|
logger.info("This script is intended to be used with Azure Database for PostgreSQL, not local PostgreSQL.") |
|
exit(1) |
|
|
|
assign_role_for_webapp(args.host, args.username, args.app_identity_name) |
|
logger.info("Role created successfully.") |
|
|
