Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import logging 

2from abc import ABC, abstractmethod 

3 

4# Try to import in Airflow environemnt for get_credentials() use 

5try: 

6 from airflow.models import Variable 

7except: 

8 pass 

9# Try to import in Glue enivornment for get_credentials() use 

10try: 

11 from sdc_etl_libs.aws_helpers.aws_helpers import AWSHelpers 

12 from sdc_etl_libs.sdc_dataframe.Dataframe import Dataframe 

13 from sdc_etl_libs.sdc_dataframe.SDCDataframeEnums import SDCDFTypes 

14except Exception as e: 

15 raise Exception(e) 

16 

17 

18class API(object): 

19 base_url = None 

20 credentials = None 

21 schema_name = None 

22 

23 def __init__(self): 

24 logging.info("Base constructor.") 

25 raise Exception("Base constructor.") 

26 

27 def get_api_response(self): 

28 """ 

29 Gets API Response. 

30 :return: SDC Dataframe or a flattened list of dictionaries. 

31 """ 

32 data = None 

33 try: 

34 data = self.get_response_data() 

35 except Exception as e: 

36 logging.error(f"Failed getting response data. {e}") 

37 return data 

38 

39 def get_credentials(self, source_, aws_secret_id_=None, aws_secret_region_='us-east-2', airflow_vars_=None): 

40 """ 

41 @param source_: Where to grab credentials from (Secrets Manager, Airflow, etc.) 

42 if source_ is 'aws_secrets', pass the following arguments: 

43 @param aws_secret_id_: AWS Secret Name 

44 @param aws_secret_region_: AWS region. Defaults to us-east-2. 

45 if source_ is 'airflow, pass the following argruments: 

46 @param airflow_vars_: List of Airflow variables 

47 @return: Dictionary of credentials 

48 """ 

49 # This method has been deprecated. Please use the SDCEndpointCredentials  

50 # class in the sdc_endpoint_credentials modules instead. 

51 if source_ == 'aws_secrets': 

52 if aws_secret_id_: 

53 try: 

54 creds = AWSHelpers.get_secrets(secret_id_=aws_secret_id_, region_=aws_secret_region_) 

55 

56 except Exception as e: 

57 logging.error(str(e)) 

58 raise Exception("Failed to get credentials.") 

59 

60 return creds 

61 else: 

62 raise Exception("aws_secret_id_ not provided.") 

63 

64 elif source_ == 'airflow': 

65 if airflow_vars_: 

66 creds = {} 

67 if isinstance(airflow_vars_, list): 

68 for var in airflow_vars_: 

69 try: 

70 creds.update({var: Variable.get(var)}) 

71 except Exception as e: 

72 logging.exception(f"Error fetching variable '{var}'") 

73 return creds 

74 else: 

75 raise Exception("Airflow variable list must be passed as a list.") 

76 else: 

77 raise Exception("No Airflow variable list passed.") 

78 

79 else: 

80 raise Exception(f"'{source_}' is an invalid credential source type.") 

81 

82 def get_google_service_credentials(self, service_name_, version_, certificate_, scopes_=None): 

83 """ 

84 Generates the proper credentials for Google service account based on provided certificate. 

85 :param service_name_: String. Name of the service account service (ex. 'doubleclicksearch'). List of 

86 services can be found here: https://developers.google.com/apis-explorer. 

87 :param version_: String. Version of the service account service (ex. 'v2'). 

88 :param certificate_: Dict. Dictionary containing credentials/certificate information needed by 

89 service account. 

90 :param scopes_: List. Credentials/certificates set of scopes. 

91 :return: Google service handle. 

92 """ 

93 

94 from google.oauth2 import service_account 

95 import googleapiclient.discovery 

96 credentials = service_account.Credentials.from_service_account_info(certificate_, scopes=scopes_) 

97 service = googleapiclient.discovery.build( 

98 serviceName=service_name_, version=version_, credentials=credentials, cache_discovery=False) 

99 return service 

100 

101 def get_access_token(self): 

102 

103 logging.error("get_access_token() from API base class not currently setup.") 

104 

105 def get_response_data(self): 

106 """ 

107 Retrieve data from the API. 

108 

109 Note: Overriding this method is optional, but is mandatory for a _source_ API. 

110 """ 

111 raise NotImplementedError(f"{type(self).__name__} Doesn't support getting response data, override this method if you want to use it as a source.") 

112 

113 def write_data(self, data_): 

114 """ 

115 Send/write data to the API. 

116 

117 Note: Overriding this method is optional, but is mandatory for a _sink_ API. 

118 """ 

119 raise NotImplementedError(f"{type(self).__name__} Doesn't support writing data, override this method if you want to use it as a sink.")