graceful shutdown for aiohttp session
This commit is contained in:
		
							
								
								
									
										30
									
								
								bot.py
									
									
									
									
									
								
							
							
						
						
									
										30
									
								
								bot.py
									
									
									
									
									
								
							| @@ -6,16 +6,40 @@ from discord.ext import commands | |||||||
| import heavynode | import heavynode | ||||||
|  |  | ||||||
|  |  | ||||||
|  | class CustomBot(commands.Bot): | ||||||
|  |     def __init__(self, *args, **kwargs): | ||||||
|  |         self.cleanup_sync = [] | ||||||
|  |         self.cleanup_async = [] | ||||||
|  |         super().__init__(*args, **kwargs) | ||||||
|  |  | ||||||
|  |     def add_cleanup_sync(self, func): | ||||||
|  |         self.cleanup_sync.append(func) | ||||||
|  |  | ||||||
|  |     def add_cleanup_async(self, coro): | ||||||
|  |         self.cleanup_async.append(coro) | ||||||
|  |  | ||||||
|  |     async def close(self): | ||||||
|  |         for func in self.cleanup_sync: | ||||||
|  |             func() | ||||||
|  |         for coro in self.cleanup_async: | ||||||
|  |             await coro() | ||||||
|  |         await super().close() | ||||||
|  |  | ||||||
|  |  | ||||||
| logging.basicConfig() | logging.basicConfig() | ||||||
|  |  | ||||||
|  | server_id = 'd030737c' | ||||||
| hn = heavynode.Client(os.environ['heavynode_token']) | hn = heavynode.Client(os.environ['heavynode_token']) | ||||||
| bot = commands.Bot(command_prefix='!') |  | ||||||
|  | bot = CustomBot(command_prefix='!') | ||||||
|  | bot.add_cleanup_async(hn.shutdown) | ||||||
|  |  | ||||||
|  |  | ||||||
| @bot.command() | @bot.command() | ||||||
| @commands.has_role('Server admin') | @commands.has_role('Server admin') | ||||||
| async def add(ctx, player): | async def add(ctx, player): | ||||||
|     """Add a player to the server whitelist.Must use exact Minecraft name.""" |     """Add a player to the server whitelist.Must use exact Minecraft name.""" | ||||||
|     await hn.send_command('d030737c', f'whitelist add {player}') |     await hn.send_command(server_id, f'whitelist add {player}') | ||||||
|     await ctx.send(f'"{player}" added to whitelist.') |     await ctx.send(f'"{player}" added to whitelist.') | ||||||
|  |  | ||||||
|  |  | ||||||
| @@ -23,7 +47,7 @@ async def add(ctx, player): | |||||||
| @commands.has_role('Server admin') | @commands.has_role('Server admin') | ||||||
| async def remove(ctx, player): | async def remove(ctx, player): | ||||||
|     """Remove a player from the server whitelist. Must use exact Minecraft name.""" |     """Remove a player from the server whitelist. Must use exact Minecraft name.""" | ||||||
|     await hn.send_command('d030737c', f'whitelist remove {player}') |     await hn.send_command(server_id, f'whitelist remove {player}') | ||||||
|     await ctx.send(f'"{player}" removed from whitelist.') |     await ctx.send(f'"{player}" removed from whitelist.') | ||||||
|  |  | ||||||
|  |  | ||||||
|   | |||||||
| @@ -34,3 +34,6 @@ class Client: | |||||||
|     async def send_command(self, serverid, cmd): |     async def send_command(self, serverid, cmd): | ||||||
|         payload = {'command': cmd} |         payload = {'command': cmd} | ||||||
|         return await self.make_request('POST', f'/client/servers/{serverid}/command', json=payload) |         return await self.make_request('POST', f'/client/servers/{serverid}/command', json=payload) | ||||||
|  |  | ||||||
|  |     async def shutdown(self): | ||||||
|  |         await self.session.close() | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user