stop using async with and make external clients use node session

This commit is contained in:
cloudwithax 2023-05-01 21:01:12 -04:00
parent b91f6ec04e
commit 394e3a3907
No known key found for this signature in database
GPG Key ID: 5DBE54E45794983E
3 changed files with 74 additions and 84 deletions

View File

@ -42,10 +42,10 @@ class Client:
self.session: aiohttp.ClientSession = None # type: ignore self.session: aiohttp.ClientSession = None # type: ignore
self._log = logging.getLogger(__name__) self._log = logging.getLogger(__name__)
async def request_token(self) -> None: async def _set_session(self, session: aiohttp.ClientSession) -> None:
if not self.session: self.session = session
self.session = aiohttp.ClientSession()
async def request_token(self) -> None:
# First lets get the raw response from the main page # First lets get the raw response from the main page
resp = await self.session.get("https://music.apple.com") resp = await self.session.get("https://music.apple.com")
@ -187,8 +187,3 @@ class Client:
next_page_url = None next_page_url = None
return Playlist(data, album_tracks) return Playlist(data, album_tracks)
async def close(self) -> None:
if self.session:
await self.session.close()
self.session = None # type: ignore

View File

@ -274,6 +274,13 @@ class Node:
"Lavalink version 3.7.0 or above is required to use this library.", "Lavalink version 3.7.0 or above is required to use this library.",
) )
async def _set_ext_client_session(self, session: aiohttp.ClientSession) -> None:
if self._spotify_client:
await self._spotify_client._set_session(session=session)
if self._apple_music_client:
await self._apple_music_client._set_session(session=session)
async def _update_handler(self, data: dict) -> None: async def _update_handler(self, data: dict) -> None:
await self._bot.wait_until_ready() await self._bot.wait_until_ready()
@ -373,12 +380,12 @@ class Node:
f'{f"?{query}" if query else ""}' f'{f"?{query}" if query else ""}'
) )
async with self._session.request( resp = await self._session.request(
method=method, method=method,
url=uri, url=uri,
headers=self._headers, headers=self._headers,
json=data or {}, json=data or {},
) as resp: )
self._log.debug( self._log.debug(
f"Making REST request to Node {self._identifier} with method {method} to {uri}", f"Making REST request to Node {self._identifier} with method {method} to {uri}",
) )
@ -428,6 +435,7 @@ class Node:
) )
await self._handle_version_check(version=version) await self._handle_version_check(version=version)
await self._set_ext_client_session(session=self._session)
self._log.debug( self._log.debug(
f"Version check from Node {self._identifier} successful. Returned version {version}", f"Version check from Node {self._identifier} successful. Returned version {version}",
@ -485,14 +493,6 @@ class Node:
await self._session.close() await self._session.close()
self._log.debug("Websocket and http session closed.") self._log.debug("Websocket and http session closed.")
if self._spotify_client:
await self._spotify_client.close()
self._log.debug("Spotify client session closed.")
if self._apple_music_client:
await self._apple_music_client.close()
self._log.debug("Apple Music client session closed.")
del self._pool._nodes[self._identifier] del self._pool._nodes[self._identifier]
self.available = False self.available = False
self._task.cancel() self._task.cancel()

View File

@ -49,13 +49,14 @@ class Client:
self._bearer_headers: Optional[Dict] = None self._bearer_headers: Optional[Dict] = None
self._log = logging.getLogger(__name__) self._log = logging.getLogger(__name__)
async def _set_session(self, session: aiohttp.ClientSession) -> None:
self.session = session
async def _fetch_bearer_token(self) -> None: async def _fetch_bearer_token(self) -> None:
_data = {"grant_type": "client_credentials"} _data = {"grant_type": "client_credentials"}
if not self.session: resp = await self.session.post(GRANT_URL, data=_data, headers=self._grant_headers)
self.session = aiohttp.ClientSession()
async with self.session.post(GRANT_URL, data=_data, headers=self._grant_headers) as resp:
if resp.status != 200: if resp.status != 200:
raise SpotifyRequestException( raise SpotifyRequestException(
f"Error fetching bearer token: {resp.status} {resp.reason}", f"Error fetching bearer token: {resp.status} {resp.reason}",
@ -83,7 +84,7 @@ class Client:
request_url = REQUEST_URL.format(type=spotify_type, id=spotify_id) request_url = REQUEST_URL.format(type=spotify_type, id=spotify_id)
async with self.session.get(request_url, headers=self._bearer_headers) as resp: resp = await self.session.get(request_url, headers=self._bearer_headers)
if resp.status != 200: if resp.status != 200:
raise SpotifyRequestException( raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}", f"Error while fetching results: {resp.status} {resp.reason}",
@ -99,10 +100,10 @@ class Client:
elif spotify_type == "album": elif spotify_type == "album":
return Album(data) return Album(data)
elif spotify_type == "artist": elif spotify_type == "artist":
async with self.session.get( resp = await self.session.get(
f"{request_url}/top-tracks?market=US", f"{request_url}/top-tracks?market=US",
headers=self._bearer_headers, headers=self._bearer_headers,
) as resp: )
if resp.status != 200: if resp.status != 200:
raise SpotifyRequestException( raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}", f"Error while fetching results: {resp.status} {resp.reason}",
@ -126,7 +127,7 @@ class Client:
next_page_url = data["tracks"]["next"] next_page_url = data["tracks"]["next"]
while next_page_url is not None: while next_page_url is not None:
async with self.session.get(next_page_url, headers=self._bearer_headers) as resp: resp = await self.session.get(next_page_url, headers=self._bearer_headers)
if resp.status != 200: if resp.status != 200:
raise SpotifyRequestException( raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}", f"Error while fetching results: {resp.status} {resp.reason}",
@ -164,19 +165,13 @@ class Client:
id=f"?seed_tracks={spotify_id}", id=f"?seed_tracks={spotify_id}",
) )
async with self.session.get(request_url, headers=self._bearer_headers) as resp: resp = await self.session.get(request_url, headers=self._bearer_headers)
if resp.status != 200: if resp.status != 200:
raise SpotifyRequestException( raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}", f"Error while fetching results: {resp.status} {resp.reason}",
) )
data: dict = await resp.json(loads=json.loads) data: dict = await resp.json(loads=json.loads)
tracks = [Track(track) for track in data["tracks"]] tracks = [Track(track) for track in data["tracks"]]
return tracks return tracks
async def close(self) -> None:
if self.session:
await self.session.close()
self.session = None # type: ignore