feat: re-enable safe quote

litestar-org · Dec 30, 2024 · 1c3f9d8 · 1c3f9d8
1 parent 13a82e5
commit 1c3f9d8
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/litestar/channels/backends/psycopg.py b/litestar/channels/backends/psycopg.py
@@ -9,6 +9,10 @@
 from .base import ChannelsBackend
 
 
+def _safe_quote(ident: str) -> str:
+    return '"{}"'.format(ident.replace('"', '""'))  # sourcery skip
+
+
 class PsycoPgChannelsBackend(ChannelsBackend):
     _listener_conn: AsyncConnection[Any]
 
@@ -32,13 +36,13 @@ async def publish(self, data: bytes, channels: Iterable[str]) -> None:
 
     async def subscribe(self, channels: Iterable[str]) -> None:
         for channel in set(channels) - self._subscribed_channels:
-            await self._listener_conn.execute(SQL("LISTEN {}").format(Identifier(channel)))
+            await self._listener_conn.execute(SQL("LISTEN {}").format(Identifier(_safe_quote(channel))))
 
             self._subscribed_channels.add(channel)
 
     async def unsubscribe(self, channels: Iterable[str]) -> None:
         for channel in channels:
-            await self._listener_conn.execute(SQL("UNLISTEN {}").format(Identifier(channel)))
+            await self._listener_conn.execute(SQL("UNLISTEN {}").format(Identifier(_safe_quote(channel))))
         self._subscribed_channels = self._subscribed_channels - set(channels)
 
     async def stream_events(self) -> AsyncGenerator[tuple[str, bytes], None]: