Add ability to attach request information and wait for flush

examples/logger/pages/_app.js

@@ -1,7 +1,8 @@
-import { log } from 'next-axiom'
+import { Logger } from 'next-axiom'
 
 export { reportWebVitals } from 'next-axiom'
 
+const log = new Logger()
 log.info('Hello from frontend', { foo: 'bar' })
 
 function MyApp({ Component, pageProps }) {

examples/logger/pages/api/edge.js

@@ -1,11 +1,11 @@
-import { log, withAxiom } from 'next-axiom'
+import { withAxiom } from 'next-axiom'
 
 export const config = {
     runtime: 'experimental-edge',
 };
 
-function handler() {
-    log.debug("message from edge", { foo: 'bar' })
+function handler(req) {
+    req.log.debug("message from edge", { foo: 'bar' })
 
     return new Response(
         JSON.stringify({

examples/logger/pages/api/hello.js

@@ -1,8 +1,8 @@
 // Next.js API route support: https://nextjs.org/docs/api-routes/introduction
-import { log, withAxiom } from 'next-axiom'
+import { withAxiom } from 'next-axiom'
 
 async function handler(req, res) {
-  log.info('Hello from function', { url: req.url });
+  req.log.info('Hello from function', { url: req.url });
   res.status(200).json({ name: 'John Doe' })
 }
 

examples/logger/pages/index.js

@@ -1,6 +1,8 @@
-import { log } from 'next-axiom'
+import { Logger } from 'next-axiom'
 import useSWR from 'swr'
 
+const log = new Logger()
+
 export async function getStaticProps(context) {
   log.info('Hello from SSR', { context })
   return {

src/index.ts

@@ -1,3 +1,3 @@
 export { reportWebVitals } from './webVitals';
-export { log } from './logger';
+export { Logger } from './logger';
 export { withAxiom } from './withAxiom';

src/logger.ts

@@ -1,84 +1,108 @@
+import { NextRequest } from 'next/server';
 import { proxyPath, isBrowser, EndpointType, getIngestURL, isEnvVarsSet, isNoPrettyPrint } from './shared';
 import { throttle } from './shared';
 
 const url = isBrowser ? `${proxyPath}/logs` : getIngestURL(EndpointType.logs);
-const throttledSendLogs = throttle(sendLogs, 1000);
-let logEvents: any[] = [];
-
-function _log(level: string, message: string, args: any = {}) {
-  if (!isEnvVarsSet) {
-    // if AXIOM ingesting url is not set, fallback to printing to console
-    // to avoid network errors in development environments
-    prettyPrint(level, message, args);
-    return;
-  }
-
-  const logEvent = { level, message, _time: new Date(Date.now()).toISOString() };
-  if (Object.keys(args).length > 0) {
-    logEvent['fields'] = args;
-  }
 
-  logEvents.push(logEvent);
-  throttledSendLogs();
+interface LogEvent {
+  level: string;
+  message: string;
+  fields: {};
+  _time: string;
+  request?: any;
 }
 
-export const log = {
-  debug: (message: string, args: any = {}) => _log('debug', message, args),
-  info: (message: string, args: any = {}) => _log('info', message, args),
-  warn: (message: string, args: any = {}) => _log('warn', message, args),
-  error: (message: string, args: any = {}) => _log('error', message, args),
-  with: (args: any) => new Logger(args),
-  flush: sendLogs,
-};
-
 export class Logger {
-  args: any = {};
+  public logEvents: any[] = [];
+  throttledSendLogs = throttle(this.sendLogs, 1000);
 
-  constructor(args: any = {}) {
-    this.args = args;
-  }
+  constructor(private args: any = {}, private req: NextRequest | null = null, private waitForFlush: Boolean = false) {}
 
   debug(message: string, args: any = {}) {
-    _log('debug', message, { ...this.args, ...args });
+    this._log('debug', message, { ...this.args, ...args });
   }
   info(message: string, args: any = {}) {
-    _log('info', message, { ...this.args, ...args });
+    this._log('info', message, { ...this.args, ...args });
   }
   warn(message: string, args: any = {}) {
-    _log('warn', message, { ...this.args, ...args });
+    this._log('warn', message, { ...this.args, ...args });
   }
   error(message: string, args: any = {}) {
-    _log('error', message, { ...this.args, ...args });
+    this._log('error', message, { ...this.args, ...args });
   }
 
   with(args: any) {
-    return new Logger({ ...this.args, ...args });
+    return new Logger({ ...this.args, ...args }, this.req, this.waitForFlush);
   }
-}
 
-async function sendLogs() {
-  if (!logEvents.length) {
-    return;
+  withRequest(req: NextRequest) {
+    return new Logger({ ...this.args }, req, this.waitForFlush);
+  }
+
+  _log(level: string, message: string, args: any = {}) {
+    const logEvent: LogEvent = { level, message, _time: new Date(Date.now()).toISOString(), fields: {} };
+    if (Object.keys(args).length > 0) {
+      logEvent['fields'] = args;
+    }
+
+    if (this.req != null) {
+      logEvent['request'] = {
+        ip: this.req.ip,
+        region: this.req.geo?.region,
+        host: this.req.nextUrl.host,
+        method: this.req.method,
+        path: this.req.nextUrl.pathname,
+        scheme: this.req.nextUrl.protocol.replace(':', ''),
+        userAgent: this.req.headers.get('user-agent'),
+      };
+    }
+
+    this.logEvents.push(logEvent);
+    if (!this.waitForFlush) {
+      this.throttledSendLogs();
+    }
   }
 
-  const method = 'POST';
-  const keepalive = true;
-  const body = JSON.stringify(logEvents);
-  // clear pending logs
-  logEvents = [];
-
-  try {
-    if (typeof fetch === 'undefined') {
-      const fetch = await require('whatwg-fetch');
-      await fetch(url, { body, method, keepalive });
-    } else if (isBrowser && navigator.sendBeacon) {
-      navigator.sendBeacon(url, body);
-    } else {
-      await fetch(url, { body, method, keepalive });
+  attachResponseStatus(status: number) {
+    this.logEvents.map((log) => {
+      log.request.satusCode = status;
+    });
+  }
+
+  async sendLogs() {
+    if (!this.logEvents.length) {
+      return;
+    }
+
+    if (!isEnvVarsSet) {
+      // if AXIOM ingesting url is not set, fallback to printing to console
+      // to avoid network errors in development environments
+      this.logEvents.forEach((ev) => prettyPrint(ev));
+      this.logEvents = [];
+      return;
+    }
+
+    const method = 'POST';
+    const keepalive = true;
+    const body = JSON.stringify(this.logEvents);
+    // clear pending logs
+    this.logEvents = [];
+
+    try {
+      if (typeof fetch === 'undefined') {
+        const fetch = await require('whatwg-fetch');
+        await fetch(url, { body, method, keepalive });
+      } else if (isBrowser && navigator.sendBeacon) {
+        navigator.sendBeacon(url, body);
+      } else {
+        await fetch(url, { body, method, keepalive });
+      }
+    } catch (e) {
+      console.error(`Failed to send logs to Axiom: ${e}`);
     }
-  } catch (e) {
-    console.error(`Failed to send logs to Axiom: ${e}`);
   }
+
+  flush = this.sendLogs;
 }
 
 const levelColors = {
@@ -100,13 +124,13 @@ const levelColors = {
   },
 };
 
-export function prettyPrint(level: string, message: string, fields: any = {}) {
-  const hasFields = Object.keys(fields).length > 0;
+export function prettyPrint(ev: LogEvent) {
+  const hasFields = Object.keys(ev.fields).length > 0;
   // check whether pretty print is disabled
   if (isNoPrettyPrint) {
-    let msg = `${level} - ${message}`;
+    let msg = `${ev.level} - ${ev.message}`;
     if (hasFields) {
-      msg += ' ' + JSON.stringify(fields);
+      msg += ' ' + JSON.stringify(ev.fields);
     }
     console.log(msg);
     return;
@@ -116,19 +140,24 @@ export function prettyPrint(level: string, message: string, fields: any = {}) {
   // object as normal text, it loses all the functionality the browser gives for viewing
   // objects in the console, such as expanding and collapsing the object.
   let msgString = '';
-  let args = [level, message];
+  let args: any[] = [ev.level, ev.message];
 
   if (isBrowser) {
     msgString = '%c%s - %s';
-    args = [`color: ${levelColors[level].browser};`, ...args];
+    args = [`color: ${levelColors[ev.level].browser};`, ...args];
   } else {
-    msgString = `\x1b[${levelColors[level].terminal}m%s\x1b[0m - %s`;
+    msgString = `\x1b[${levelColors[ev.level].terminal}m%s\x1b[0m - %s`;
   }
   // we check if the fields object is not empty, otherwise its printed as <empty string>
   // or just "".
   if (hasFields) {
     msgString += ' %o';
-    args.push(fields);
+    args.push(ev.fields);
+  }
+
+  if (ev.request) {
+    msgString += ' %o';
+    args.push(ev.request);
   }
 
   console.log.apply(console, [msgString, ...args]);

src/webVitals.ts

@@ -1,8 +1,6 @@
 import { NextWebVitalsMetric } from 'next/app';
 import { isBrowser, proxyPath, isEnvVarsSet, throttle } from './shared';
 
-export { log } from './logger';
-
 const url = `${proxyPath}/web-vitals`;
 
 export declare type WebVitalsMetric = NextWebVitalsMetric & { route: string };

src/withAxiom.ts

@@ -1,8 +1,8 @@
-import { NextConfig, NextApiHandler, NextApiResponse } from 'next';
+import { NextConfig, NextApiHandler, NextApiResponse, NextApiRequest } from 'next';
 import { proxyPath, EndpointType, getIngestURL } from './shared';
 import { NextFetchEvent, NextMiddleware, NextRequest } from 'next/server';
 
-import { log, Logger } from './logger';
+import { Logger } from './logger';
 import { NextMiddlewareResult } from 'next/dist/server/web/types';
 
 declare global {
@@ -18,6 +18,7 @@ function withAxiomNextConfig(nextConfig: NextConfig): NextConfig {
       const webVitalsEndpoint = getIngestURL(EndpointType.webVitals);
       const logsEndpoint = getIngestURL(EndpointType.logs);
       if (!webVitalsEndpoint && !logsEndpoint) {
+        const log = new Logger();
         log.warn(
           'axiom: Envvars not detected. If this is production please see https://github.com/axiomhq/next-axiom for help'
         );
@@ -52,14 +53,14 @@ function withAxiomNextConfig(nextConfig: NextConfig): NextConfig {
 // Sending logs after res.{json,send,end} is very unreliable.
 // This function overwrites these functions and makes sure logs are sent out
 // before the response is sent.
-function interceptNextApiResponse(res: NextApiResponse): [NextApiResponse, Promise<void>[]] {
+function interceptNextApiResponse(req: AxiomAPIRequest, res: NextApiResponse): [NextApiResponse, Promise<void>[]] {
   const allPromises: Promise<void>[] = [];
 
   const resSend = res.send;
   res.send = (body: any) => {
     allPromises.push(
       (async () => {
-        await log.flush();
+        await req.log.flush();
         resSend(body);
       })()
     );
@@ -69,7 +70,7 @@ function interceptNextApiResponse(res: NextApiResponse): [NextApiResponse, Promi
   res.json = (json: any) => {
     allPromises.push(
       (async () => {
-        await log.flush();
+        await req.log.flush();
         resJson(json);
       })()
     );
@@ -79,7 +80,7 @@ function interceptNextApiResponse(res: NextApiResponse): [NextApiResponse, Promi
   res.end = (cb?: () => undefined): NextApiResponse => {
     allPromises.push(
       (async () => {
-        await log.flush();
+        await req.log.flush();
         resEnd(cb);
       })()
     );
@@ -89,17 +90,26 @@ function interceptNextApiResponse(res: NextApiResponse): [NextApiResponse, Promi
   return [res, allPromises];
 }
 
+export type AxiomAPIRequest = NextApiRequest & { log: Logger };
+export type AxiomApiHandler = (
+  request: AxiomAPIRequest,
+  response: NextApiResponse
+) => NextApiHandler | Promise<NextApiHandler> | Promise<void>;
+
 function withAxiomNextApiHandler(handler: NextApiHandler): NextApiHandler {
   return async (req, res) => {
-    const [wrappedRes, allPromises] = interceptNextApiResponse(res);
+    const logger = new Logger();
+    const axiomRequest = req as AxiomAPIRequest;
+    axiomRequest.log = logger;
+    const [wrappedRes, allPromises] = interceptNextApiResponse(axiomRequest, res);
 
     try {
-      await handler(req, wrappedRes);
-      await log.flush();
+      await handler(axiomRequest, wrappedRes);
+      await logger.flush();
       await Promise.all(allPromises);
     } catch (error) {
-      log.error('Error in API handler', { error });
-      await log.flush();
+      logger.error('Error in API handler', { error });
+      await logger.flush();
       await Promise.all(allPromises);
       throw error;
     }
@@ -114,32 +124,19 @@ export type AxiomMiddleware = (
 
 function withAxiomNextEdgeFunction(handler: NextMiddleware): NextMiddleware {
   return async (req, ev) => {
-    const report = {
-      request: {
-        ip: req.ip,
-        region: req.geo?.region,
-        host: req.nextUrl.host,
-        method: req.method,
-        path: req.nextUrl.pathname,
-        scheme: req.nextUrl.protocol.replace(':', ''),
-        statusCode: 0,
-        userAgent: req.headers.get('user-agent'),
-      },
-    };
-    const logger = log.with({
-      request: report.request,
-    });
+    const logger = new Logger({}, req, true);
     const axiomRequest = req as AxiomRequest;
     axiomRequest.log = logger;
 
     try {
       const res = await handler(axiomRequest, ev);
-      report.request.statusCode = res?.status || 0;
-      ev.waitUntil(log.flush());
+      logger.attachResponseStatus(res?.status || 0);
+      ev.waitUntil(logger.flush());
       return res;
     } catch (error) {
-      log.error('Error in edge function', { error });
-      ev.waitUntil(log.flush());
+      logger.error('Error in edge function', { error });
+      logger.attachResponseStatus(500);
+      ev.waitUntil(logger.flush());
       throw error;
     }
   };