1#!/usr/bin/env python3 

2""" 

3############################################################################## 

4## gjb 

5## 06/23/20 

6## https://smiledirectclub.atlassian.net/browse/DE-691 

7## odbc Database class 

8## Connects To odbc databases and provides basic select functionality 

9## 

10# -------------------------------------------------------------------------- 

11############################################################################## 

12""" 

13import logging 

14import pyodbc 

15 

16from sdc_etl_libs.database_helpers.Database import Database, SDCDBTypes 

17from sdc_etl_libs.sdc_credentials.sdc_endpoint_credentials import SDCEndpointCredentials 

18 

19try: 

20 from airflow import AirflowException 

21except: 

22 pass 

23 

24 

25class ODBCDatabase(Database): 

26 # Set logger up 

27 logging.basicConfig(format='%(levelname)s: %(asctime)s: ' '%(funcName)s: %(message)s') 

28 logger = logging.getLogger(__name__) 

29 logger.setLevel(logging.INFO) 

30 

31 def __init__(self, schema_, endpoint_schema_, **kwargs): 

32 self.connection = None 

33 self.cursor = None 

34 self.airflow = None 

35 self.engine = None 

36 self.query_results = None 

37 self.schema = schema_ 

38 self.endpoint_schema = endpoint_schema_ 

39 self.sqlalchemy = kwargs.get("sqlalchemy_", False) 

40 

41 # data schema keys 

42 database_type = self.endpoint_schema.get("info", {}).get("access", {}).get("database_type") 

43 self.type = SDCDBTypes[database_type.upper()] 

44 self.driver = self.endpoint_schema.get("info", {}).get("access", {}).get("driver") 

45 self.host = self.endpoint_schema.get("info", {}).get("access", {}).get("host") 

46 self.port = self.endpoint_schema.get("info", {}).get("access", {}).get("port") 

47 self.dsn = self.endpoint_schema.get("info", {}).get("access", {}).get("dsn") 

48 self.database = self.endpoint_schema.get("info", {}).get("access", {}).get("database_name") 

49 

50 # credentials section 

51 self.credentials = SDCEndpointCredentials.get_credentials(endpoint_schema_["info"]["access"]["credentials"]) 

52 self.user = self.credentials.get("username") 

53 self.password = self.credentials.get("password") 

54 

55 logging.info("ODBC SERVER constructor.") 

56 

57 def connect(self, airflow_=False): 

58 """ 

59 Connects to ODBC database. 

60 :airflow_ Boolean to indicate running from airflow or not 

61 :return: Nothing 

62 """ 

63 self.airflow = airflow_ 

64 logging.info("%s", "Connecting to ODBC Server") 

65 

66 if self.sqlalchemy: 

67 logging.exception("sqlalchemy is not yet implemented for this class") 

68 elif self.dsn: 

69 ## DSN NOT YET TESTED.. But will be need for nexus(cadigital) refactor 

70 logging.info("%s", "Connecting to ODBC Server via DSN") 

71 self.connection = \ 

72 pyodbc.connect("""DSN={dsn}; DRIVER ={driver}""" 

73 .format(dsn=self.dsn, driver=self.driver)) 

74 else: 

75 logging.info("%s", "Connecting to ODBC Server") 

76 self.connection = \ 

77 pyodbc.connect("""DRIVER={driver}; 

78 SERVER={server}; 

79 PORT={port}; 

80 DATABASE={database}; 

81 UID={uid}; 

82 PWD={pwd};""".format(driver=self.driver, server=self.host, port=self.port, 

83 database=self.database, uid=self.user, 

84 pwd=self.password) 

85 ) 

86 

87 try: 

88 if self.sqlalchemy: 

89 self.cursor = self.connection 

90 else: 

91 self.cursor = self.connection.cursor() 

92 logging.info("Connected to ODBC Server.") 

93 

94 except Exception as e: 

95 if self.airflow: 

96 # For Airflow, forces task to fail and sets it up for re-try 

97 raise AirflowException("Error connecting to ODBC. {}".format(str(e))) 

98 else: 

99 logging.exception("Error connecting to ODBC.") 

100 raise Exception 

101 

102 def execute_query(self, query_, return_results_=False): 

103 """ 

104 Executes query against ODBC Database.. 

105 

106 :param query_: Query string to run in ODBC Database. 

107 :param return_results_: Boolean. Defaults to False. If True, 

108 returns results of query to self.query_results 

109 :return: If return_results is True, returns the results of a query 

110 to self.query_results as a tuple or tuple-like proxy object, 

111 which supports index access. 

112 """ 

113 

114 try: 

115 if return_results_: 

116 logging.info("Fetching results...") 

117 self.query_results = self.cursor.execute(query_) 

118 else: 

119 self.cursor.execute(query_) 

120 logging.info("Query ran successfully.") 

121 

122 except Exception as e: 

123 if self.airflow: 

124 # For Airflow, forces task to fail and set it up for re-try 

125 raise AirflowException("Error connecting to ODBC. {}".format(str(e))) 

126 else: 

127 logging.exception("Error running query.") 

128 raise e 

129 

130 def get_results(self, fetch_method_='all', fetch_amount_=None): 

131 """ 

132 Returns results from ODBC Database cursor object at self.query_results. 

133 

134 When running execute_query() with return_results_=True, a 

135 cursor/proxy object is returned at self.query_results_ which 

136 can be used to return records. When the records are exhausted, 

137 the cursor/proxy object will return 0. 

138 

139 :param fetch_method_: Method by which to fetch results from 

140 the cursor/proxy object and place into memory. 

141 Can be one of the following options: 

142 'all' - Returns all records at once 

143 'one' - Returns a single record 

144 'many' - Returns a specified number of records at a time 

145 :param fetch_amount_: Number of rows to return at a time. If 'many' 

146 is selected for the fetch_method_, must provide a fetch_amount_ 

147 as a positive integer. Does not work for 'all' or 'one'. 

148 :yields: Tuples or tuple-like objects for each record 

149 in the iteration. 

150 

151 :returns: This function will return a generator object that can be 

152 looped through to get results. 

153 

154 # Example of simple iteration over rows 

155 while cnx.query_results: 

156 for row in cnx.get_results('many', 10): 

157 print(row) 

158 

159 """ 

160 

161 if not self.query_results: 

162 raise Exception("Cannot get results. Missing cusor/proxy object " 

163 "with query. Use execute_query() where " 

164 "return_results_=True.") 

165 

166 if fetch_method_.lower() == 'all': 

167 results = self.query_results.fetchall() 

168 for result in results: 

169 yield result 

170 self.query_results = None 

171 

172 elif fetch_method_.lower() == 'one': 

173 yield self.query_results.fetchone() 

174 self.query_results = None 

175 

176 elif fetch_method_.lower() == 'many': 

177 if not isinstance(fetch_amount_, int) or fetch_amount_ < 0: 

178 raise Exception("Query results fetch type is set to 'many', " 

179 "however, a positive integer was not passed " 

180 "for fetch amount.") 

181 

182 results = self.query_results.fetchmany(fetch_amount_) 

183 if len(results) > 0: 

184 for row in results: 

185 yield row 

186 else: 

187 self.query_results = None 

188 

189 else: 

190 raise Exception("{} if not a valid fetch method type.".format(fetch_method_)) 

191 

192 def close_connection(self): 

193 """ 

194 Closes database connection. 

195 """ 

196 

197 try: 

198 self.connection.close() 

199 logging.info("Connection closed.") 

200 

201 except Exception: 

202 # We want to try and close connection, but, don't want to trigger 

203 # a fail state/re-try in Airflow if something goes wrong here. 

204 if self.airflow: 

205 logging.exception("Class variable airflow is set to true. So no Exception action will be taken here") 

206 else: 

207 logging.exception("Did not properly close connection.") 

208 raise Exception 

209 

210 # Will never insert into ODBC Database.. i dont think 

211 def insert_data(self, data_list_, upsert=False): 

212 """ 

213 Inserts data into ODBC SERVER database. 

214 :param data_list_: 

215 :param upsert: 

216 """ 

217 raise NotImplementedError