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

1from pandas import DataFrame 

2 

3import sqlanydb 

4import sqlalchemy_sqlany.base as sqlany_dialect 

5 

6from mindsdb_sql_parser import parse_sql 

7from mindsdb_sql_parser.ast.base import ASTNode 

8from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender 

9 

10from mindsdb.utilities import log 

11from mindsdb.integrations.libs.base import DatabaseHandler 

12from mindsdb.integrations.libs.response import ( 

13 HandlerStatusResponse as StatusResponse, 

14 HandlerResponse as Response, 

15 RESPONSE_TYPE 

16) 

17 

18 

19logger = log.getLogger(__name__) 

20 

21 

22class SQLAnyHandler(DatabaseHandler): 

23 """ 

24 This handler handles connection and execution of the SAP SQL Anywhere statements. 

25 """ 

26 

27 name = 'sqlany' 

28 

29 def __init__(self, name: str, connection_data: dict, **kwargs): 

30 super().__init__(name) 

31 

32 self.dialect = 'sqlany' 

33 self.parser = parse_sql 

34 self.connection_data = connection_data 

35 self.renderer = SqlalchemyRender(sqlany_dialect.SQLAnyDialect) 

36 self.host = self.connection_data.get('host') 

37 self.port = self.connection_data.get('port') 

38 self.userid = self.connection_data.get('user') 

39 self.password = self.connection_data.get('password') 

40 self.server = self.connection_data.get('server') 

41 self.databaseName = self.connection_data.get('database') 

42 self.encryption = self.connection_data.get('encrypt', False) 

43 self.connection = None 

44 self.is_connected = False 

45 

46 def __del__(self): 

47 if self.is_connected is True: 

48 self.disconnect() 

49 

50 def connect(self): 

51 """ 

52 Handles the connection to a SAP SQL Anywhere database insance. 

53 """ 

54 

55 if self.is_connected is True: 

56 return self.connection 

57 

58 if self.port.strip().isnumeric(): 

59 self.host += ":" + self.port.strip() 

60 

61 if self.encryption: 

62 self.encryption = "SIMPLE" 

63 else: 

64 self.encryption = "NONE" 

65 

66 connection = sqlanydb.connect( 

67 host=self.host, 

68 userid=self.userid, 

69 password=self.password, 

70 server=self.server, 

71 databaseName=self.databaseName, 

72 encryption=self.encryption 

73 ) 

74 self.is_connected = True 

75 self.connection = connection 

76 return self.connection 

77 

78 def disconnect(self): 

79 """ 

80 Disconnects from the SAP SQL Anywhere database 

81 """ 

82 

83 if self.is_connected is True: 

84 self.connection.close() 

85 self.is_connected = False 

86 

87 def check_connection(self) -> StatusResponse: 

88 """ 

89 Check the connection of the SAP SQL Anywhere database 

90 :return: success status and error message if error occurs 

91 """ 

92 

93 response = StatusResponse(False) 

94 need_to_close = self.is_connected is False 

95 

96 try: 

97 connection = self.connect() 

98 cur = connection.cursor() 

99 cur.execute('SELECT 1 FROM SYS.DUMMY;') 

100 response.success = True 

101 except sqlanydb.Error as e: 

102 logger.error(f'Error connecting to SAP SQL Anywhere {self.host}, {e}!') 

103 response.error_message = e 

104 

105 if response.success is True and need_to_close: 

106 self.disconnect() 

107 if response.success is False and self.is_connected is True: 

108 self.is_connected = False 

109 

110 return response 

111 

112 def native_query(self, query: str) -> Response: 

113 """ 

114 Receive SQL query and runs it 

115 :param query: The SQL query to run in SAP SQL Anywhere 

116 :return: returns the records from the current recordset 

117 """ 

118 

119 need_to_close = self.is_connected is False 

120 

121 connection = self.connect() 

122 cur = connection.cursor() 

123 try: 

124 cur.execute(query) 

125 if not cur.description: 

126 response = Response(RESPONSE_TYPE.OK) 

127 else: 

128 result = cur.fetchall() 

129 response = Response( 

130 RESPONSE_TYPE.TABLE, 

131 DataFrame( 

132 result, 

133 columns=[x[0] for x in cur.description] 

134 ) 

135 ) 

136 connection.commit() 

137 except Exception as e: 

138 logger.error(f'Error running query: {query} on {self.connection}!') 

139 response = Response( 

140 RESPONSE_TYPE.ERROR, 

141 error_code=0, 

142 error_message=str(e) 

143 ) 

144 connection.rollback() 

145 

146 if need_to_close is True: 

147 self.disconnect() 

148 

149 return response 

150 

151 def query(self, query: ASTNode) -> Response: 

152 """ 

153 Retrieve the data from the SQL statement with eliminated rows that dont satisfy the WHERE condition 

154 """ 

155 

156 query_str = self.renderer.get_string(query, with_failback=True) 

157 return self.native_query(query_str) 

158 

159 def get_tables(self) -> Response: 

160 """ 

161 List all tables in SAP SQL Anywhere in the current schema 

162 """ 

163 

164 return self.native_query(""" 

165 SELECT USER_NAME(ob.UID) AS SCHEMA_NAME 

166 , st.TABLE_NAME 

167 , st.TABLE_TYPE 

168 FROM SYSOBJECTS ob 

169 INNER JOIN SYS.SYSTABLE st on ob.ID = st.OBJECT_ID 

170 WHERE ob.TYPE='U' AND st.TABLE_TYPE <> 'GBL TEMP' 

171 """) 

172 

173 def get_columns(self, table_name: str) -> Response: 

174 """ 

175 List all columns in a table in SAP SQL Anywhere in the current schema 

176 :param table_name: the table name for which to list the columns 

177 :return: returns the columns in the table 

178 """ 

179 

180 return self.renderer.dialect.get_columns(table_name)