# Copyright 2022, Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Amazon Software License (the "License"). # You may not use this file except in compliance with the License. # A copy of the License is located at # # http://aws.amazon.com/asl/ # # or in the "license" file accompanying this file. This file is distributed # on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either # express or implied. See the License for the specific language governing # permissions and limitations under the License. # Glue Job to collect Lake Formation permissions data for the given account/region. # Collected data is stored as Parquet in S3 bucket and Athena tables are created on the data. # Job takes the following parameters: # catalog-id : Current AWS Account ID # databasename : Glue Database Name # region : Current AWS region # s3bucket: S3 bucket Name where data is stored. Data is stored under the prefix:lfpermissions-data. # createtable: yes (This parameter enables Athena table creation) from pyathena import connect import sys import json from typing import Tuple, List, Dict import boto3 import pandas as pd from boto3.session import Session import awswrangler from pandas import json_normalize,DataFrame from datetime import datetime import time import re import argparse import sys import os class LakeFormation: def __init__(self,region: str): self.lf_client: LakeFormationClient = Session(region_name=region).client("lakeformation") self.athena_client = Session(region_name=region).client("athena") self.s3 = boto3.resource('s3',region_name=region) # Removes objects under the prefix in given S3 bucket def cleanup_lf_schema(self, s3bucket: str, prefix : str): print("[INFO] Permissions Data cleanup - Bucket: {} , Prefix: {}".format(s3bucket,prefix)) bucket = self.s3.Bucket(s3bucket) for obj in bucket.objects.filter(Prefix=prefix): self.s3.Object(bucket.name,obj.key).delete() # Submits the query to Athena and returns the status def _run_athena_sql(self,athenaparams): s3stagingdir='s3://' + athenaparams['bucket'] + '/staging/' cursor = connect(s3_staging_dir=s3stagingdir,region_name=athenaparams['region'],schema_name=athenaparams['database']).cursor() cursor.execute(athenaparams['query']) print("[INFO] Executing Athena query with description :{}".format(cursor.description)) cursor.fetchall() # Creates Summary Table on the permissions data with provided partition keys. def create_summary_table(self,region: str,databasename: str, s3bucket : str, datalocation:str): print("[INFO] Permissions Summary Table creation in Database {} Begins".format(databasename)) query_string= '''CREATE EXTERNAL TABLE IF NOT EXISTS `lfpermissions_summary`( `permissions` array, `permissionswithgrantoption` array, `database.name` string, `database.catalogid` string, `table.databasename` string, `table.catalogid` string, `table.name` string, `lfcatalogid` string, `principal` string, `lfresourcetype` string, `principal.datalakeprincipalidentifier` string, `resource.datacellsfilter.tablecatalogid` string, `resource.datacellsfilter.databasename` string, `resource.datacellsfilter.tablename` string, `resource.datacellsfilter.name` string, `resource.database.catalogid` string, `resource.database.name` string, `resource.datalocation.catalogid` string, `resource.datalocation.resourcearn` string, `resource.table.catalogid` string, `resource.table.databasename` string, `resource.table.name` string, `resource.tablewithcolumns.catalogid` string, `resource.tablewithcolumns.databasename` string, `resource.tablewithcolumns.name` string, `resource.tablewithcolumns.columnnames` array, `resource.lftag.catalogid` string, `resource.lftag.tagkey` string, `resource.lftag.tagvalues` array, `additionaldetails.resourceshare` array, `resource.tablewithcolumns.columnwildcard.excludedcolumnnames` array, `lftagexpression` array>>) PARTITIONED BY ( catalog_id string, year string, month string, day string, resource_type string) STORED AS PARQUET LOCATION '{}' TBLPROPERTIES ( 'classification'='parquet' ) ''' new_query_string = query_string.format(datalocation) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} athenaStatus= self._run_athena_sql(athenaparams=params) query_string='''MSCK REPAIR TABLE lfpermissions_summary''' params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': query_string} self._run_athena_sql(athenaparams=params) print("[INFO] Permissions Summary Table creation in Database {} Completed".format(databasename)) # Creates permissions table with latest snapshot of permissions data for the dashboard def create_resourcepermission_tables(self,region: str,databasename: str, s3bucket : str, year:str, month: str,day : str): print("[INFO] Permissions Dashboard Table-lfpermissions in Database {} Begins".format(databasename)) milliseconds = str(round(time.time() * 1000)) tablepath=f's3://{s3bucket}/athenaoutput/data{year}{month}{day}{milliseconds}' query_string= '''UNLOAD (select distinct permission, coalesce("principal.datalakeprincipalidentifier","principal") as principal, coalesce("lfcatalogid","resource.datacellsfilter.tablecatalogid","resource.database.catalogid","resource.datalocation.catalogid","resource.table.catalogid", "resource.tablewithcolumns.catalogid","resource.lftag.catalogid") as owner_catalog_id, coalesce("resource.datacellsfilter.databasename","database.name","table.databasename","resource.database.name", "resource.table.databasename", "resource.tablewithcolumns.databasename") as databasename, coalesce("resource.datacellsfilter.tablename" , "table.name","resource.table.name", "resource.tablewithcolumns.name" ) as tablename, "resource.datalocation.resourcearn" as datalocation, coalesce("lfresourcetype","resource_type") as lftype, "lftagexpression", catalog_id, resource_type from lfpermissions_summary cross join unnest(permissions) as t(permission) where coalesce("principal.datalakeprincipalidentifier","principal") != 'IAM_ALLOWED_PRINCIPALS' and year = '{}' and month ='{}' and day ='{}') TO '{}' WITH (format = 'PARQUET',compression = 'SNAPPY')''' new_query_string = query_string.format(year,month,day,tablepath) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} self._run_athena_sql(athenaparams=params) query_string= ''' CREATE EXTERNAL TABLE IF NOT EXISTS `lfpermissions`( `permission` string , `principal` string , `owner_catalog_id` string , `databasename` string , `tablename` string , `datalocation` string , `lftype` string , `lftagexpression` array>> , `catalog_id` string , `resource_type` string ) ROW FORMAT SERDE 'org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe' STORED AS INPUTFORMAT 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat' OUTPUTFORMAT 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat' LOCATION '{}' TBLPROPERTIES ( 'classification'='parquet') ''' new_query_string = query_string.format(tablepath) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} self._run_athena_sql(athenaparams=params) query_string='''ALTER TABLE lfpermissions SET LOCATION '{}' ''' new_query_string = query_string.format(tablepath) print("INFO] Permissions Dashboard Table-lfpermissions Alter Statement: {}".format(new_query_string)) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} self._run_athena_sql(athenaparams=params) print("[INFO] Permissions Dashboard Table-lfpermissions creation completed") # Creates grantable permissions table with latest snapshot of grantable permissions data for the dashboard def create_resourcepermissionwithgrant_tables(self,region: str,databasename: str, s3bucket : str, year:str, month: str,day : str): print("[INFO] Permissions Dashboard Table-lfpermissionswithgrant in Database {} Begins".format(databasename)) milliseconds = str(round(time.time() * 1000)) tablepath=f's3://{s3bucket}/athenaoutput/grantdata{year}{month}{day}{milliseconds}' query_string= '''UNLOAD (select distinct permission, coalesce("principal.datalakeprincipalidentifier","principal") as principal, coalesce("lfcatalogid","resource.datacellsfilter.tablecatalogid","resource.database.catalogid","resource.datalocation.catalogid","resource.table.catalogid", "resource.tablewithcolumns.catalogid","resource.lftag.catalogid") as owner_catalog_id, coalesce("resource.datacellsfilter.databasename","database.name","table.databasename","resource.database.name", "resource.table.databasename", "resource.tablewithcolumns.databasename") as databasename, coalesce("resource.datacellsfilter.tablename" , "table.name","resource.table.name", "resource.tablewithcolumns.name" ) as tablename, "resource.datalocation.resourcearn" as datalocation, coalesce("lfresourcetype","resource_type") as lftype, "lftagexpression", catalog_id, resource_type from lfpermissions_summary cross join unnest(permissionswithgrantoption) as t(permission) where coalesce("principal.datalakeprincipalidentifier","principal") != 'IAM_ALLOWED_PRINCIPALS' and year = '{}' and month ='{}' and day ='{}') TO '{}' WITH (format = 'PARQUET',compression = 'SNAPPY')''' new_query_string = query_string.format(year,month,day,tablepath) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} self._run_athena_sql(athenaparams=params) query_string= ''' CREATE EXTERNAL TABLE IF NOT EXISTS `lfpermissionswithgrant`( `permission` string , `principal` string , `owner_catalog_id` string , `databasename` string , `tablename` string , `datalocation` string , `lftype` string , `lftagexpression` array>> , `catalog_id` string , `resource_type` string ) ROW FORMAT SERDE 'org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe' STORED AS INPUTFORMAT 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat' OUTPUTFORMAT 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat' LOCATION '{}' TBLPROPERTIES ( 'classification'='parquet') ''' new_query_string = query_string.format(tablepath) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} self._run_athena_sql(athenaparams=params) query_string='''ALTER TABLE lfpermissionswithgrant SET LOCATION '{}' ''' new_query_string = query_string.format(tablepath) print("INFO] Permissions Dashboard Table-lfpermissionswithgrant Alter Statement: {}".format(new_query_string)) params = {'region': region,'database': databasename,'bucket': s3bucket,'path': 'athenaresults','query': new_query_string} athenaquerystatus = self._run_athena_sql(athenaparams=params) print("[INFO] Permissions Dashboard Table-lfpermissionswithgrant creation completed") # Retrieves the list of permissions for given resource type def _get_permissions(self, catalog_id: str, resource_type: str = None, next_token: str = "") -> Tuple[str, DataFrame]: print("_get_permissions :: start") print(f"Next token before = {next_token}") request = {"CatalogId": catalog_id, "NextToken": next_token,"MaxResults": 500} if resource_type: request["ResourceType"] = resource_type list_permission: ListPermissionsResponseTypeDef = self.lf_client.list_permissions(**request) next_token = list_permission.get("NextToken") df = json_normalize(list_permission, "PrincipalResourcePermissions") pd.set_option('display.max_columns', None) print(f"Next token after = {next_token}") print("_get_permissions :: end") return next_token, df # Retrieves the list of resources matching LF Tag Policy def _get_resource_by_lf_tags(self, catalog_id: str, lf_tag_expression: List, resource_type: str, next_token: str = "") -> Tuple[str, DataFrame]: print("_get_resource_by_lf_tags :: start") print(f"Searching for expression = {lf_tag_expression}, resource_type = {resource_type}") if resource_type == "DATABASE": response = self.lf_client.search_databases_by_lf_tags(Expression=lf_tag_expression, CatalogId=catalog_id,NextToken=next_token) next_token = response.get("NextToken") df = json_normalize(response, "DatabaseList") elif resource_type == "TABLE": response = self.lf_client.search_tables_by_lf_tags(Expression=lf_tag_expression, CatalogId=catalog_id,NextToken=next_token) next_token = response.get("NextToken") df = json_normalize(response, "TableList") else: print("[ERROR]: _get_resource_by_lf_tags :: Invalid resource type {}:".format(resource_type)) pd.set_option('display.max_columns', None) print("_get_resource_by_lf_tags :: end") return next_token, df # Retrieves the list of permissions granted via Named resources def download_resource_named_resourcetype(self, catalog_id: str, resource_type: str, s3_path: str = None): print("download_resource_named_resourcetype for {} :: start".format(resource_type)) next_token = "" now = datetime.now() dt_string = now.strftime("%d%m%Y%H%M%S") while True: next_token, df = self._get_permissions(catalog_id=catalog_id, resource_type=resource_type, next_token=next_token) pd.set_option('display.max_columns', None) if df.size == 0: print("download_resource_named_resourcetype for {} :: Empty Result".format(resource_type)) else: milliseconds = str(round(time.time() * 1000)) awswrangler.s3.to_parquet(df=df, path=s3_path+ "/resource_type="+ resource_type+"/data"+dt_string+milliseconds+".parquet", index=False) if next_token is None: break print("download_resource_named_resourcetype for {} :: End".format(resource_type)) # Retrieves the list of permissions granted via LF Tag Policy def download_resource_lftags_resourcetype(self, catalog_id: str, lftag_resource_type: str, s3_path: str = None): print("download_resource_lftags_resourcetype for {} :: Start".format(lftag_resource_type)) next_token = "" now = datetime.now() dt_string = now.strftime("%d%m%Y%H%M%S") while True: next_token, df = self._get_permissions(catalog_id=catalog_id, resource_type=lftag_resource_type, next_token=next_token) pd.set_option('display.max_columns', None) if df.size == 0: print("download_resource_lftags_resourcetype for {} :: Empty Result".format(lftag_resource_type)) else: df = df[['Resource.LFTagPolicy.ResourceType', 'Resource.LFTagPolicy.Expression','Resource.LFTagPolicy.CatalogId','Principal.DataLakePrincipalIdentifier','Permissions','PermissionsWithGrantOption']] df['Resource.LFTagPolicy.Expression'] = df['Resource.LFTagPolicy.Expression'].astype(str) result_df = pd.DataFrame() for index, row in df.iterrows(): expression = list(eval(row['Resource.LFTagPolicy.Expression'])) lfresourcetype = row['Resource.LFTagPolicy.ResourceType'] lfcatalogid = row['Resource.LFTagPolicy.CatalogId'] if(lfcatalogid == catalog_id): lf_next_token = "" while True: lf_next_token, iter_df = self._get_resource_by_lf_tags(catalog_id=lfcatalogid, lf_tag_expression=expression, resource_type=lfresourcetype, next_token=lf_next_token) if iter_df.size == 0: print("_get_resource_by_lf_tags for {} :: Empty Result".format(expression)) else: iter_df['lftagexpression'] = [expression for _ in range(len(iter_df))] iter_df['lfcatalogid'] = [row['Resource.LFTagPolicy.CatalogId'] for _ in range(len(iter_df))] iter_df['principal'] = [row['Principal.DataLakePrincipalIdentifier'] for _ in range(len(iter_df))] iter_df['permissions'] = [row['Permissions'] for _ in range(len(iter_df))] iter_df['lfresourcetype'] = [row['Resource.LFTagPolicy.ResourceType'] for _ in range(len(iter_df))] iter_df['permissionswithgrantoption'] = [row['PermissionsWithGrantOption'] for _ in range(len(iter_df))] result_df = pd.concat([result_df, iter_df]) if lf_next_token is None: break if result_df.size >0 : milliseconds = str(round(time.time() * 1000)) awswrangler.s3.to_parquet(df=result_df, path=s3_path+ "/resource_type="+ lftag_resource_type+"/data"+dt_string+milliseconds+".parquet", index=False) if next_token is None: break print("download_resource_lftags_resourcetype for {} :: End".format(lftag_resource_type)) def get_job_args(required_args, optional_args): lst = sys.argv reqd_args_map = get_resolved_options(lst, required_args, True) opt_args_map = {} for a in optional_args: opt_args_map[a] = "" try: r = get_resolved_options(lst, [a], False) if a in r: opt_args_map[a] = r[a] elif a.replace("-", "_") in r: opt_args_map[a.replace("-", "_")] = r[a.replace("-", "_")] except: print("Ignoring key " + a) merged_map = {} for k, v in opt_args_map.items(): merged_map[k] = v for k,v in reqd_args_map.items(): merged_map[k] = v return merged_map def get_resolved_options(args, options, required=True): parser = argparse.ArgumentParser() for option in options: parser.add_argument('--' + option, required=required) parsed, extra = parser.parse_known_args(args) return vars(parsed) def main(): now = datetime.now() dt_year = now.strftime("%Y") dt_month = now.strftime("%m") dt_day = now.strftime("%d") print("starting") args = get_job_args(['catalog-id','region','s3bucket','databasename','createtable'], []) print("startingm1") region = args['region'] catalog_id = args['catalog_id'] s3bucket = args['s3bucket'] databasename = args['databasename'] createtable = args['createtable'] datapath = "lfpermissions-data" schemapath = "lfpermissions-schema" s3outputbucket = "s3://"+s3bucket+"/" s3_data_prefix = datapath+"/catalog_id="+ catalog_id+"/year=" +dt_year+"/month="+dt_month+"/day="+dt_day s3_path = s3outputbucket+s3_data_prefix s3_data_location=s3outputbucket+datapath print("[INFO] Lake Formation Permission Exporter Begins For Date - {}/{}/{}".format(dt_month,dt_day,dt_year)) lf = LakeFormation(region) lf.cleanup_lf_schema(s3bucket=s3bucket,prefix=s3_data_prefix) lf.download_resource_lftags_resourcetype(catalog_id=catalog_id, lftag_resource_type="LF_TAG_POLICY_TABLE",s3_path=s3_path) lf.download_resource_lftags_resourcetype(catalog_id=catalog_id, lftag_resource_type="LF_TAG_POLICY_DATABASE",s3_path=s3_path) lf.download_resource_named_resourcetype(catalog_id=catalog_id, resource_type="DATABASE",s3_path=s3_path) lf.download_resource_named_resourcetype(catalog_id=catalog_id, resource_type="TABLE",s3_path=s3_path) lf.download_resource_named_resourcetype(catalog_id=catalog_id, resource_type="DATA_LOCATION",s3_path=s3_path) if(createtable == "yes"): print("[INFO] Lake Formation Permission Dashboard Tables Begins") lf.create_summary_table(region=region,databasename=databasename, s3bucket=s3bucket, datalocation=s3_data_location) lf.create_resourcepermission_tables(region=region,databasename=databasename, s3bucket=s3bucket, year=dt_year,month=dt_month,day=dt_day) lf.create_resourcepermissionwithgrant_tables(region=region,databasename=databasename, s3bucket=s3bucket, year=dt_year,month=dt_month,day=dt_day) print("[INFO] Lake Formation Permission Dashboard Tables End") if __name__ == "__main__": main()