from io import StringIO, BytesIO import psycopg2 import sqlalchemy import pandas.io.sql as sql import pandas as pd import boto3 import os import operator from numpy import ndarray from psycopg2.extras import execute_values from internal.helpers import get_rds_connection import logging component_logger = logging.getLogger().getChild("utils.aws_connectors") if os.getenv('PROFILE') == 'local': boto3.setup_default_session(profile_name='fansifter') def pd_read_sql(query_sql, **kwargs) -> pd.DataFrame: attempts = 10 # workaround for sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) SSL connection has been closed # unexpectedly while attempts > 0: try: attempts -= 1 df = pd.read_sql(sql=query_sql, con=get_rds_connection(), **kwargs) return df except sqlalchemy.exc.OperationalError as e: component_logger.exception(e) except psycopg2.OperationalError as e: component_logger.exception(e) except Exception as e: component_logger.exception(e) raise e def run_query(query_sql, *args, return_type=None, update_values=False, fetch=True, **kwargs): """ Runs query against AWS RDS PostgresSQL instance (version >10). Creates connection, runs query, ommmits, terminates connection. If there are rows, returns them. Avoid creating connection for each individual INSERT, or refactor this function. Need to explicitly grant all privileges on any new schemas and tables when not superuser. """ try: with get_rds_connection() if return_type == 'dict' else get_rds_connection(cursor_factory=None) as connection: with connection.cursor() as cursor: # Using parameterized query to avoid sql injection and other similar issues # E.g. it fixes this, which before caused syntax error: "text": "Economy of the People\'s <-- the apostrophe # If there are args, it comes from insert and we insert a tuple. This seems easy to break, so add test cases # print(cursor.mogrify(query_sql, (args,))) # for when we want to see the sql records = None if update_values and args and isinstance(args[0], list): # TODO: probably works smoother with pandas? execute_values(cursor, query_sql, args[0]) else: if return_type is None or return_type == 'dict': if args and isinstance(args[0], dict): cursor.execute(query_sql, args[0]) elif args and isinstance(args[0], (list, ndarray)): #cursor.executemany(query_sql, args[0]) records = execute_values(cursor, query_sql, args[0], page_size=len(args[0]), fetch=fetch) return_type = "toomany" elif isinstance(query_sql, list): for ql in query_sql: cursor.execute(ql, *args, **kwargs) else: cursor.execute(query_sql, *args, **kwargs) try: # See if we want to return dataframe instead of list. Consider only using dataframes. if return_type == 'df': if args and isinstance(args[0], dict): records = sql.read_sql(query_sql, connection, params=args[0]) else: records = sql.read_sql(query_sql, connection) elif return_type == 'toomany': pass elif fetch: records = cursor.fetchall() except psycopg2.ProgrammingError as e: pass except Exception as e: # TODO! this needs to be handled more gracefully (instead pass) # We want don't want to fetch when we insert component_logger.exception(f"FAILED run_query {query_sql}", exc_info=e) return records except sqlalchemy.exc.OperationalError as e: component_logger.exception(e) except psycopg2.OperationalError as e: component_logger.exception(e) except Exception as e: component_logger.exception(e) raise e def s3_to_pandas(bucket, key, header=None, file_format='csv'): """ Downloads S3 file and returns it in a pandas DataFrame. If you want to get header from file, pass header="infer". """ component_logger.info(f'Attempting to read {key} file from {bucket} bucket') client = boto3.client('s3') temp_object = client.get_object(Bucket=bucket, Key=key)['Body'].read(100) obj = client.get_object(Bucket=bucket, Key=key)['Body'] """ Find a deliminter. Although csv-sniffer in C and Python should do that automatically, they can fail sometimes, as was evident with "devel---useruploadoriginals/2c1a9dd1e9f518e1d4fee2f2a14917722b80da4b029e084fa64ceaeab322beb3/testcase.csv" We see the occuranes of ; , \t and decide what should be delimiter. We do this on temp_object, else we already start reading the obj stream. """ chars = str(temp_object) char_stats = {';': chars.count(";"), ",": chars.count(","), "\t": chars.count("\t"), } delimiter = max(char_stats.items(), key=operator.itemgetter(1))[0] # TODO! handle encoding guessing # import chardet # result = chardet.detect(temp_object) # charenc = result['encoding'] # component_logger.info(f' ENCODING = {charenc}') if not delimiter: raise Exception("We could not determine supported delimiters within the first 100 chars from file") if file_format == 'csv': return pd.read_csv(obj, sep=delimiter, header=header, dtype=str) elif file_format == 'xlsx' or file_format == 'xls': return pd.read_excel(BytesIO(obj.read())) # TODO! Add Error Handling for unsupported file type return False def pandas_to_s3(bucket, df, key, index=False): """ Creates a CSV from pandas DataFrame and uploads it to S3""" csv_buffer = StringIO() df.to_csv(csv_buffer, index=index) s3_resource = boto3.resource('s3') s3_resource.Object(bucket, key).put(Body=csv_buffer.getvalue(), ServerSideEncryption='AES256') def df_to_db(df: pd.DataFrame, schema: str, table_name: str, if_exists='append', index_label=None): """ Using df.to_sql to insert the df in to table :param df: source_df :param schema: schema name :param table_name: table name :param if_exists: if exists parameter default: 'append' :param index_label: index label :return: """ df.to_sql(name=table_name, con=get_rds_connection(), schema=schema, if_exists=if_exists, chunksize=16000, index=False, index_label=index_label ) return if __name__ == '__main__': """ # S3 test bucket = 'fansifter-model-data' key = 'data.csv' header = 'infer' df = s3_to_pandas(bucket, key, header) print(df.head())"""