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

1import requests 

2import pandas as pd 

3 

4from mindsdb.integrations.handlers.discord_handler.discord_tables import MessagesTable 

5 

6from mindsdb.utilities import log 

7 

8from mindsdb.integrations.libs.api_handler import APIHandler, FuncParser 

9from mindsdb.integrations.utilities.date_utils import parse_utc_date 

10 

11from mindsdb.integrations.libs.response import ( 

12 HandlerStatusResponse as StatusResponse, 

13 HandlerResponse as Response, 

14 RESPONSE_TYPE, 

15) 

16 

17discord_bot = None 

18logger = log.getLogger(__name__) 

19 

20 

21class DiscordHandler(APIHandler): 

22 """ 

23 The Discord handler implementation. 

24 """ 

25 

26 name = 'discord' 

27 

28 def __init__(self, name: str, **kwargs): 

29 """ 

30 Initialize the handler. 

31 Args: 

32 name (str): name of particular handler instance 

33 **kwargs: arbitrary keyword arguments. 

34 """ 

35 super().__init__(name) 

36 

37 connection_data = kwargs.get("connection_data", {}) 

38 self.connection_data = connection_data 

39 self.kwargs = kwargs 

40 

41 self.is_connected = False 

42 

43 messages = MessagesTable(self) 

44 self._register_table('messages', messages) 

45 

46 def connect(self): 

47 """ 

48 Set up the connection required by the handler. 

49 Returns 

50 ------- 

51 StatusResponse 

52 connection object 

53 """ 

54 

55 if self.is_connected: 

56 return StatusResponse(True) 

57 

58 url = 'https://discord.com/api/v10/applications/@me' 

59 result = requests.get( 

60 url, 

61 headers={ 

62 'Authorization': f'Bot {self.connection_data["token"]}', 

63 'Content-Type': 'application/json', 

64 }, 

65 ) 

66 

67 if result.status_code != 200: 

68 raise ValueError(result.text) 

69 

70 self.is_connected = True 

71 

72 def check_connection(self) -> StatusResponse: 

73 """ 

74 Check connection to the handler. 

75 Returns: 

76 HandlerStatusResponse 

77 """ 

78 

79 response = StatusResponse(False) 

80 

81 try: 

82 self.connect() 

83 response.success = True 

84 except Exception as e: 

85 response.error_message = e 

86 logger.error(f'Error connecting to Discord: {response.error_message}') 

87 

88 self.is_connected = response.success 

89 

90 return response 

91 

92 def native_query(self, query: str = None) -> StatusResponse: 

93 """Receive and process a raw query. 

94 Parameters 

95 ---------- 

96 query : str 

97 query in a native format 

98 Returns 

99 ------- 

100 StatusResponse 

101 Request status 

102 """ 

103 operation, params = FuncParser().from_string(query) 

104 

105 df = self.call_discord_api(operation, params) 

106 

107 return Response(RESPONSE_TYPE.TABLE, data_frame=df) 

108 

109 def utc_to_snowflake(self, utc_date: str) -> int: 

110 """ 

111 Convert a UTC date to a Snowflake date. 

112 Args: 

113 utc_date (str): the UTC date 

114 Returns: 

115 int 

116 """ 

117 # https://discord.com/developers/docs/reference#snowflakes 

118 return str( 

119 int(parse_utc_date(utc_date).timestamp() * 1000 - 1420070400000) << 22 

120 ) 

121 

122 def call_discord_api( 

123 self, operation: str, params: dict = None, filters: list = None 

124 ): 

125 """ 

126 Call a Discord API method. 

127 Args: 

128 method_name (str): the method name 

129 params (dict): the method parameters 

130 Returns: 

131 pd.DataFrame 

132 """ 

133 

134 if operation == 'get_messages': 

135 param_strings = {'limit': params['limit']} 

136 if 'after' in params: 

137 param_strings['after'] = self.utc_to_snowflake(params['after']) 

138 if 'before' in params: 

139 param_strings['before'] = self.utc_to_snowflake(params['before']) 

140 

141 url = ( 

142 f'https://discord.com/api/v10/channels/{params["channel_id"]}/messages' 

143 ) 

144 result = requests.get( 

145 url, 

146 headers={ 

147 'Authorization': f'Bot {self.connection_data["token"]}', 

148 'Content-Type': 'application/json', 

149 }, 

150 params=param_strings, 

151 ) 

152 

153 if result.status_code != 200: 

154 raise ValueError(f'Error calling Discord API: {result.json()}') 

155 

156 json = result.json() 

157 for message in json: 

158 author = message.get('author') 

159 if author is not None: 

160 message['author_id'] = author.get('id') 

161 message['author_username'] = author.get('username') 

162 message['author_global_name'] = author.get('global_name') 

163 

164 df = pd.DataFrame.from_records(json) 

165 return df 

166 elif operation == 'send_message': 

167 url = ( 

168 f'https://discord.com/api/v10/channels/{params["channel_id"]}/messages' 

169 ) 

170 result = requests.post( 

171 url, 

172 headers={ 

173 'Authorization': f'Bot {self.connection_data["token"]}', 

174 'Content-Type': 'application/json', 

175 }, 

176 json={ 

177 'content': params['text'], 

178 }, 

179 ) 

180 

181 if result.status_code != 200: 

182 raise ValueError(f'Error calling Discord API: {result.json()}') 

183 

184 df = pd.DataFrame.from_records([result.json()]) 

185 return df 

186 else: 

187 raise ValueError(f"Unsupported method: {operation}")