Skip to content

Commit fd00b24

Browse files
yash-pouranikyash-pouranik
committed
feat: implement aggregation api with rls support
1 parent 009dc33 commit fd00b24

5 files changed

Lines changed: 227 additions & 2 deletions

File tree

‎apps/public-api/src/__tests__/aggregation.test.js‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
'use strict';
2+
3+
const mockAggregate = jest.fn();
4+
5+
jest.mock('@urbackend/common', () => ({
6+
sanitize: (v) => v,
7+
Project: {},
8+
getConnection: jest.fn().mockResolvedValue({}),
9+
getCompiledModel: jest.fn(() => ({
10+
aggregate: mockAggregate,
11+
})),
12+
QueryEngine: jest.fn(),
13+
validateData: jest.fn(),
14+
validateUpdateData: jest.fn(),
15+
aggregateSchema: require('../../../../packages/common/src/utils/input.validation').aggregateSchema,
16+
}));
17+
18+
const { aggregateData } = require('../controllers/data.controller');
19+
20+
function makeReq(overrides = {}) {
21+
return {
22+
params: { collectionName: 'posts' },
23+
project: {
24+
_id: 'proj_1',
25+
resources: { db: { isExternal: false } },
26+
collections: [{ name: 'posts', model: [] }],
27+
},
28+
body: { pipeline: [{ $group: { _id: '$status', count: { $sum: 1 } } }] },
29+
rlsFilter: {},
30+
...overrides,
31+
};
32+
}
33+
34+
function makeRes() {
35+
const res = {
36+
statusCode: null,
37+
body: null,
38+
status: jest.fn().mockReturnThis(),
39+
json: jest.fn().mockReturnThis(),
40+
};
41+
res.status.mockImplementation((code) => {
42+
res.statusCode = code;
43+
return res;
44+
});
45+
res.json.mockImplementation((data) => {
46+
res.body = data;
47+
return res;
48+
});
49+
return res;
50+
}
51+
52+
describe('aggregateData controller', () => {
53+
beforeEach(() => {
54+
jest.clearAllMocks();
55+
mockAggregate.mockResolvedValue([{ _id: 'published', count: 2 }]);
56+
});
57+
58+
test('executes a valid aggregation pipeline', async () => {
59+
const req = makeReq();
60+
const res = makeRes();
61+
62+
await aggregateData(req, res);
63+
64+
expect(mockAggregate).toHaveBeenCalledWith([
65+
{ $group: { _id: '$status', count: { $sum: 1 } } },
66+
]);
67+
expect(res.statusCode).toBe(200);
68+
expect(res.body).toEqual({
69+
success: true,
70+
data: [{ _id: 'published', count: 2 }],
71+
message: 'Aggregation executed successfully.',
72+
});
73+
});
74+
75+
test('prepends the RLS match stage when req.rlsFilter is set', async () => {
76+
const req = makeReq({
77+
rlsFilter: { userId: 'user_1' },
78+
body: {
79+
pipeline: [
80+
{ $match: { status: 'published' } },
81+
{ $sort: { createdAt: -1 } },
82+
],
83+
},
84+
});
85+
const res = makeRes();
86+
87+
await aggregateData(req, res);
88+
89+
expect(mockAggregate).toHaveBeenCalledWith([
90+
{ $match: { userId: 'user_1' } },
91+
{ $match: { status: 'published' } },
92+
{ $sort: { createdAt: -1 } },
93+
]);
94+
expect(res.statusCode).toBe(200);
95+
});
96+
97+
test('blocks write-capable aggregation stages', async () => {
98+
const req = makeReq({
99+
body: { pipeline: [{ $match: { status: 'published' } }, { $out: 'archive' }] },
100+
});
101+
const res = makeRes();
102+
103+
await aggregateData(req, res);
104+
105+
expect(mockAggregate).not.toHaveBeenCalled();
106+
expect(res.statusCode).toBe(400);
107+
expect(res.body).toEqual({
108+
success: false,
109+
data: {},
110+
message: 'Aggregation pipeline contains blocked stage.',
111+
});
112+
});
113+
114+
test('rejects invalid pipeline payloads', async () => {
115+
const req = makeReq({
116+
body: { pipeline: { $group: { _id: '$status' } } },
117+
});
118+
const res = makeRes();
119+
120+
await aggregateData(req, res);
121+
122+
expect(mockAggregate).not.toHaveBeenCalled();
123+
expect(res.statusCode).toBe(400);
124+
expect(res.body).toEqual({
125+
success: false,
126+
data: {},
127+
message: 'Invalid input: expected array, received object',
128+
});
129+
});
130+
});

‎apps/public-api/src/controllers/data.controller.js‎

Lines changed: 80 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@ const { Project } = require("@urbackend/common");
44
const { getConnection } = require("@urbackend/common");
55
const { getCompiledModel } = require("@urbackend/common");
66
const { QueryEngine } = require("@urbackend/common");
7-
const { validateData, validateUpdateData } = require("@urbackend/common");
7+
const { validateData, validateUpdateData, aggregateSchema } = require("@urbackend/common");
88
const { performance } = require('perf_hooks');
99
const { dispatchWebhooks } = require('../utils/webhookDispatcher');
10+
const { z } = require("zod");
1011

1112
const isDebug = process.env.DEBUG === 'true';
1213

@@ -17,6 +18,14 @@ const isDuplicateKeyError = (err) => {
1718
return err && err.code === 11000;
1819
};
1920

21+
const BLOCKED_AGGREGATION_STAGES = new Set(["$out", "$merge"]);
22+
23+
const containsBlockedAggregationStage = (pipeline = []) => {
24+
return pipeline.some((stage) =>
25+
Object.keys(stage || {}).some((key) => BLOCKED_AGGREGATION_STAGES.has(key)),
26+
);
27+
};
28+
2029
// INSERT DATA
2130
module.exports.insertData = async (req, res) => {
2231
try {
@@ -181,6 +190,76 @@ module.exports.getSingleDoc = async (req, res) => {
181190
}
182191
};
183192

193+
// AGGREGATE DATA
194+
module.exports.aggregateData = async (req, res) => {
195+
try {
196+
let start;
197+
if (isDebug) start = performance.now();
198+
const { collectionName } = req.params;
199+
const project = req.project;
200+
201+
const collectionConfig = project.collections.find(
202+
(c) => c.name === collectionName,
203+
);
204+
if (!collectionConfig) {
205+
return res.status(404).json({
206+
success: false,
207+
data: {},
208+
message: "Collection not found",
209+
});
210+
}
211+
212+
const { pipeline } = aggregateSchema.parse(req.body || {});
213+
214+
if (containsBlockedAggregationStage(pipeline)) {
215+
return res.status(400).json({
216+
success: false,
217+
data: {},
218+
message: "Aggregation pipeline contains blocked stage.",
219+
});
220+
}
221+
222+
const connection = await getConnection(project._id);
223+
const Model = getCompiledModel(
224+
connection,
225+
collectionConfig,
226+
project._id,
227+
project.resources.db.isExternal,
228+
);
229+
230+
const baseFilter =
231+
req.rlsFilter && typeof req.rlsFilter === "object" ? req.rlsFilter : {};
232+
const effectivePipeline = Object.keys(baseFilter).length > 0
233+
? [{ $match: baseFilter }, ...pipeline]
234+
: pipeline;
235+
236+
const data = await Model.aggregate(effectivePipeline);
237+
238+
if (isDebug) console.log(`[DEBUG] aggregate took ${(performance.now() - start).toFixed(2)}ms`);
239+
return res.status(200).json({
240+
success: true,
241+
data,
242+
message: "Aggregation executed successfully.",
243+
});
244+
} catch (err) {
245+
console.error(err);
246+
247+
if (err instanceof z.ZodError) {
248+
return res.status(400).json({
249+
success: false,
250+
data: {},
251+
message: err.issues?.[0]?.message || "Invalid aggregation payload.",
252+
});
253+
}
254+
255+
return res.status(500).json({
256+
success: false,
257+
data: {},
258+
message: err.message || "Failed to execute aggregation.",
259+
});
260+
}
261+
};
262+
184263
// UPDATE DATA
185264
module.exports.updateSingleData = async (req, res) => {
186265
try {

‎apps/public-api/src/routes/data.js‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ const authorizeWriteOperation = require('../middlewares/authorizeWriteOperation'
66
const authorizeReadOperation = require('../middlewares/authorizeReadOperation');
77
const projectRateLimiter = require('../middlewares/projectRateLimiter');
88
const blockUsersCollectionDataAccess = require('../middlewares/blockUsersCollectionDataAccess');
9-
const { insertData, getAllData, getSingleDoc, updateSingleData, deleteSingleDoc } = require("../controllers/data.controller")
9+
const { insertData, getAllData, getSingleDoc, updateSingleData, deleteSingleDoc, aggregateData } = require("../controllers/data.controller")
1010

1111

1212
// POST REQ TO INSERT DATA
@@ -16,6 +16,9 @@ router.post('/:collectionName', verifyApiKey, blockUsersCollectionDataAccess, re
1616
// GET REQ ALL DATA
1717
router.get('/:collectionName', verifyApiKey, blockUsersCollectionDataAccess, resolvePublicAuthContext, projectRateLimiter, authorizeReadOperation, getAllData);
1818

19+
// POST REQ AGGREGATION DATA
20+
router.post('/:collectionName/aggregate', verifyApiKey, blockUsersCollectionDataAccess, resolvePublicAuthContext, projectRateLimiter, authorizeReadOperation, aggregateData);
21+
1922

2023
// GET REQ SINGLE DATA
2124
router.get('/:collectionName/:id', verifyApiKey, blockUsersCollectionDataAccess, resolvePublicAuthContext, projectRateLimiter, authorizeReadOperation, getSingleDoc);

‎packages/common/src/index.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ const {
5454
createProjectSchema,
5555
createCollectionSchema,
5656
createSchemaApiKeySchema,
57+
aggregateSchema,
5758
sanitize,
5859
userSignupSchema,
5960
updateExternalConfigSchema,
@@ -113,6 +114,7 @@ module.exports = {
113114
createProjectSchema,
114115
createCollectionSchema,
115116
createSchemaApiKeySchema,
117+
aggregateSchema,
116118
sanitize,
117119
updateExternalConfigSchema,
118120
updateAuthProvidersSchema,

‎packages/common/src/utils/input.validation.js‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,17 @@ module.exports.createSchemaApiKeySchema = z.object({
273273
fields: z.array(buildApiFieldSchemaZod(1)).optional(),
274274
});
275275

276+
module.exports.aggregateSchema = z.object({
277+
pipeline: z
278+
.array(
279+
z.record(z.string(), z.unknown()).refine(
280+
(stage) => Object.keys(stage).length > 0,
281+
{ message: "Each aggregation stage must be a non-empty object." },
282+
),
283+
)
284+
.min(1, "Pipeline must contain at least one stage."),
285+
});
286+
276287
module.exports.sanitize = (obj) => {
277288
const clean = {};
278289
for (const key in obj) {

0 commit comments

Comments
 (0)