diff --git a/lib/cache.ts b/lib/cache.ts index 2aad7ed8..d1a56b5e 100644 --- a/lib/cache.ts +++ b/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 { +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 { + 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 { + const monthDate = startOfMonth(new Date()); + const monthKey = lightFormat(monthDate, 'yyyy-MM'); + + return deleteObject(`collectLimit:${userId}:${monthKey}`); +} + +async function incrementCollectLimit(userId): Promise { + 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, }; diff --git a/lib/clickhouse.js b/lib/clickhouse.js index 46cdaabc..b986cc39 100644 --- a/lib/clickhouse.js +++ b/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); }); diff --git a/lib/middleware.ts b/lib/middleware.ts index ce6aea25..e8f56b1a 100644 --- a/lib/middleware.ts +++ b/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(); }); diff --git a/lib/prisma.ts b/lib/prisma.ts index 4f63bf7d..417af00e 100644 --- a/lib/prisma.ts +++ b/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', diff --git a/lib/session.js b/lib/session.ts similarity index 69% rename from lib/session.js rename to lib/session.ts index f6698480..08a54fc8 100644 --- a/lib/session.js +++ b/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 }; } diff --git a/package.json b/package.json index f28eb038..e878a1eb 100644 --- a/package.json +++ b/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", diff --git a/pages/api/collect.ts b/pages/api/collect.ts index 68c4afb2..2f70bc08 100644 --- a/pages/api/collect.ts +++ b/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(/\/$/, ''); diff --git a/queries/admin/website.ts b/queries/admin/website.ts index cf2425a3..e820bc84 100644 --- a/queries/admin/website.ts +++ b/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 { +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 { +export async function getWebsites( + where: Prisma.WebsiteFindManyArgs, + showDeleted = false, +): Promise { return prisma.client.website.findMany({ + where: { ...where, deletedAt: showDeleted ? { not: null } : null }, + orderBy: { + name: 'asc', + }, + }); +} + +export async function getAllWebsitesByUser(userId): Promise { + return prisma.client.website.findMany({ + where: { + OR: [ + { userId }, + { + team: { + userId, + }, + }, + ], + }, orderBy: { name: 'asc', }, diff --git a/queries/analytics/session/getSession.ts b/queries/analytics/session/getSession.ts index 19875117..f52e92ef 100644 --- a/queries/analytics/session/getSession.ts +++ b/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 { return prisma.client.session.findUnique({ where, }); diff --git a/queries/analytics/stats/getViewTotals.ts b/queries/analytics/stats/getViewTotals.ts new file mode 100644 index 00000000..7ad5b3b2 --- /dev/null +++ b/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)); +} diff --git a/queries/analytics/stats/getWebsiteStats.ts b/queries/analytics/stats/getWebsiteStats.ts index bf5cdd96..1a1d289a 100644 --- a/queries/analytics/stats/getWebsiteStats.ts +++ b/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( diff --git a/queries/index.js b/queries/index.js index 2d1931ee..e30d3f4b 100644 --- a/queries/index.js +++ b/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'; diff --git a/yarn.lock b/yarn.lock index 450c4744..83e65ace 100644 --- a/yarn.lock +++ b/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"