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

1from typing import Optional 

2 

3import pandas as pd 

4import duckdb 

5 

6from mindsdb_sql_parser import parse_sql 

7from mindsdb_sql_parser.ast.base import ASTNode 

8from mindsdb.utilities import log 

9from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender 

10from mindsdb.integrations.libs.base import DatabaseHandler 

11from mindsdb.integrations.libs.response import ( 

12 HandlerStatusResponse as StatusResponse, 

13 HandlerResponse as Response, 

14 RESPONSE_TYPE, 

15) 

16 

17 

18logger = log.getLogger(__name__) 

19 

20 

21class SheetsHandler(DatabaseHandler): 

22 """ 

23 This handler handles connection and execution of queries against the Excel Sheet. 

24 TODO: add authentication for private sheets 

25 """ 

26 

27 name = "sheets" 

28 

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

30 """ 

31 Initialize the handler. 

32 Args: 

33 name (str): name of particular handler instance 

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

35 **kwargs: arbitrary keyword arguments. 

36 """ 

37 super().__init__(name) 

38 self.parser = parse_sql 

39 self.renderer = SqlalchemyRender("postgresql") 

40 self.connection_data = connection_data 

41 self.kwargs = kwargs 

42 

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) -> StatusResponse: 

51 """ 

52 Set up the connection required by the handler. 

53 Returns: 

54 HandlerStatusResponse 

55 """ 

56 url = f"https://docs.google.com/spreadsheets/d/{self.connection_data['spreadsheet_id']}/gviz/tq?tqx=out:csv&sheet={self.connection_data['sheet_name']}" 

57 self.sheet = pd.read_csv(url, on_bad_lines="skip") 

58 self.connection = duckdb.connect() 

59 self.connection.register(self.connection_data["sheet_name"], self.sheet) 

60 self.is_connected = True 

61 

62 return self.connection 

63 

64 def disconnect(self): 

65 """ 

66 Close any existing connections. 

67 """ 

68 if self.is_connected is False: 

69 return 

70 

71 self.connection.close() 

72 self.is_connected = False 

73 return self.is_connected 

74 

75 def check_connection(self) -> StatusResponse: 

76 """ 

77 Check connection to the handler. 

78 Returns: 

79 HandlerStatusResponse 

80 """ 

81 response = StatusResponse(False) 

82 need_to_close = self.is_connected is False 

83 

84 try: 

85 self.connect() 

86 response.success = True 

87 except Exception as e: 

88 logger.error(f"Error connecting to the Google Sheet with ID {self.connection_data['spreadsheet_id']}, {e}!") 

89 response.error_message = str(e) 

90 finally: 

91 if response.success is True and need_to_close: 

92 self.disconnect() 

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

94 self.is_connected = False 

95 

96 return response 

97 

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

99 """ 

100 Receive raw query and act upon it somehow. 

101 Args: 

102 query (str): query in native format 

103 Returns: 

104 HandlerResponse 

105 """ 

106 

107 need_to_close = self.is_connected is False 

108 connection = self.connect() 

109 try: 

110 result = connection.execute(query).fetchdf() 

111 if not result.empty: 

112 response = Response(RESPONSE_TYPE.TABLE, result) 

113 else: 

114 response = Response(RESPONSE_TYPE.OK) 

115 connection.commit() 

116 except Exception as e: 

117 logger.error( 

118 f"Error running query: {query} on the Google Sheet with ID {self.connection_data['spreadsheet_id']}!" 

119 ) 

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

121 

122 if need_to_close is True: 

123 self.disconnect() 

124 

125 return response 

126 

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

128 """ 

129 Receive query as AST (abstract syntax tree) and act upon it somehow. 

130 Args: 

131 query (ASTNode): sql query represented as AST. May be any kind 

132 of query: SELECT, INTSERT, DELETE, etc 

133 Returns: 

134 HandlerResponse 

135 """ 

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

137 return self.native_query(query_str) 

138 

139 def get_tables(self) -> StatusResponse: 

140 """ 

141 Return list of entities that will be accessible as tables. 

142 Returns: 

143 HandlerResponse 

144 """ 

145 response = Response( 

146 RESPONSE_TYPE.TABLE, data_frame=pd.DataFrame([self.connection_data["sheet_name"]], columns=["table_name"]) 

147 ) 

148 

149 return response 

150 

151 def get_columns(self) -> StatusResponse: 

152 """ 

153 Returns a list of entity columns. 

154 Args: 

155 table_name (str): name of one of tables returned by self.get_tables() 

156 Returns: 

157 HandlerResponse 

158 """ 

159 response = Response( 

160 RESPONSE_TYPE.TABLE, 

161 data_frame=pd.DataFrame({"column_name": list(self.sheet.columns), "data_type": self.sheet.dtypes}), 

162 ) 

163 

164 return response