fix logger link + Axum feature flag

specta-rs · Jul 19, 2023 · b8255af · b8255af
1 parent 1d1a876
commit b8255af
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 18 deletions.
diff --git a/packages/client/src/links/loggerLink.ts b/packages/client/src/links/loggerLink.ts
@@ -2,35 +2,42 @@ import { Link } from "./link";
 
 // TODO: Pretty log output like tRPC's logger link
 
+export type LoggerLinkOpts = {
+  enabled: boolean | (() => boolean);
+};
+
 /**
  * Link for logging operations.
  *
  * This must go before the terminating link for it to work!
  *
  */
-export function loggerLink(): Link {
+export function loggerLink(opts?: LoggerLinkOpts): Link {
+  const { enabled = true } = opts ?? {};
+  const isEnabled = () => (typeof enabled === "function" ? enabled() : enabled);
+
   return ({ op, next }) => {
     const result = next({
       op,
     });
 
-    console.log("REQUEST", op, next);
+    if (isEnabled()) console.log("REQUEST", op, next);
 
     return {
       exec: (resolve, reject) => {
         result.exec(
           (data) => {
-            console.log("RESPONSE", op, data);
+            if (isEnabled()) console.log("RESPONSE", op, data);
             resolve(data);
           },
           (err) => {
-            console.error("RESPONSE ERROR", op, err);
+            if (isEnabled()) console.error("RESPONSE ERROR", op, err);
             reject(err);
           }
         );
       },
       abort: () => {
-        console.log("ABORT OP", op);
+        if (isEnabled()) console.log("ABORT OP", op);
         result.abort();
       },
     };

diff --git a/src/error.rs b/src/error.rs
@@ -104,6 +104,7 @@ impl From<ExecError> for ResponseError {
                 ExecError::OperationNotFound => ErrorCode::NotFound,
                 ExecError::DeserializingArgErr(_) => ErrorCode::BadRequest,
                 ExecError::SerializingResultErr(_) => ErrorCode::InternalServerError,
+                #[cfg(feature = "axum")]
                 ExecError::AxumExtractorError => ErrorCode::BadRequest,
                 ExecError::ErrResolverError(err) => err.code,
                 ExecError::ErrSubscriptionWithNullId => ErrorCode::BadRequest,

diff --git a/src/internal/exec/connection.rs b/src/internal/exec/connection.rs
@@ -330,20 +330,21 @@ impl<
         cx: &mut Context<'_>,
     ) -> Poll<PollResult> {
         let mut conn = this.conn.as_mut().project();
-        match ready!(conn.streams.as_mut().poll_next(cx)) {
-            Some((a, _)) => match a {
-                StreamYield::Item(resp) => {
-                    this.batch.as_mut().insert(resp);
-                    return PollResult::QueueSend.into();
-                }
-                StreamYield::Finished(f) => {
-                    f.take(conn.streams.as_mut());
-                }
-            },
-            // If no streams, fall asleep until a new subscription is queued
-            None => {}
+        for _ in 0..conn.streams.len() {
+            match ready!(conn.streams.as_mut().poll_next(cx)) {
+                Some((a, _)) => match a {
+                    StreamYield::Item(resp) => {
+                        this.batch.as_mut().insert(resp);
+                        return PollResult::QueueSend.into();
+                    }
+                    StreamYield::Finished(f) => {
+                        f.take(conn.streams.as_mut());
+                    }
+                },
+                // If no streams, fall asleep until a new subscription is queued
+                None => {}
+            }
         }
-
         PollResult::Progressed.into()
     }