Add collect limits.

lib/cache.ts

@@ -1,6 +1,7 @@
-import { User, Website } from '@prisma/client';
+import { User, Website, Team } from '@prisma/client';
 import redis from '@umami/redis-client';
-import { getSession, getUser, getWebsite } from '../queries';
+import { lightFormat, startOfMonth } from 'date-fns';
+import { getAllWebsitesByUser, getSession, getUser, getViewTotals, getWebsite } from '../queries';
 
 const DELETED = 'DELETED';
 
@@ -32,7 +33,11 @@ async function deleteObject(key, soft = false) {
   return soft ? redis.set(key, DELETED) : redis.del(key);
 }
 
-async function fetchWebsite(id): Promise<Website> {
+async function fetchWebsite(id): Promise<
+  Website & {
+    team?: Team;
+  }
+> {
   return fetchObject(`website:${id}`, () => getWebsite({ id }));
 }
 
@@ -58,6 +63,31 @@ async function storeUser(data) {
   return storeObject(key, data);
 }
 
+async function fetchCollectLimit(userId): Promise<Number> {
+  const monthDate = startOfMonth(new Date());
+  const monthKey = lightFormat(monthDate, 'yyyy-MM');
+
+  return fetchObject(`collectLimit:${userId}:${monthKey}`, async () => {
+    const websiteIds = (await getAllWebsitesByUser(userId)).map(a => a.id);
+
+    return getViewTotals(websiteIds, monthDate).then(data => data.views);
+  });
+}
+
+async function deleteCollectLimit(userId): Promise<Number> {
+  const monthDate = startOfMonth(new Date());
+  const monthKey = lightFormat(monthDate, 'yyyy-MM');
+
+  return deleteObject(`collectLimit:${userId}:${monthKey}`);
+}
+
+async function incrementCollectLimit(userId): Promise<Number> {
+  const monthDate = startOfMonth(new Date());
+  const monthKey = lightFormat(monthDate, 'yyyy-MM');
+
+  return redis.incr(`collectLimit:${userId}:${monthKey}`);
+}
+
 async function deleteUser(id) {
   return deleteObject(`user:${id}`);
 }
@@ -87,5 +117,8 @@ export default {
   fetchSession,
   storeSession,
   deleteSession,
+  fetchCollectLimit,
+  incrementCollectLimit,
+  deleteCollectLimit,
   enabled: redis.enabled,
 };

lib/clickhouse.js

@@ -177,6 +177,10 @@ function formatQuery(str, params = []) {
       replace = `'${replace}'`;
     }
 
+    if (param instanceof Date) {
+      replace = getDateFormat(param);
+    }
+
     formattedString = formattedString.replace(`$${i + 1}`, replace);
   });
 

lib/middleware.ts

@@ -13,7 +13,7 @@ const log = debug('umami:middleware');
 
 export const useCors = createMiddleware(cors());
 
-export const useSession = createMiddleware(async (req, res, next) => {
+export const useSession = createMiddleware(async (req: any, res, next) => {
   const session = await findSession(req);
 
   if (!session) {
@@ -21,7 +21,7 @@ export const useSession = createMiddleware(async (req, res, next) => {
     return badRequest(res);
   }
 
-  (req as any).session = session;
+  req.session = session;
   next();
 });
 

lib/prisma.ts

@@ -1,7 +1,7 @@
 import prisma from '@umami/prisma-client';
-import moment from 'moment-timezone';
-import { MYSQL, POSTGRESQL, getDatabaseType } from 'lib/db';
 import { FILTER_IGNORED } from 'lib/constants';
+import { getDatabaseType, MYSQL, POSTGRESQL } from 'lib/db';
+import moment from 'moment-timezone';
 
 const MYSQL_DATE_FORMATS = {
   minute: '%Y-%m-%d %H:%i:00',

lib/session.ts

@@ -1,38 +1,41 @@
-import { parseToken } from 'next-basics';
-import { validate } from 'uuid';
-import { secret, uuid } from 'lib/crypto';
+import { Session, Team, Website } from '@prisma/client';
 import cache from 'lib/cache';
 import clickhouse from 'lib/clickhouse';
+import { secret, uuid } from 'lib/crypto';
 import { getClientInfo, getJsonBody } from 'lib/detect';
+import { parseToken } from 'next-basics';
 import { createSession, getSession, getWebsite } from 'queries';
+import { validate } from 'uuid';
 
-export async function findSession(req) {
+export async function findSession(req): Promise<{
+  error?: {
+    status: number;
+    message: string;
+  };
+  session?: {
+    id: string;
+    websiteId: string;
+    hostname: string;
+    browser: string;
+    os: string;
+    device: string;
+    screen: string;
+    language: string;
+    country: string;
+  };
+  website?: Website & { team?: Team };
+}> {
   const { payload } = getJsonBody(req);
 
   if (!payload) {
     return null;
   }
 
-  // Check if cache token is passed
-  const cacheToken = req.headers['x-umami-cache'];
-
-  if (cacheToken) {
-    const result = await parseToken(cacheToken, secret());
-
-    if (result) {
-      return result;
-    }
-  }
-
   // Verify payload
   const { website: websiteId, hostname, screen, language } = payload;
 
-  if (!validate(websiteId)) {
-    return null;
-  }
-
   // Find website
-  let website;
+  let website: Website & { team?: Team } = null;
 
   if (cache.enabled) {
     website = await cache.fetchWebsite(websiteId);
@@ -44,26 +47,44 @@ export async function findSession(req) {
     throw new Error(`Website not found: ${websiteId}`);
   }
 
+  // Check if cache token is passed
+  const cacheToken = req.headers['x-umami-cache'];
+
+  if (cacheToken) {
+    const result = await parseToken(cacheToken, secret());
+
+    if (result) {
+      return { session: result, website };
+    }
+  }
+
+  if (!validate(websiteId)) {
+    return null;
+  }
+
   const { userAgent, browser, os, ip, country, device } = await getClientInfo(req, payload);
   const sessionId = uuid(websiteId, hostname, ip, userAgent);
 
   // Clickhouse does not require session lookup
   if (clickhouse.enabled) {
     return {
-      id: sessionId,
-      websiteId,
-      hostname,
-      browser,
-      os,
-      device,
-      screen,
-      language,
-      country,
+      session: {
+        id: sessionId,
+        websiteId,
+        hostname,
+        browser,
+        os,
+        device,
+        screen,
+        language,
+        country,
+      },
+      website,
     };
   }
 
   // Find session
-  let session;
+  let session: Session;
 
   if (cache.enabled) {
     session = await cache.fetchSession(sessionId);
@@ -85,12 +106,12 @@ export async function findSession(req) {
         language,
         country,
       });
-    } catch (e) {
+    } catch (e: any) {
       if (!e.message.toLowerCase().includes('unique constraint')) {
         throw e;
       }
     }
   }
 
-  return session;
+  return { session, website };
 }

package.json

@@ -61,7 +61,7 @@
     "@fontsource/inter": "4.5.7",
     "@prisma/client": "4.8.0",
     "@tanstack/react-query": "^4.16.1",
-    "@umami/prisma-client": "^0.1.0",
+    "@umami/prisma-client": "^0.2.0",
     "@umami/redis-client": "^0.1.0",
     "chalk": "^4.1.1",
     "chart.js": "^2.9.4",

pages/api/collect.ts

@@ -1,24 +1,40 @@
 const { Resolver } = require('dns').promises;
 import isbot from 'isbot';
 import ipaddr from 'ipaddr.js';
-import { createToken, unauthorized, send, badRequest, forbidden } from 'next-basics';
+import {
+  createToken,
+  unauthorized,
+  send,
+  badRequest,
+  forbidden,
+  tooManyRequest,
+} from 'next-basics';
 import { savePageView, saveEvent } from 'queries';
 import { useCors, useSession } from 'lib/middleware';
 import { getJsonBody, getIpAddress } from 'lib/detect';
 import { secret } from 'lib/crypto';
 import { NextApiRequest, NextApiResponse } from 'next';
+import cache from 'lib/cache';
+import { Team, Website } from '@prisma/client';
 
 export interface NextApiRequestCollect extends NextApiRequest {
   session: {
-    id: string;
-    websiteId: string;
-    hostname: string;
-    browser: string;
-    os: string;
-    device: string;
-    screen: string;
-    language: string;
-    country: string;
+    error?: {
+      status: number;
+      message: string;
+    };
+    session?: {
+      id: string;
+      websiteId: string;
+      hostname: string;
+      browser: string;
+      os: string;
+      device: string;
+      screen: string;
+      language: string;
+      country: string;
+    };
+    website?: Website & { team?: Team };
   };
 }
 
@@ -88,7 +104,21 @@ export default async (req: NextApiRequestCollect, res: NextApiResponse) => {
 
   await useSession(req, res);
 
-  const session = req.session;
+  const { session, website } = req.session;
+
+  // Check collection limit
+  if (process.env.ENABLE_COLLECT_LIMIT) {
+    const userId = website.userId ? website.userId : website.team.userId;
+
+    const limit = await cache.fetchCollectLimit(userId);
+
+    // To-do: Need to implement logic to find user-specific limit. Defaulted to 10k.
+    if (limit > 10000) {
+      return tooManyRequest(res, 'Collect currently exceeds monthly limit of 10000.');
+    }
+
+    await cache.incrementCollectLimit(userId);
+  }
 
   if (process.env.REMOVE_TRAILING_SLASH) {
     url = url.replace(/\/$/, '');

queries/admin/website.ts

@@ -1,16 +1,45 @@
-import { Prisma, Website } from '@prisma/client';
+import { Prisma, Team, Website } from '@prisma/client';
 import cache from 'lib/cache';
 import prisma from 'lib/prisma';
 import { runQuery, CLICKHOUSE, PRISMA } from 'lib/db';
 
-export async function getWebsite(where: Prisma.WebsiteWhereUniqueInput): Promise<Website> {
+export async function getWebsite(where: Prisma.WebsiteWhereUniqueInput): Promise<
+  Website & {
+    team?: Team;
+  }
+> {
   return prisma.client.website.findUnique({
     where,
+    include: {
+      team: true,
+    },
   });
 }
 
-export async function getWebsites(): Promise<Website[]> {
+export async function getWebsites(
+  where: Prisma.WebsiteFindManyArgs,
+  showDeleted = false,
+): Promise<Website[]> {
   return prisma.client.website.findMany({
+    where: { ...where, deletedAt: showDeleted ? { not: null } : null },
+    orderBy: {
+      name: 'asc',
+    },
+  });
+}
+
+export async function getAllWebsitesByUser(userId): Promise<Website[]> {
+  return prisma.client.website.findMany({
+    where: {
+      OR: [
+        { userId },
+        {
+          team: {
+            userId,
+          },
+        },
+      ],
+    },
     orderBy: {
       name: 'asc',
     },

queries/analytics/session/getSession.ts

@@ -1,7 +1,7 @@
 import clickhouse from 'lib/clickhouse';
 import { CLICKHOUSE, PRISMA, runQuery } from 'lib/db';
 import prisma from 'lib/prisma';
-import { Prisma } from '@prisma/client';
+import { Prisma, Session } from '@prisma/client';
 
 export async function getSession(args: { id: string }) {
   return runQuery({
@@ -10,7 +10,7 @@ export async function getSession(args: { id: string }) {
   });
 }
 
-async function relationalQuery(where: Prisma.SessionWhereUniqueInput) {
+async function relationalQuery(where: Prisma.SessionWhereUniqueInput): Promise<Session> {
   return prisma.client.session.findUnique({
     where,
   });

queries/analytics/stats/getViewTotals.ts

@@ -0,0 +1,38 @@
+import clickhouse from 'lib/clickhouse';
+import { CLICKHOUSE, PRISMA, runQuery } from 'lib/db';
+import prisma from 'lib/prisma';
+
+export async function getViewTotals(...args: [websites: string[], startDate: Date]) {
+  return runQuery({
+    [PRISMA]: () => relationalQuery(...args),
+    [CLICKHOUSE]: () => clickhouseQuery(...args),
+  });
+}
+
+async function relationalQuery(websites: string[], startDate: Date) {
+  return prisma.client.websiteEvent.count({
+    where: {
+      websiteId: {
+        in: websites,
+      },
+      createdAt: {
+        gte: startDate,
+      },
+    },
+  });
+}
+
+async function clickhouseQuery(websiteIds: string[], startDate: Date) {
+  const { rawQuery, getDateFormat, getCommaSeparatedStringFormat, findFirst } = clickhouse;
+
+  return rawQuery(
+    `
+    select        
+      count(*) as views
+    from event
+    where
+      website_id in (${getCommaSeparatedStringFormat(websiteIds)})
+      and created_at >= ${getDateFormat(startDate)}
+     `,
+  ).then(data => findFirst(data));
+}

queries/analytics/stats/getWebsiteStats.ts

@@ -18,7 +18,7 @@ async function relationalQuery(
 ) {
   const { startDate, endDate, filters = {} } = data;
   const { getDateQuery, getTimestampInterval, parseFilters, rawQuery } = prisma;
-  const params = [startDate, endDate];
+  const params: any = [startDate, endDate];
   const { filterQuery, joinSession } = parseFilters(filters, params);
 
   return rawQuery(

queries/index.js

@@ -16,4 +16,5 @@ export * from './analytics/session/getSessionMetrics';
 export * from './analytics/session/getSessions';
 export * from './analytics/stats/getActiveVisitors';
 export * from './analytics/stats/getRealtimeData';
+export * from './analytics/stats/getViewTotals';
 export * from './analytics/stats/getWebsiteStats';

yarn.lock

@@ -2290,10 +2290,10 @@
     "@typescript-eslint/types" "5.45.0"
     eslint-visitor-keys "^3.3.0"
 
-"@umami/prisma-client@^0.1.0":
-  version "0.1.0"
-  resolved "https://registry.yarnpkg.com/@umami/prisma-client/-/prisma-client-0.1.0.tgz#4ecc17b1998cb738b8f2027b80b003864beef6a4"
-  integrity sha512-2qTXN0dtMUT+HVhz37eVC02lDSwz9TGUGHRRO0LPRHdfJPfSkUF0D8pgksX7ETy3IZoIQP1h45tiXsM0pKBMZA==
+"@umami/prisma-client@^0.2.0":
+  version "0.2.0"
+  resolved "https://registry.yarnpkg.com/@umami/prisma-client/-/prisma-client-0.2.0.tgz#b9de1f28be67ccfb9e2544f23c69c392c5b26ea7"
+  integrity sha512-+27dd4DLl8SvbbIYG1mdm6pIZd+UzQI7eZGNjQ9ONeWO0jr+/wiVnPIXUzd8w4R/OoM4ChpI3mBZPqcWa5MAOw==
   dependencies:
     debug "^4.3.4"