1  import boto3
   2  
   3  
   4  class DataExtractor:
   5      def __init__(self, spark_session,cred_manager):
   6          self.spark_session = spark_session
   7          self.cred_manager = cred_manager 
   8          
   9      
  10      def extract_data(self,source,  options, pushdown_predicate=None, columns=None):
  11          # Connect to Snowflake using the AWS Snowflake Connector
  12          # Extract data from Snowflake using Glue
  13          
  14          columns = columns if columns else '*'
  15          pushdown_predicate= 'where '+pushdown_predicate if pushdown_predicate else ''
  16          
  17          if 'glue' == source:
  18              
  19              query = f"SELECT {columns} FROM {options['database']}.{options['table']} {pushdown_predicate}".strip()
  20              
  21              df= self.spark_session.sql(query)
  22              
  23          elif 'snowflake' in source:
  24              creds = self.cred_manager.get_credentials(source)
  25              sfOptions = {
  26                            "sfURL" : creds['account']+'.snowflakecomputing.com',
  27                            "sfUser" : creds['user'],
  28                            "sfPassword" : creds['password'],
  29                            "sfDatabase" : options['database'],
  30                            "sfSchema" : options['schema'],
  31                            "sfWarehouse" : options['warehouse']
  32                          }
  33                          
  34              
  35              
  36              query = f"SELECT {columns} FROM {options['schema']}.{options['table']} {pushdown_predicate}".strip()
  37              
  38              df = self.spark_session.read.format('net.snowflake.spark.snowflake') \
  39                    .options(**sfOptions) \
  40                    .option("query",  query) \
  41                    .load()
  42          elif 'aurora' in source:
  43              creds = self.cred_manager.get_credentials(source)
  44              table = options['table']
  45              database= options['database']
  46              try:
  47                  host=creds['host']
  48              except:
  49                  host=creds['base_url']
  50              try:
  51                  port=creds['port']
  52              except:
  53                  port='5432'
  54              try:
  55                  user=creds['user']
  56              except:
  57                  user='postgres'
  58              try:    
  59                  password = creds['password']
  60              except:
  61                  raise ValueError(f"Password Not Defined in secret")
  62                  
  63              url= f'jdbc:postgresql://{host}:{port}/{database}'
  64         
  65              query = f"SELECT {columns} FROM {options['table']} {pushdown_predicate}".strip()
  66              
  67              
  68              df = self.spark_session.read.format("jdbc").option("url", url)\
  69                  .option("user", user).option("password", password)\
  70                      .option("query", query).load()
  71                  
  72              #df = self.spark_session.read.jdbc(url=AURORA_DATABASE_URL, table=query, properties=properties)
  73              
  74          return df