1  from typing import  Dict, List
   2  import boto3
   3  from config import configManager
   4  
   5  class DataExtractor:
   6      def __init__(self, spark_session, source):
   7          self.spark_session = spark_session
   8          self.options= configManager().get_details(source)
   9          self.creds= self.options['creds']
  10          self.source = source
  11          
  12      def get_transformed(self,df):
  13          try:
  14              logic=configManager().get_config(f"{self.source.split('_')[0]}_transformation")
  15              logic = ','.join([f" {i[1]} as {i[0]}" for i in logic.items()])
  16              df.createOrReplaceTempView(f'__df__{self.source}')
  17              return  self.spark_session.sql(f"select {logic} from __df__{self.source}")
  18          except Exception as e:
  19              raise e
  20              #return logic
  21              
  22      def push_to_prod_d(df, source=None):
  23          ####
  24          """DataFrame
  25          Customer Data pushed to PROD D
  26          """
  27      
  28      def extract_data(self,columns :str =None, pushdown_predicate : str =None, limit :int=None ):
  29          # Connect to Snowflake using the AWS Snowflake Connector
  30          # Extract data from Snowflake using Glue
  31          """
  32          source eg: petco-snowflake, retailscape-aurora, glue etc
  33          columns : ','.join(list_of_columns)
  34          pushdown_predicate : source ='xyz' and cast(date as varchar) = '2023-01-08'
  35          """
  36          source = self.source
  37          columns = columns if columns else '*'
  38          pushdown_predicate= 'where '+pushdown_predicate if pushdown_predicate else ''
  39          limit = f'limit {str(limit)}' if limit else ''
  40          
  41          if 'glue' == source:
  42              
  43              query = f"SELECT {columns} FROM {self.options['database']}.{self.options['table']}  {pushdown_predicate} {limit}".strip()
  44              
  45              df= self.spark_session.sql(query)
  46              
  47          elif 'snowflake' in source:
  48              
  49              sfoptions = {
  50                            "sfURL" : self.creds['account']+'.snowflakecomputing.com',
  51                            "sfUser" : self.creds['user'],
  52                            "sfPassword" : self.creds['password'],
  53                            "sfDatabase" : self.options['database'],
  54                            "sfSchema" : self.options['schema'],
  55                            "sfWarehouse" : self.options['warehouse']
  56                          }
  57                          
  58              
  59              
  60              query = f"SELECT {columns} FROM {self.options['schema']}.{self.options['table']} {pushdown_predicate} {limit}".strip()
  61              
  62              df = self.spark_session.read.format('net.snowflake.spark.snowflake') \
  63                    .options(**sfoptions) \
  64                    .option("query",  query) \
  65                    .load()
  66          elif 'aurora' in source:
  67              
  68              table = self.options['table']
  69              database= self.options['database']
  70              try:
  71                  host=self.creds['host']
  72              except:
  73                  host=self.creds['base_url']
  74              try:
  75                  port=self.creds['port']
  76              except:
  77                  port='5432'
  78              try:
  79                  user=self.creds['user']
  80              except:
  81                  user='postgres'
  82              try:    
  83                  password = self.creds['password']
  84              except:
  85                  raise ValueError(f"Password Not Defined in secret")
  86                  
  87              url= f'jdbc:postgresql://{host}:{port}/{database}'
  88         
  89              query = f"SELECT {columns} FROM {self.options['table']} {pushdown_predicate} {limit}".strip()
  90              
  91              
  92              df = self.spark_session.read.format("jdbc").option("url", url)\
  93                  .option("user", user).option("password", password)\
  94                      .option("query", query).load()
  95                  
  96              #df = self.spark_session.read.jdbc(url=AURORA_DATABASE_URL, table=query, properties=properties)
  97              
  98          return df