1  import csv
   2  import io
   3  from datetime import datetime
   4  from typing import List
   5  
   6  import boto3
   7  from botocore.client import Config
   8  from configmanager import configManager
   9  
  10  
  11  
  12  
  13  config =Config(connect_timeout=9999)
  14  
  15  
  16  class read_s3:
  17      def __intit__(self):
  18          self.configmng= configManager()
  19          self.s3_resoyurce= boto3.resource('s3', aws_access_key_id='',
  20                                      aws_secret_access_key='',
  21                                      region_name='us-east-1', config=config)
  22  
  23      def rename_file_on_s3(full_input_file_path, target_file_path, 
  24                              target_bucket_same=True, target_bucket_name : str=None):
  25                        
  26                              
  27          bucket_name = target_bucket_name if target_bucket_name else full_input_file_path.split('//')[1].split('/', 1)[0]
  28          file_path = full_input_file_path.split('//')[1].split('/', 1)[1]
  29  
  30          original_obj = s3_resource_bt.Object(bucket_name, file_path)
  31  
  32          target_obj = s3_resource_bt.Object(target_bucket_name, target_file_path)
  33  
  34          target_obj.put(Body=original_obj.get()['Body'].read())
  35          print("renamed file to ", target_file_path)
  36  
  37  
  38      def delete_file_on_s3(full_input_file_path):
  39          bucket_name = full_input_file_path.split('//')[1].split('/', 1)[0]
  40          file_path = full_input_file_path.split('//')[1].split('/', 1)[1]
  41          self.s3_resource.Object(bucket_name, file_path).delete()
  42          print('deleted temp file', file_path)
  43  
  44  
  45  
  46      def split_s3_full_path(s3_path):
  47          path_parts = s3_path.replace("s3://", "").split("/")
  48          bucket = path_parts.pop(0)
  49          key = "/".join(path_parts)
  50          return bucket, key
  51  
  52  
  53      def list_s3_files(bucket_name, bucket_path, file_name_wildcard):
  54          file_name_wildcard = file_name_wildcard.split('*')[0]
  55          file_name_wildcard = file_name_wildcard.split('/')[-1]
  56  
  57          
  58          bucket = self.s3_resource.Bucket(bucket_name)
  59          list_of_files = []
  60          for object_summary in bucket.objects.filter(Prefix=bucket_path):
  61              if file_name_wildcard in object_summary.key:
  62                  list_of_files.append(object_summary.key)
  63          return list_of_files
  64  
  65      def update_tables_with_new_partitions(database , table, partition_values, s3_path):
  66          import boto3
  67          client = boto3.client("athena", aws_access_key_id="",
  68                              aws_secret_access_key="", region_name="us-east-1")
  69  
  70          
  71          try:
  72              update_query_response = client.start_query_execution(
  73                  QueryString="ALTER TABLE {} ADD PARTITION ({});".format(table_name, partition_values),
  74                  QueryExecutionContext={
  75                      'Database': database
  76                  },
  77                  ResultConfiguration={
  78                      'OutputLocation': s3_path
  79                  }
  80              )
  81          except Exception as e:
  82              raise e
  83          print("updated partitions for table {} : {}!".format(table_name, full_date))
  84  
  85      def run_athena_query(query, db, output_path):
  86          each_query_response = athena_client_bt.start_query_execution(
  87              QueryString=query,
  88              QueryExecutionContext={
  89                  'Database': db
  90              },
  91              ResultConfiguration={
  92                  'OutputLocation': output_path,
  93              }
  94          )
  95          query_id = each_query_response['QueryExecutionId']
  96          i = 0
  97          while True:
  98              i += 1
  99  
 100              if i > 10:
 101                  time.sleep(10)
 102              else:
 103                  time.sleep(2)
 104              each_query_execution_status = athena_client_bt.get_query_execution(
 105                  QueryExecutionId=query_id)
 106  
 107              if each_query_execution_status['QueryExecution']['Status']['State'] == 'SUCCEEDED':
 108                  print("SUCCEEDED")
 109                  output_file_name = each_query_execution_status['QueryExecution']['ResultConfiguration'][
 110                      'OutputLocation']
 111                  return output_file_name
 112              elif each_query_execution_status['QueryExecution']['Status']['State'] in ('FAILED', 'CANCELLED'):
 113                  print("FAILED/CANCELLED", query_id)
 114                  raise Exception("FAILED/CANCELLED{}".format(query_id))
 115              else:
 116                  continue
 117