Coverage for mindsdb / integrations / handlers / derby_handler / derby_handler.py: 0%

1from typing import Optional 

2from mindsdb_sql_parser.ast.base import ASTNode 

3from mindsdb.integrations.libs.base import DatabaseHandler 

4from mindsdb.utilities import log 

5from mindsdb_sql_parser import parse_sql 

6from mindsdb.integrations.libs.response import ( 

7 HandlerStatusResponse as StatusResponse, 

8 HandlerResponse as Response, 

9 RESPONSE_TYPE, 

10) 

11import pandas as pd 

12import jaydebeapi as jdbcconnector 

13 

14logger = log.getLogger(__name__) 

15 

16 

17class DerbyHandler(DatabaseHandler): 

18 name = "derby" 

19 

20 def __init__(self, name: str, connection_data: Optional[dict], **kwargs): 

21 """Initialize the handler 

22 Args: 

23 name (str): name of particular handler instance 

24 connection_data (dict): parameters for connecting to the database 

25 **kwargs: arbitrary keyword arguments. 

26 """ 

27 super().__init__(name) 

28 

29 self.kwargs = kwargs 

30 self.parser = parse_sql 

31 self.database = connection_data["database"] 

32 self.connection_config = connection_data 

33 self.host = connection_data["host"] 

34 self.port = connection_data["port"] 

35 self.schema = "APP" 

36 self.connection = None 

37 self.is_connected = False 

38 

39 def connect(self): 

40 """Set up any connections required by the handler 

41 Should return output of check_connection() method after attempting 

42 connection. Should switch self.is_connected. 

43 Returns: 

44 Connection Object 

45 """ 

46 if self.is_connected is True: 

47 return self.connection 

48 

49 user = self.connection_config.get("user") 

50 password = self.connection_config.get("password") 

51 jdbc_class = self.connection_config.get("jdbcClass") 

52 jar_location = self.connection_config.get("jdbcJarLocation") 

53 

54 jdbc_url = "jdbc:derby://" + self.host + ":" + self.port + "/" + self.database + ";" 

55 

56 if not jdbc_class: 

57 jdbc_class = "org.apache.derby.jdbc.ClientDriver" 

58 

59 if user: 

60 self.schema = user 

61 

62 try: 

63 if user and password and jar_location: 

64 self.connection = jdbcconnector.connect( 

65 jclassname=jdbc_class, url=jdbc_url, driver_args=[user, password], jars=jar_location.split(",") 

66 ) 

67 elif user and password: 

68 self.connection = jdbcconnector.connect( 

69 jclassname=jdbc_class, url=jdbc_url, driver_args=[user, password] 

70 ) 

71 elif jar_location: 

72 self.connection = jdbcconnector.connect( 

73 jclassname=jdbc_class, url=jdbc_url, jars=jar_location.split(",") 

74 ) 

75 else: 

76 self.connection = jdbcconnector.connect(jdbc_class, jdbc_url) 

77 except Exception as e: 

78 logger.error(f"Error while connecting to {self.database}, {e}") 

79 

80 return self.connection 

81 

82 def disconnect(self): 

83 """Close any existing connections 

84 Should switch self.is_connected. 

85 """ 

86 if self.is_connected is False: 

87 return 

88 try: 

89 self.connection.close() 

90 self.is_connected = False 

91 except Exception as e: 

92 logger.error(f"Error while disconnecting to {self.database}, {e}") 

93 

94 return 

95 

96 def check_connection(self) -> StatusResponse: 

97 """Check connection to the handler 

98 Returns: 

99 HandlerStatusResponse 

100 """ 

101 responseCode = StatusResponse(False) 

102 need_to_close = self.is_connected is False 

103 

104 try: 

105 self.connect() 

106 responseCode.success = True 

107 except Exception as e: 

108 logger.error(f"Error connecting to database {self.database}, {e}!") 

109 responseCode.error_message = str(e) 

110 finally: 

111 if responseCode.success is True and need_to_close: 

112 self.disconnect() 

113 if responseCode.success is False and self.is_connected is True: 

114 self.is_connected = False 

115 

116 return responseCode 

117 

118 def native_query(self, query: str) -> StatusResponse: 

119 """Receive raw query and act upon it somehow. 

120 Args: 

121 query (Any): query in native format (str for sql databases, 

122 etc) 

123 Returns: 

124 HandlerResponse 

125 """ 

126 need_to_close = self.is_connected is False 

127 conn = self.connect() 

128 with conn.cursor() as cur: 

129 try: 

130 cur.execute(query) 

131 if cur.description: 

132 result = cur.fetchall() 

133 response = Response( 

134 RESPONSE_TYPE.TABLE, data_frame=pd.DataFrame(result, columns=[x[0] for x in cur.description]) 

135 ) 

136 else: 

137 response = Response(RESPONSE_TYPE.OK) 

138 self.connection.commit() 

139 except Exception as e: 

140 logger.error(f"Error running query: {query} on {self.database}!") 

141 response = Response(RESPONSE_TYPE.ERROR, error_message=str(e)) 

142 self.connection.rollback() 

143 

144 if need_to_close is True: 

145 self.disconnect() 

146 

147 return response 

148 

149 def query(self, query: ASTNode) -> StatusResponse: 

150 """Render and execute a SQL query. 

151 

152 Args: 

153 query (ASTNode): The SQL query. 

154 

155 Returns: 

156 Response: The query result. 

157 """ 

158 if isinstance(query, ASTNode): 

159 query_str = query.to_string() 

160 else: 

161 query_str = str(query) 

162 

163 # Replace backticks with double quotes for Derby compatibility 

164 query_str = query_str.replace("`", '"') 

165 

166 return self.native_query(query_str) 

167 

168 def get_tables(self) -> StatusResponse: 

169 """Get a list of all the tables in the database. 

170 

171 Returns: 

172 Response: Names of the tables in the database. 

173 """ 

174 query = f""" 

175 SELECT st.tablename FROM sys.systables st LEFT OUTER JOIN sys.sysschemas ss ON (st.schemaid = ss.schemaid) WHERE ss.schemaname ='{self.schema}' """ 

176 

177 result = self.native_query(query) 

178 df = result.data_frame 

179 result.data_frame = df.rename(columns={df.columns[0]: "table_name"}) 

180 return result 

181 

182 def get_columns(self, table_name: str) -> StatusResponse: 

183 """Get details about a table. 

184 

185 Args: 

186 table_name (str): Name of the table to retrieve details of. 

187 

188 Returns: 

189 Response: Details of the table. 

190 """ 

191 

192 query = f""" SELECT COLUMNNAME FROM SYS.SYSCOLUMNS INNER JOIN SYS.SYSTABLES ON SYS.SYSCOLUMNS.REFERENCEID=SYS.SYSTABLES.TABLEID WHERE TABLENAME='{table_name}' """ 

193 return self.native_query(query)