import boto3


class DataExtractor:
    def __init__(self, spark_session,cred_manager):
        self.spark_session = spark_session
        self.cred_manager = cred_manager 
        
    
    def extract_data(self,source,  options, pushdown_predicate=None, columns=None):
        # Connect to Snowflake using the AWS Snowflake Connector
        # Extract data from Snowflake using Glue
        
        columns = columns if columns else '*'
        pushdown_predicate= 'where '+pushdown_predicate if pushdown_predicate else ''
        
        if 'glue' == source:
            
            query = f"SELECT {columns} FROM {options['database']}.{options['table']} {pushdown_predicate}".strip()
            
            df= self.spark_session.sql(query)
            
        elif 'snowflake' in source:
            creds = self.cred_manager.get_credentials(source)
            sfOptions = {
                          "sfURL" : creds['account']+'.snowflakecomputing.com',
                          "sfUser" : creds['user'],
                          "sfPassword" : creds['password'],
                          "sfDatabase" : options['database'],
                          "sfSchema" : options['schema'],
                          "sfWarehouse" : options['warehouse']
                        }
                        
            
            
            query = f"SELECT {columns} FROM {options['schema']}.{options['table']} {pushdown_predicate}".strip()
            
            df = self.spark_session.read.format('net.snowflake.spark.snowflake') \
                  .options(**sfOptions) \
                  .option("query",  query) \
                  .load()
        elif 'aurora' in source:
            creds = self.cred_manager.get_credentials(source)
            table = options['table']
            database= options['database']
            try:
                host=creds['host']
            except:
                host=creds['base_url']
            try:
                port=creds['port']
            except:
                port='5432'
            try:
                user=creds['user']
            except:
                user='postgres'
            try:    
                password = creds['password']
            except:
                raise ValueError(f"Password Not Defined in secret")
                
            url= f'jdbc:postgresql://{host}:{port}/{database}'
       
            query = f"SELECT {columns} FROM {options['table']} {pushdown_predicate}".strip()
            
            
            df = self.spark_session.read.format("jdbc").option("url", url)\
                .option("user", user).option("password", password)\
                    .option("query", query).load()
                
            #df = self.spark_session.read.jdbc(url=AURORA_DATABASE_URL, table=query, properties=properties)
            
        return df