chore(api): allow device deletion

api/src/billing/billing.service.ts

@@ -16,7 +16,6 @@ import { User, UserDocument } from '../users/schemas/user.schema'
 import { CheckoutResponseDTO, PlanDTO } from './billing.dto'
 import { SMSDocument } from '../gateway/schemas/sms.schema'
 import { SMS } from '../gateway/schemas/sms.schema'
-import { Device, DeviceDocument } from '../gateway/schemas/device.schema'
 import { validateEvent } from '@polar-sh/sdk/webhooks'
 import {
   PolarWebhookPayload,
@@ -41,7 +40,6 @@ export class BillingService {
     private subscriptionModel: Model<SubscriptionDocument>,
     @InjectModel(User.name) private userModel: Model<UserDocument>,
     @InjectModel(SMS.name) private smsModel: Model<SMSDocument>,
-    @InjectModel(Device.name) private deviceModel: Model<DeviceDocument>,
     @InjectModel(PolarWebhookPayload.name)
     private polarWebhookPayloadModel: Model<PolarWebhookPayloadDocument>,
     @InjectModel(CheckoutSession.name)
@@ -69,17 +67,13 @@ export class BillingService {
       })
       .populate('plan')
 
-    // Get user's devices and usage data
-    const userDevices = await this.deviceModel.find({ user: user._id }, '_id')
-    const deviceIds = userDevices.map((d) => d._id)
-
     const processedSmsToday = await this.smsModel.countDocuments({
-      device: { $in: deviceIds },
+      user: user._id,
       createdAt: { $gte: new Date(new Date().setHours(0, 0, 0, 0)) },
     })
 
     const processedSmsLastMonth = await this.smsModel.countDocuments({
-      device: { $in: deviceIds },
+      user: user._id,
       createdAt: {
         $gte: new Date(new Date().setMonth(new Date().getMonth() - 1)),
       },
@@ -568,16 +562,12 @@ export class BillingService {
       let hasReachedLimit = false
       let message = ''
 
-      // Get user's devices and then count SMS
-      const userDevices = await this.deviceModel.find({ user: user._id }, '_id')
-      const deviceIds = userDevices.map((d) => d._id)
-
       const processedSmsToday = await this.smsModel.countDocuments({
-        device: { $in: deviceIds },
+        user: user._id,
         createdAt: { $gte: new Date(new Date().setHours(0, 0, 0, 0)) },
       })
       const processedSmsLastMonth = await this.smsModel.countDocuments({
-        device: { $in: deviceIds },
+        user: user._id,
         createdAt: {
           $gte: new Date(new Date().setMonth(new Date().getMonth() - 1)),
         },
@@ -724,19 +714,13 @@ export class BillingService {
 
     const effectiveLimits = this.getEffectiveLimits(subscription, plan)
 
-    // First get all devices belonging to the user
-    const userDevices = await this.deviceModel
-      .find({ user: new Types.ObjectId(userId) })
-      .select('_id')
-    const deviceIds = userDevices.map((device) => device._id)
-
     const processedSmsToday = await this.smsModel.countDocuments({
-      device: { $in: deviceIds },
+      user: new Types.ObjectId(userId),
       createdAt: { $gte: new Date(new Date().setHours(0, 0, 0, 0)) },
     })
 
     const processedSmsLastMonth = await this.smsModel.countDocuments({
-      device: { $in: deviceIds },
+      user: new Types.ObjectId(userId),
       createdAt: {
         $gte: new Date(new Date().setMonth(new Date().getMonth() - 1)),
       },

api/src/gateway/gateway.module.ts

@@ -1,6 +1,7 @@
 import { forwardRef, Module } from '@nestjs/common'
 import { MongooseModule } from '@nestjs/mongoose'
 import { Device, DeviceSchema } from './schemas/device.schema'
+import { DeviceTombstone, DeviceTombstoneSchema } from './schemas/device-tombstone.schema'
 import { GatewayController } from './gateway.controller'
 import { GatewayService } from './gateway.service'
 import { AuthModule } from '../auth/auth.module'
@@ -23,6 +24,10 @@ import { HeartbeatCheckTask } from './tasks/heartbeat-check.task'
         name: Device.name,
         schema: DeviceSchema,
       },
+      {
+        name: DeviceTombstone.name,
+        schema: DeviceTombstoneSchema,
+      },
       {
         name: SMS.name,
         schema: SMSSchema,

api/src/gateway/gateway.service.spec.ts

@@ -3,6 +3,7 @@ import { GatewayService } from './gateway.service'
 import { AuthModule } from '../auth/auth.module'
 import { getModelToken } from '@nestjs/mongoose'
 import { Device, DeviceDocument } from './schemas/device.schema'
+import { DeviceTombstone } from './schemas/device-tombstone.schema'
 import { SMS } from './schemas/sms.schema'
 import { SMSBatch } from './schemas/sms-batch.schema'
 import { AuthService } from '../auth/auth.service'
@@ -29,6 +30,7 @@ jest.mock('firebase-admin', () => ({
 describe('GatewayService', () => {
   let service: GatewayService
   let deviceModel: Model<DeviceDocument>
+  let deviceTombstoneModel: Model<any>
   let smsModel: Model<SMS>
   let smsBatchModel: Model<SMSBatch>
   let authService: AuthService
@@ -60,6 +62,10 @@ describe('GatewayService', () => {
     findByIdAndUpdate: jest.fn(),
   }
 
+  const mockDeviceTombstoneModel = {
+    updateOne: jest.fn(),
+  }
+
   const mockAuthService = {
     getUserApiKeys: jest.fn(),
   }
@@ -85,6 +91,10 @@ describe('GatewayService', () => {
           provide: getModelToken(Device.name),
           useValue: mockDeviceModel,
         },
+        {
+          provide: getModelToken(DeviceTombstone.name),
+          useValue: mockDeviceTombstoneModel,
+        },
         {
           provide: getModelToken(SMS.name),
           useValue: mockSmsModel,
@@ -115,6 +125,9 @@ describe('GatewayService', () => {
 
     service = module.get<GatewayService>(GatewayService)
     deviceModel = module.get<Model<DeviceDocument>>(getModelToken(Device.name))
+    deviceTombstoneModel = module.get<Model<any>>(
+      getModelToken(DeviceTombstone.name),
+    )
     smsModel = module.get<Model<SMS>>(getModelToken(SMS.name))
     smsBatchModel = module.get<Model<SMSBatch>>(getModelToken(SMSBatch.name))
     authService = module.get<AuthService>(AuthService)
@@ -299,16 +312,18 @@ describe('GatewayService', () => {
   })
 
   describe('deleteDevice', () => {
-    const mockDeviceId = 'device123'
+    const mockDeviceId = '507f1f77bcf86cd799439011'
     const mockDevice = { _id: mockDeviceId, model: 'Pixel 6' }
 
-    it('should return empty object when device exists', async () => {
+    it('should tombstone and delete when device exists', async () => {
       mockDeviceModel.findById.mockResolvedValue(mockDevice)
 
       const result = await service.deleteDevice(mockDeviceId)
 
       expect(mockDeviceModel.findById).toHaveBeenCalledWith(mockDeviceId)
-      expect(result).toEqual({})
+      expect(mockDeviceTombstoneModel.updateOne).toHaveBeenCalled()
+      expect(mockDeviceModel.findByIdAndDelete).toHaveBeenCalledWith(mockDeviceId)
+      expect(result).toEqual({ success: true })
     })
 
     it('should throw an error if device does not exist', async () => {

api/src/gateway/gateway.service.ts

@@ -3,6 +3,7 @@ import { InjectModel } from '@nestjs/mongoose'
 import { Device, DeviceDocument } from './schemas/device.schema'
 import { Model, Types } from 'mongoose'
 import * as firebaseAdmin from 'firebase-admin'
+import { DeviceTombstone, DeviceTombstoneDocument } from './schemas/device-tombstone.schema'
 import {
   ReceivedSMSDTO,
   RegisterDeviceInputDTO,
@@ -28,6 +29,8 @@ import { SmsQueueService } from './queue/sms-queue.service'
 export class GatewayService {
   constructor(
     @InjectModel(Device.name) private deviceModel: Model<DeviceDocument>,
+    @InjectModel(DeviceTombstone.name)
+    private deviceTombstoneModel: Model<DeviceTombstoneDocument>,
     @InjectModel(SMS.name) private smsModel: Model<SMS>,
     @InjectModel(SMSBatch.name) private smsBatchModel: Model<SMSBatch>,
     private authService: AuthService,
@@ -141,8 +144,21 @@ export class GatewayService {
       )
     }
 
-    return {}
-    // return await this.deviceModel.findByIdAndDelete(deviceId)
+    await this.deviceTombstoneModel.updateOne(
+      { deviceId: new Types.ObjectId(deviceId) },
+      {
+        $setOnInsert: {
+          deviceId: new Types.ObjectId(deviceId),
+          userId: device.user,
+          deletedAt: new Date(),
+        },
+      },
+      { upsert: true },
+    )
+
+    await this.deviceModel.findByIdAndDelete(deviceId)
+
+    return { success: true }
   }
 
   private calculateDelayFromScheduledAt(scheduledAt?: string): number | undefined {
@@ -256,6 +272,7 @@ export class GatewayService {
 
     try {
       smsBatch = await this.smsBatchModel.create({
+        user: device.user,
         device: device._id,
         message,
         recipientCount: recipients.length,
@@ -278,6 +295,7 @@ export class GatewayService {
     for (let recipient of recipients) {
       recipient = recipient.replace(/\s+/g, "")
       const sms = await this.smsModel.create({
+        user: device.user,
         device: device._id,
         smsBatch: smsBatch._id,
         message: message,
@@ -465,6 +483,7 @@ export class GatewayService {
     const { messageTemplate, messages } = body
 
     const smsBatch = await this.smsBatchModel.create({
+      user: device.user,
       device: device._id,
       message: messageTemplate,
       recipientCount: messages
@@ -504,6 +523,7 @@ export class GatewayService {
       for (let recipient of recipients) {
         recipient = recipient.replace(/\s+/g, "")
         smsDocumentsToInsert.push({
+          user: device.user,
           device: device._id,
           smsBatch: smsBatch._id,
           message: message,
@@ -765,6 +785,7 @@ export class GatewayService {
     }
 
     const sms = await this.smsModel.create({
+      user: device.user,
       device: device._id,
       message: dto.message,
       type: SMSType.RECEIVED,

api/src/gateway/schemas/device-tombstone.schema.ts

@@ -0,0 +1,25 @@
+import { Prop, Schema, SchemaFactory } from '@nestjs/mongoose'
+import { Document, Types } from 'mongoose'
+import { User } from '../../users/schemas/user.schema'
+
+export type DeviceTombstoneDocument = DeviceTombstone & Document
+
+@Schema({ timestamps: true })
+export class DeviceTombstone {
+  _id?: Types.ObjectId
+
+  @Prop({ type: Types.ObjectId, required: true, unique: true, index: true })
+  deviceId: Types.ObjectId
+
+  @Prop({ type: Types.ObjectId, ref: User.name, required: true, index: true })
+  userId: User | Types.ObjectId
+
+  @Prop({ type: Date, required: true })
+  deletedAt: Date
+}
+
+export const DeviceTombstoneSchema =
+  SchemaFactory.createForClass(DeviceTombstone)
+
+DeviceTombstoneSchema.index({ userId: 1, deletedAt: -1 })
+

api/src/gateway/schemas/sms-batch.schema.ts

@@ -1,6 +1,7 @@
 import { Prop, Schema, SchemaFactory } from '@nestjs/mongoose'
 import { Document, Types } from 'mongoose'
 import { Device } from './device.schema'
+import { User } from '../../users/schemas/user.schema'
 
 export type SMSBatchDocument = SMSBatch & Document
 
@@ -8,6 +9,9 @@ export type SMSBatchDocument = SMSBatch & Document
 export class SMSBatch {
   _id?: Types.ObjectId
 
+  @Prop({ type: Types.ObjectId, ref: User.name, required: true, index: true })
+  user: User | Types.ObjectId
+
   @Prop({ type: Types.ObjectId, ref: Device.name })
   device: Device
 
@@ -53,3 +57,5 @@ export class SMSBatch {
 }
 
 export const SMSBatchSchema = SchemaFactory.createForClass(SMSBatch)
+
+SMSBatchSchema.index({ user: 1, createdAt: -1 })

api/src/gateway/schemas/sms.schema.ts

@@ -2,6 +2,7 @@ import { Prop, Schema, SchemaFactory } from '@nestjs/mongoose'
 import { Document, Types } from 'mongoose'
 import { Device } from './device.schema'
 import { SMSBatch } from './sms-batch.schema'
+import { User } from '../../users/schemas/user.schema'
 
 export type SMSDocument = SMS & Document
 
@@ -9,6 +10,9 @@ export type SMSDocument = SMS & Document
 export class SMS {
   _id?: Types.ObjectId
 
+  @Prop({ type: Types.ObjectId, ref: User.name, required: true, index: true })
+  user: User | Types.ObjectId
+
   @Prop({ type: Types.ObjectId, ref: Device.name, required: true })
   device: Device | Types.ObjectId
 
@@ -84,3 +88,4 @@ export const SMSSchema = SchemaFactory.createForClass(SMS)
 
 
 SMSSchema.index({ device: 1, type: 1, receivedAt: -1 })
+SMSSchema.index({ user: 1, createdAt: -1, type: 1 })