 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: MIT-0
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this
 # software and associated documentation files (the "Software"), to deal in the Software
 # without restriction, including without limitation the rights to use, copy, modify,
 # merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 # permit persons to whom the Software is furnished to do so.
 #
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 # INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
 # PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 # HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import boto3
import sys
import json
import time
import botocore
from tqdm import tqdm

def execute_query(clusterid, db, secret, query):
    try:
        client = boto3.client('redshift-data')
        response = client.execute_statement(
            ClusterIdentifier=clusterid,
            Database=db,
            SecretArn=secret,
            Sql=query
        )

        while True:
            response=client.describe_statement(Id=response['Id'])
            if (response["Status"] == "FINISHED" or response["Status"] == "ABORTED" or response["Status"] == "FAILED"):
                #print("Query Status: {0}".format(response["Status"]))
                if(response["Status"] == "FAILED"):
                    print("\nEXECUTION ERROR {0}".format(response["Error"]))
                    raise Exception
                elif(response["Status"] == "ABORTED"):
                    print("\n\tQuery has been aborted\n")
                    raise Exception
                #else:
                #    print("EXECUTION ERROR: None")
                break

            else:
                #print("Query Status: {0}".format(response["Status"]))
                time.sleep(2)

        if (response["HasResultSet"] == True):
            result=client.get_statement_result(
                Id=response['Id']
            )
            return result

    except client.exceptions.ValidationException as Ve:
        sys.exit("Validation Exception: \n\n\t{0}".format(Ve))

    except client.exceptions.ExecuteStatementException as ESe:
        sys.exit("\nExecute Statement Exception: \n\n\t{0}\n".format(ESe))
    
    except client.exceptions.ActiveStatementsExceededException as ASEe:
        sys.exit("\nActive Statements Exceeded Exception: \n\n\t{0}\n".format(ASEe))
    
    except client.exceptions.ResourceNotFoundException as RNFe:
        sys.exit("\n Resource not Found Exception: \n\n\t{0}\n".format(RNFe))

    except client.exceptions.InternalServerException as ISe:
        sys.exit("\nInternal Server Exception: \n\n\t{0}\n".format(ISe))
    
    except client.exceptions.ClientError as CEe:
        sys.exit("\nClient Error Exception: \n\n\t{0}\n".format(CEe))

    except:
        sys.exit("{0}".format(sys.exc_info()[1]))

def create_control_table(clusterid, db, secret):

    createQuery = "CREATE TABLE IF NOT EXISTS public.encoding_control \
    (  \
	schema_name character varying(100) NOT NULL ENCODE lzo, \
	table_name character varying(100) NOT NULL ENCODE lzo,   \
	column_name character varying(100) NOT NULL ENCODE lzo,  \
	curr_encoding_type character varying(20) NOT NULL ENCODE lzo, \
    recmd_encoding_type character varying(20) NOT NULL DEFAULT 'unknown' ENCODE lzo, \
    est_red_pct character varying(20) NOT NULL DEFAULT '0.00' ENCODE lzo, \
    analyze_ind character(1) NOT NULL DEFAULT 'N'::bpchar ENCODE lzo, \
	alter_status character(1) NOT NULL DEFAULT 'N'::bpchar ENCODE lzo,   \
	update_timestamp timestamp without time zone ENCODE az64 \
    )  \
    DISTSTYLE AUTO;"

    for i in tqdm (range (1), desc="Creating Control Table if not exists.."):
        execute_query(clusterid, db, secret, createQuery)

def truncate_control_table(clusterid, db, secret):

    truncateQuery = "truncate public.encoding_control"

    for i in tqdm (range (1), desc="Truncating Control Table:.."):
        execute_query(clusterid, db, secret, truncateQuery)

def load_control_table(clusterid, db, secret, schemaName):

    listTables = "select schemaname \
    ,tablename \
    ,"'"column"'" as columnname \
    ,encoding \
    from pg_table_def \
    where schemaname='"+schemaName+"' \
    and tablename<>'encoding_control'"

    result = execute_query(clusterid, db, secret, listTables)
    if(result["Records"] != []):
        for i in tqdm (range (1), desc="Loading Control Table for Schema {0}..".format(schemaName)):
            for datarecord in result["Records"]:
                loadControlQuery="insert into public.encoding_control \
                values("
                for column in datarecord:
                    loadControlQuery+="'"+column['stringValue']+"',"
                
                loadControlQuery+="default,default,default,default,current_timestamp);"
                execute_query(clusterid, db, secret, loadControlQuery)
    else:
        sys.exit("The schema does not exist or does not have any tables")

def is_table_large(clusterid, db, secret, schema, table, threshold):
    tableSizeQuery="select "'"schema"'" as schemaname \
    ,"'"table"'" as tablename \
    ,size as sizeinMB \
    FROM SVV_TABLE_INFO \
    where schemaname = '"+schema+"' \
    and tablename = '"+table+"';"

    result = execute_query(clusterid, db, secret, tableSizeQuery)
    if(result["Records"] != []):
        tablesize=result["Records"][0][2]['longValue']
        if((tablesize/1024/1024)>int(threshold)):
            return True
    
    return False

def first_sort_key(clusterid, db, secret, schema, table):
    sortKeyQuery="select "'"column"'" \
    FROM pg_table_def \
    where schemaname = '"+schema+"' \
    and tablename = '"+table+"' \
    and sortkey=1;"

    result = execute_query(clusterid, db, secret, sortKeyQuery)
    if(result["Records"] != []):
        return result["Records"][0][0]['stringValue']
    
    return ''



def update_control_table(clusterid, db, secret, schema_name, table_name, column_name, encoding, est_red_pct, analyze_ind, alter_status):
    updateQuery="update public.encoding_control \
    set recmd_encoding_type='"+encoding+"',   \
    est_red_pct='"+est_red_pct+"',   \
    analyze_ind='"+analyze_ind+"',    \
    alter_status='"+alter_status+"',  \
    update_timestamp=current_timestamp   \
    where schema_name='"+schema_name+"' \
    and table_name='"+table_name+"'    \
    and column_name='"+column_name+"';"

    result = execute_query(clusterid, db, secret, updateQuery)
    return result


def analyze_compression(clusterid, db, secret, schema, table):
    analyzeQuery='analyze compression '+schema+'.'+table+';'
    result = execute_query(clusterid, db, secret, analyzeQuery)
    if(result != None):
        return result["Records"]
    
    return result

def alter_compression(clusterid, db, secret, schema, table, column, encoding):

    alterQuery=" alter table "+schema+"."+table+" alter column "+column+" encode "+encoding+";"
    result = execute_query(clusterid, db, secret, alterQuery)
    return result

def get_control_table(clusterid, db, secret, schema_name):
    
    getTableQuery="select distinct table_name \
    from public.encoding_control    \
    where schema_name='"+schema_name+"';"
    
    result = execute_query(clusterid, db, secret, getTableQuery)
    if(result["Records"] != []):
        return result["Records"]
    

def compress_all(clusterid, db, secret,schema_name):
    result=get_control_table(clusterid, db, secret, schema_name)
    if(result != None):
        for i in tqdm (range (1), desc="Applying Compression for all Required Tables in Schema {0}..".format(schema_name)):
            for table in result:
                table_name=table[0]["stringValue"]
                fSortKey=first_sort_key(clusterid, db, secret, schema_name,table_name)
                analyze_result=analyze_compression(clusterid, db, secret, schema_name, table_name)
                if(analyze_result != None):
                    for encoding in analyze_result:
                        if(float(encoding[3]["stringValue"])>0 and encoding[1]["stringValue"]!=fSortKey):
                            alter_compression(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"])
                            update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'Y')
                        else:
                            update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'N')

    
def compress_large(clusterid, db, secret,schema_name,threshold):
    result=get_control_table(clusterid, db, secret, schema_name)
    if(result != None):
        for i in tqdm (range (1), desc="Applying Compression for Required Tables > {0} TB in Schema {1}..".format(threshold,schema_name)):
            for table in result:
                table_name=table[0]["stringValue"]
                fSortKey=first_sort_key(clusterid, db, secret, schema_name,table_name)
                if(is_table_large(clusterid, db, secret, schema_name, table_name, threshold)):
                    analyze_result=analyze_compression(clusterid, db, secret, schema_name, table_name)
                    if(analyze_result != None):
                        for encoding in analyze_result:
                            if(float(encoding[3]["stringValue"])>0 and encoding[1]["stringValue"]!=fSortKey):
                                alter_compression(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"])
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'Y')
                            else:
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'N')  

def compress_small(clusterid, db, secret,schema_name, threshold):
    result=get_control_table(clusterid, db, secret, schema_name)
    if(result != None):
        for i in tqdm (range (1), desc="Applying Compression for Required Tables < {0} TB in Schema {1}..".format(threshold,schema_name)):
            for table in result:
                table_name=table[0]["stringValue"]
                fSortKey=first_sort_key(clusterid, db, secret, schema_name,table_name)
                if(not is_table_large(clusterid, db, secret, schema_name, table_name, threshold)):
                    analyze_result=analyze_compression(clusterid, db, secret, schema_name, table_name)
                    if(analyze_result != None):
                        for encoding in analyze_result:
                            if(float(encoding[3]["stringValue"])>0 and encoding[1]["stringValue"]!=fSortKey):
                                alter_compression(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"])
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'Y')
                            else:
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'N')

        
        
    
    

def main(argv):
    
    try:
        if(len(sys.argv) == 1):
            print(
                '\nUsage: Alter-Compression-Utility.py   \
                \n\tThe utility can automatically analyze schema, identify tables and alter column encoding compression leveraging Alter Compression feature of Redshift released in Oct, 20   \
                \n\t\tReference Document: https://aws.amazon.com/about-aws/whats-new/2020/10/amazon-redshift-supports-modifying-column-comprression-encodings-to-optimize-storage-utilization-query-performance/ \
                \n\nArguments:\n \
                \n\targv[1] - Cluster Identifier: Name of the Cluster \
                \n\targv[2] - Database Name: Default database name that is required for connecting to cluster \
                \n\targv[3] - Secret Identifier: Secret ARN/Friendly name of the secret for retrieving user and password of the cluster \
                \n\targv[4] - Schema Name: Name of the schema that user wants to analyze and compress tables for  \
                \n\targv[5] - Compression Mode: Three modes of compression for users to choose from and pass as argument \
                \n\n\t\tcompress-all: Automatically analyze and compress required columns for all the tables for the schema    \
                \n\t\tcompress-small: Based on provided threshold value, Automatically analyze and compress required columns for tables < threshold value in TB for the schema   \
                \n\t\tcompress-large: Based on provided threshold value,  Automatically analyze and compress required columns for tables > threshold value in TB for the schema  \
                \n\n\targv[6] - Threshold value: Threshold value upon which small/large table size in TBs will be defined.The value is not required for compress-all mode.\n'
            )

        elif(len(sys.argv) > 7):
            print('\n\tPassed arguments have exceeded the limit. Please refer to the utility documentation\n')

        else:
            print('\n\n')
            if (argv[5]=='compress-all'):
                create_control_table(argv[1], argv[2], argv[3])
                truncate_control_table(argv[1], argv[2], argv[3])
                load_control_table(argv[1], argv[2], argv[3], argv[4])
                compress_all(argv[1], argv[2], argv[3],argv[4])
                print("\nComplete: \n\tCompression has successfully been altered for all required tables in schema {0} and public.encoding_control has been updated for detail. \n".format(argv[4]))

            elif(int(argv[6])>0 and (argv[5]=='compress-small' or argv[5]=='compress-large')):

                create_control_table(argv[1], argv[2], argv[3])
                truncate_control_table(argv[1], argv[2], argv[3])
                load_control_table(argv[1], argv[2], argv[3], argv[4])
                if(argv[5]=='compress-small'):
                    compress_small(argv[1], argv[2], argv[3],argv[4],argv[6])
                    print("\nComplete: \n\tCompression has successfully been altered for all required tables < {0} TB in schema {1} and public.encoding_control has been updated for detail. \n".format(argv[6], argv[4]))

                elif(argv[5]=='compress-large'):
                    compress_large(argv[1], argv[2], argv[3],argv[4],argv[6])
                    print("\nComplete: \n\tCompression has successfully been altered for all required tables > {0} TB in schema {1} \
                    and public.encoding_control has been updated for detail. \n".format(argv[6], argv[4]))

            else:
                if (argv[5]!='compress-all' and int(argv[6])<=0):
                    print("\nPlease pass a table size (TB) threshold value greater than zero. Refer to the utility documentation !\n")
                else:
                    print("\nPlease pass a valid compression mode. Refer to the utility documentation !\n")

    except IndexError as e:
        print("\nPlease pass valid command line arguments. Refer to the utility documentation ! \
        \n\n\tIndex Error: {0}\n".format(e))
    except:
        print('\n\tAn error has occurred - error reason', sys.exc_info()[0])
        print('\n\tThe error details 1: {0}'.format(sys.exc_info()[1]))
        print('\n\tThe error details 2:{0} \n'.format(sys.exc_info()[2]))


#...........................
# set the main function
#...........................
if __name__ == "__main__":
    main(sys.argv[0:])
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: MIT-0
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this
 # software and associated documentation files (the "Software"), to deal in the Software
 # without restriction, including without limitation the rights to use, copy, modify,
 # merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 # permit persons to whom the Software is furnished to do so.
 #
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 # INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
 # PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 # HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import boto3
import sys
import json
import time
import botocore
from tqdm import tqdm

def execute_query(clusterid, db, secret, query):
    try:
        client = boto3.client('redshift-data')
        response = client.execute_statement(
            ClusterIdentifier=clusterid,
            Database=db,
            SecretArn=secret,
            Sql=query
        )

        while True:
            response=client.describe_statement(Id=response['Id'])
            if (response["Status"] == "FINISHED" or response["Status"] == "ABORTED" or response["Status"] == "FAILED"):
                #print("Query Status: {0}".format(response["Status"]))
                if(response["Status"] == "FAILED"):
                    print("\nEXECUTION ERROR {0}".format(response["Error"]))
                    raise Exception
                elif(response["Status"] == "ABORTED"):
                    print("\n\tQuery has been aborted\n")
                    raise Exception
                #else:
                #    print("EXECUTION ERROR: None")
                break

            else:
                #print("Query Status: {0}".format(response["Status"]))
                time.sleep(2)

        if (response["HasResultSet"] == True):
            result=client.get_statement_result(
                Id=response['Id']
            )
            return result

    except client.exceptions.ValidationException as Ve:
        sys.exit("Validation Exception: \n\n\t{0}".format(Ve))

    except client.exceptions.ExecuteStatementException as ESe:
        sys.exit("\nExecute Statement Exception: \n\n\t{0}\n".format(ESe))
    
    except client.exceptions.ActiveStatementsExceededException as ASEe:
        sys.exit("\nActive Statements Exceeded Exception: \n\n\t{0}\n".format(ASEe))
    
    except client.exceptions.ResourceNotFoundException as RNFe:
        sys.exit("\n Resource not Found Exception: \n\n\t{0}\n".format(RNFe))

    except client.exceptions.InternalServerException as ISe:
        sys.exit("\nInternal Server Exception: \n\n\t{0}\n".format(ISe))
    
    except client.exceptions.ClientError as CEe:
        sys.exit("\nClient Error Exception: \n\n\t{0}\n".format(CEe))

    except:
        sys.exit("{0}".format(sys.exc_info()[1]))

def create_control_table(clusterid, db, secret):

    createQuery = "CREATE TABLE IF NOT EXISTS public.encoding_control \
    (  \
	schema_name character varying(100) NOT NULL ENCODE lzo, \
	table_name character varying(100) NOT NULL ENCODE lzo,   \
	column_name character varying(100) NOT NULL ENCODE lzo,  \
	curr_encoding_type character varying(20) NOT NULL ENCODE lzo, \
    recmd_encoding_type character varying(20) NOT NULL DEFAULT 'unknown' ENCODE lzo, \
    est_red_pct character varying(20) NOT NULL DEFAULT '0.00' ENCODE lzo, \
    analyze_ind character(1) NOT NULL DEFAULT 'N'::bpchar ENCODE lzo, \
	alter_status character(1) NOT NULL DEFAULT 'N'::bpchar ENCODE lzo,   \
	update_timestamp timestamp without time zone ENCODE az64 \
    )  \
    DISTSTYLE AUTO;"

    for i in tqdm (range (1), desc="Creating Control Table if not exists.."):
        execute_query(clusterid, db, secret, createQuery)

def truncate_control_table(clusterid, db, secret):

    truncateQuery = "truncate public.encoding_control"

    for i in tqdm (range (1), desc="Truncating Control Table:.."):
        execute_query(clusterid, db, secret, truncateQuery)

def load_control_table(clusterid, db, secret, schemaName):

    listTables = "select schemaname \
    ,tablename \
    ,"'"column"'" as columnname \
    ,encoding \
    from pg_table_def \
    where schemaname='"+schemaName+"' \
    and tablename<>'encoding_control'"

    result = execute_query(clusterid, db, secret, listTables)
    if(result["Records"] != []):
        for i in tqdm (range (1), desc="Loading Control Table for Schema {0}..".format(schemaName)):
            for datarecord in result["Records"]:
                loadControlQuery="insert into public.encoding_control \
                values("
                for column in datarecord:
                    loadControlQuery+="'"+column['stringValue']+"',"
                
                loadControlQuery+="default,default,default,default,current_timestamp);"
                execute_query(clusterid, db, secret, loadControlQuery)
    else:
        sys.exit("The schema does not exist or does not have any tables")

def is_table_large(clusterid, db, secret, schema, table, threshold):
    tableSizeQuery="select "'"schema"'" as schemaname \
    ,"'"table"'" as tablename \
    ,size as sizeinMB \
    FROM SVV_TABLE_INFO \
    where schemaname = '"+schema+"' \
    and tablename = '"+table+"';"

    result = execute_query(clusterid, db, secret, tableSizeQuery)
    if(result["Records"] != []):
        tablesize=result["Records"][0][2]['longValue']
        if((tablesize/1024/1024)>int(threshold)):
            return True
    
    return False

def first_sort_key(clusterid, db, secret, schema, table):
    sortKeyQuery="select "'"column"'" \
    FROM pg_table_def \
    where schemaname = '"+schema+"' \
    and tablename = '"+table+"' \
    and sortkey=1;"

    result = execute_query(clusterid, db, secret, sortKeyQuery)
    if(result["Records"] != []):
        return result["Records"][0][0]['stringValue']
    
    return ''



def update_control_table(clusterid, db, secret, schema_name, table_name, column_name, encoding, est_red_pct, analyze_ind, alter_status):
    updateQuery="update public.encoding_control \
    set recmd_encoding_type='"+encoding+"',   \
    est_red_pct='"+est_red_pct+"',   \
    analyze_ind='"+analyze_ind+"',    \
    alter_status='"+alter_status+"',  \
    update_timestamp=current_timestamp   \
    where schema_name='"+schema_name+"' \
    and table_name='"+table_name+"'    \
    and column_name='"+column_name+"';"

    result = execute_query(clusterid, db, secret, updateQuery)
    return result


def analyze_compression(clusterid, db, secret, schema, table):
    analyzeQuery='analyze compression '+schema+'.'+table+';'
    result = execute_query(clusterid, db, secret, analyzeQuery)
    if(result != None):
        return result["Records"]
    
    return result

def alter_compression(clusterid, db, secret, schema, table, column, encoding):

    alterQuery=" alter table "+schema+"."+table+" alter column "+column+" encode "+encoding+";"
    result = execute_query(clusterid, db, secret, alterQuery)
    return result

def get_control_table(clusterid, db, secret, schema_name):
    
    getTableQuery="select distinct table_name \
    from public.encoding_control    \
    where schema_name='"+schema_name+"';"
    
    result = execute_query(clusterid, db, secret, getTableQuery)
    if(result["Records"] != []):
        return result["Records"]
    

def compress_all(clusterid, db, secret,schema_name):
    result=get_control_table(clusterid, db, secret, schema_name)
    if(result != None):
        for i in tqdm (range (1), desc="Applying Compression for all Required Tables in Schema {0}..".format(schema_name)):
            for table in result:
                table_name=table[0]["stringValue"]
                fSortKey=first_sort_key(clusterid, db, secret, schema_name,table_name)
                analyze_result=analyze_compression(clusterid, db, secret, schema_name, table_name)
                if(analyze_result != None):
                    for encoding in analyze_result:
                        if(float(encoding[3]["stringValue"])>0 and encoding[1]["stringValue"]!=fSortKey):
                            alter_compression(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"])
                            update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'Y')
                        else:
                            update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'N')

    
def compress_large(clusterid, db, secret,schema_name,threshold):
    result=get_control_table(clusterid, db, secret, schema_name)
    if(result != None):
        for i in tqdm (range (1), desc="Applying Compression for Required Tables > {0} TB in Schema {1}..".format(threshold,schema_name)):
            for table in result:
                table_name=table[0]["stringValue"]
                fSortKey=first_sort_key(clusterid, db, secret, schema_name,table_name)
                if(is_table_large(clusterid, db, secret, schema_name, table_name, threshold)):
                    analyze_result=analyze_compression(clusterid, db, secret, schema_name, table_name)
                    if(analyze_result != None):
                        for encoding in analyze_result:
                            if(float(encoding[3]["stringValue"])>0 and encoding[1]["stringValue"]!=fSortKey):
                                alter_compression(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"])
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'Y')
                            else:
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'N')  

def compress_small(clusterid, db, secret,schema_name, threshold):
    result=get_control_table(clusterid, db, secret, schema_name)
    if(result != None):
        for i in tqdm (range (1), desc="Applying Compression for Required Tables < {0} TB in Schema {1}..".format(threshold,schema_name)):
            for table in result:
                table_name=table[0]["stringValue"]
                fSortKey=first_sort_key(clusterid, db, secret, schema_name,table_name)
                if(not is_table_large(clusterid, db, secret, schema_name, table_name, threshold)):
                    analyze_result=analyze_compression(clusterid, db, secret, schema_name, table_name)
                    if(analyze_result != None):
                        for encoding in analyze_result:
                            if(float(encoding[3]["stringValue"])>0 and encoding[1]["stringValue"]!=fSortKey):
                                alter_compression(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"])
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'Y')
                            else:
                                update_control_table(clusterid, db, secret, schema_name, table_name, encoding[1]["stringValue"], encoding[2]["stringValue"], encoding[3]["stringValue"], 'Y', 'N')

        
        
    
    

def main(argv):
    
    try:
        if(len(sys.argv) == 1):
            print(
                '\nUsage: Alter-Compression-Utility.py   \
                \n\tThe utility can automatically analyze schema, identify tables and alter column encoding compression leveraging Alter Compression feature of Redshift released in Oct, 20   \
                \n\t\tReference Document: https://aws.amazon.com/about-aws/whats-new/2020/10/amazon-redshift-supports-modifying-column-comprression-encodings-to-optimize-storage-utilization-query-performance/ \
                \n\nArguments:\n \
                \n\targv[1] - Cluster Identifier: Name of the Cluster \
                \n\targv[2] - Database Name: Default database name that is required for connecting to cluster \
                \n\targv[3] - Secret Identifier: Secret ARN/Friendly name of the secret for retrieving user and password of the cluster \
                \n\targv[4] - Schema Name: Name of the schema that user wants to analyze and compress tables for  \
                \n\targv[5] - Compression Mode: Three modes of compression for users to choose from and pass as argument \
                \n\n\t\tcompress-all: Automatically analyze and compress required columns for all the tables for the schema    \
                \n\t\tcompress-small: Based on provided threshold value, Automatically analyze and compress required columns for tables < threshold value in TB for the schema   \
                \n\t\tcompress-large: Based on provided threshold value,  Automatically analyze and compress required columns for tables > threshold value in TB for the schema  \
                \n\n\targv[6] - Threshold value: Threshold value upon which small/large table size in TBs will be defined.The value is not required for compress-all mode.\n'
            )

        elif(len(sys.argv) > 7):
            print('\n\tPassed arguments have exceeded the limit. Please refer to the utility documentation\n')

        else:
            print('\n\n')
            if (argv[5]=='compress-all'):
                create_control_table(argv[1], argv[2], argv[3])
                truncate_control_table(argv[1], argv[2], argv[3])
                load_control_table(argv[1], argv[2], argv[3], argv[4])
                compress_all(argv[1], argv[2], argv[3],argv[4])
                print("\nComplete: \n\tCompression has successfully been altered for all required tables in schema {0} and public.encoding_control has been updated for detail. \n".format(argv[4]))

            elif(int(argv[6])>0 and (argv[5]=='compress-small' or argv[5]=='compress-large')):

                create_control_table(argv[1], argv[2], argv[3])
                truncate_control_table(argv[1], argv[2], argv[3])
                load_control_table(argv[1], argv[2], argv[3], argv[4])
                if(argv[5]=='compress-small'):
                    compress_small(argv[1], argv[2], argv[3],argv[4],argv[6])
                    print("\nComplete: \n\tCompression has successfully been altered for all required tables < {0} TB in schema {1} and public.encoding_control has been updated for detail. \n".format(argv[6], argv[4]))

                elif(argv[5]=='compress-large'):
                    compress_large(argv[1], argv[2], argv[3],argv[4],argv[6])
                    print("\nComplete: \n\tCompression has successfully been altered for all required tables > {0} TB in schema {1} \
                    and public.encoding_control has been updated for detail. \n".format(argv[6], argv[4]))

            else:
                if (argv[5]!='compress-all' and int(argv[6])<=0):
                    print("\nPlease pass a table size (TB) threshold value greater than zero. Refer to the utility documentation !\n")
                else:
                    print("\nPlease pass a valid compression mode. Refer to the utility documentation !\n")

    except IndexError as e:
        print("\nPlease pass valid command line arguments. Refer to the utility documentation ! \
        \n\n\tIndex Error: {0}\n".format(e))
    except:
        print('\n\tAn error has occurred - error reason', sys.exc_info()[0])
        print('\n\tThe error details 1: {0}'.format(sys.exc_info()[1]))
        print('\n\tThe error details 2:{0} \n'.format(sys.exc_info()[2]))


#...........................
# set the main function
#...........................
if __name__ == "__main__":
    main(sys.argv[0:])
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: MIT-0 # # Permission is hereby granted, free of charge, to any person obtaining a copy of this # software and associated documentation files (the "Software"), to deal in the Software # without restriction, including without limitation the rights to use, copy, modify, # merge, publish, distribute, sublicense, and/or sell copies of the Software, and to # permit persons to whom the Software is furnished to do so. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A # PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT # HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. #print("Query Status: {0}".format(response["Status"])) #else: #    print("EXECUTION ERROR: None") #print("Query Status: {0}".format(response["Status"])) #........................... # set the main function #...........................