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 };
 }