Coverage for mindsdb / integrations / handlers / discord_handler / discord_tables.py: 0%

1import pandas as pd 

2from typing import Text, List, Dict, Any 

3 

4from mindsdb_sql_parser import ast 

5 

6from mindsdb.integrations.libs.api_handler import APITable 

7from mindsdb.integrations.utilities.sql_utils import extract_comparison_conditions 

8 

9from mindsdb.integrations.libs.response import HandlerResponse as Response 

10 

11from mindsdb.integrations.utilities.handlers.query_utilities.insert_query_utilities import ( 

12 INSERTQueryParser 

13) 

14 

15from mindsdb.utilities import log 

16 

17logger = log.getLogger(__name__) 

18 

19 

20class MessagesTable(APITable): 

21 """The Discord Table implementation""" 

22 

23 def get_columns(self): 

24 return [ 

25 'id', 

26 'type', 

27 'content', 

28 'author_id', 

29 'author_username', 

30 'author_global_name', 

31 'author_avatar', 

32 'author_banner_color', 

33 'attachments', 

34 'embeds', 

35 'mentions', 

36 'mention_roles', 

37 'pinned', 

38 'mention_everyone', 

39 'tts', 

40 'timestamp', 

41 'edited_timestamp', 

42 'flags', 

43 'components', 

44 'nonce', 

45 'referenced_message', 

46 ] 

47 

48 def select(self, query: ast.Select) -> Response: 

49 """Selects data from the Discord channel. 

50 

51 Parameters 

52 ---------- 

53 query : ast.Select 

54 Given SQL SELECT query. 

55 

56 Returns 

57 ------- 

58 Response 

59 Response object representing collected data from Discord. 

60 """ 

61 

62 conditions = extract_comparison_conditions(query.where) 

63 

64 params = {} 

65 filters = [] 

66 user_filter_flag = channel_filter_flag = False 

67 for op, arg1, arg2 in conditions: 

68 if op == 'or': 

69 raise NotImplementedError('OR is not supported') 

70 if arg1 == 'timestamp': 

71 if op == '>': 

72 params['after'] = arg2 

73 elif op == '<': 

74 params['before'] = arg2 

75 else: 

76 raise NotImplementedError( 

77 f'Unsupported operator {op} for timestamp' 

78 ) 

79 

80 elif arg1 in ['author_id', 'author_username', 'author_global_name']: 

81 if user_filter_flag: 

82 raise NotImplementedError('Multiple user filters are not supported') 

83 user_filter_flag = True 

84 

85 if op != '=': 

86 raise NotImplementedError(f'Unsupported operator {op} for {arg1}') 

87 

88 # if arg1 == 'author_id': 

89 # filters.append(lambda x: x.author.id == int(arg2)) 

90 # elif arg1 == 'author_username': 

91 # filters.append(lambda x: x.author.username == arg2) 

92 # elif arg1 == 'author_global_name': 

93 # filters.append(lambda x: x.author.global_name == arg2) 

94 

95 elif arg1 == 'channel_id': 

96 if op != '=': 

97 raise NotImplementedError(f'Unsupported operator {op} for {arg1}') 

98 channel_filter_flag = True 

99 params['channel_id'] = int(arg2) 

100 

101 else: 

102 filters.append([op, arg1, arg2]) 

103 

104 if query.limit is not None: 

105 params['limit'] = query.limit 

106 else: 

107 params['limit'] = 100 

108 

109 if not channel_filter_flag: 

110 raise NotImplementedError('Channel filter is required') 

111 

112 result = self.handler.call_discord_api( 

113 'get_messages', params=params, filters=filters 

114 ) 

115 

116 # filter targets 

117 columns = [] 

118 for target in query.targets: 

119 if isinstance(target, ast.Star): 

120 columns = [] 

121 break 

122 elif isinstance(target, ast.Identifier): 

123 columns.append(target.parts[-1]) 

124 else: 

125 raise NotImplementedError 

126 

127 if len(columns) == 0: 

128 columns = self.get_columns() 

129 

130 # columns to lower case 

131 columns = [name.lower() for name in columns] 

132 

133 if len(result) == 0: 

134 result = pd.DataFrame([], columns=columns) 

135 else: 

136 # add absent columns 

137 for col in set(columns) & set(result.columns) ^ set(columns): 

138 result[col] = None 

139 

140 # filter by columns 

141 result = result[columns] 

142 return result 

143 

144 def insert(self, query: ast.Insert) -> None: 

145 """Sends messages to a Discord channel. 

146 

147 Parameters 

148 ---------- 

149 query : ast.Insert 

150 Given SQL INSERT query. 

151 

152 Returns 

153 ------- 

154 None 

155 

156 Raises 

157 ------ 

158 ValueError 

159 If the query contains an unsupported condition 

160 """ 

161 insert_statement_parser = INSERTQueryParser( 

162 query, 

163 supported_columns=['channel_id', 'text'], 

164 mandatory_columns=['channel_id', 'text'], 

165 all_mandatory=False, 

166 ) 

167 message_data = insert_statement_parser.parse_query() 

168 self.send_message(message_data) 

169 

170 def send_message(self, message_data: List[Dict[Text, Any]]) -> None: 

171 """Sends messages to a Discord Channel using the parsed message data. 

172 

173 Parameters 

174 ---------- 

175 message_data : List[Dict[Text, Any]] 

176 List of dictionaries containing the messages to send. 

177 

178 Returns 

179 ------- 

180 None 

181 """ 

182 for message in message_data: 

183 try: 

184 params = {'channel_id': message['channel_id'], 'text': message['text']} 

185 self.handler.call_discord_api('send_message', params=params) 

186 logger.info("Message sent to Discord channel successfully.") 

187 except Exception as e: 

188 logger.error(f"Error sending message to Discord channel: {e}") 

189 raise e