feat: use a single transaction for measurement writes

app/db.server.ts

@@ -1,4 +1,10 @@
-import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js'
+import {
+	drizzle,
+	type PostgresJsDatabase,
+	type PostgresJsQueryResultHKT,
+} from 'drizzle-orm/postgres-js'
+import { type ExtractTablesWithRelations } from 'drizzle-orm'
+import { type PgTransaction } from 'drizzle-orm/pg-core'
 import postgres, { type Sql } from 'postgres'
 import invariant from 'tiny-invariant'
 import * as schema from './db/schema'
@@ -57,3 +63,9 @@ function parsePoolSize(value: string | undefined): number {
 }
 
 export { drizzleClient, pg }
+
+export type DatabaseTransaction = PgTransaction<
+	PostgresJsQueryResultHKT,
+	typeof schema,
+	ExtractTablesWithRelations<typeof schema>
+>

app/db/models/device.server.ts

@@ -29,7 +29,7 @@ import {
 	type Sensor,
 } from '~/db/schema'
 import type * as schema from '~/db/schema/index'
-import { drizzleClient } from '~/db.server'
+import { drizzleClient, type DatabaseTransaction } from '~/db.server'
 import BaseNewDeviceEmail, {
 	messages as BaseNewDeviceMessages,
 } from '~/emails/base-new-device'
@@ -155,56 +155,66 @@ export type DeviceForSingleMeasurementWrite = Awaited<
 	ReturnType<typeof getDeviceForSingleMeasurementWrite>
 >
 
-export function getDeviceForMeasurementWrite({ id }: Pick<Device, 'id'>) {
-	return drizzleClient.query.device.findFirst({
-		where: (device, { eq }) => eq(device.id, id),
-		columns: {
-			id: true,
-			archivedAt: true,
-			useAuth: true,
-			apiKey: true,
-		},
-		with: {
-			sensors: {
-				columns: {
-					id: true,
-					title: true,
-					sensorType: true,
-				},
-			},
-		},
-	})
+export async function getDeviceForMeasurementWrite(
+	{ id }: Pick<Device, 'id'>,
+	tx: DatabaseTransaction,
+) {
+	const currentDevice = await lockDeviceForMeasurementWrite(id, tx)
+	if (!currentDevice) return undefined
+
+	const sensors = await tx
+		.select({
+			id: sensor.id,
+			title: sensor.title,
+			sensorType: sensor.sensorType,
+		})
+		.from(sensor)
+		.where(eq(sensor.deviceId, id))
+		.orderBy(sensor.id)
+		.for('update')
+
+	return { ...currentDevice, sensors }
 }
 
-export async function getDeviceForSingleMeasurementWrite({
-	id,
-	sensorId,
-}: Pick<Device, 'id'> & { sensorId: string }) {
-	const [row] = await drizzleClient
+export async function getDeviceForSingleMeasurementWrite(
+	{ id, sensorId }: Pick<Device, 'id'> & { sensorId: string },
+	tx: DatabaseTransaction,
+) {
+	const currentDevice = await lockDeviceForMeasurementWrite(id, tx)
+	if (!currentDevice) return undefined
+
+	const [currentSensor] = await tx
+		.select({
+			id: sensor.id,
+		})
+		.from(sensor)
+		.where(and(eq(sensor.deviceId, currentDevice.id), eq(sensor.id, sensorId)))
+		.limit(1)
+		.for('update')
+
+	return {
+		...currentDevice,
+		sensors: currentSensor ? [currentSensor] : [],
+	}
+}
+
+async function lockDeviceForMeasurementWrite(
+	id: Device['id'],
+	tx: DatabaseTransaction,
+) {
+	const [currentDevice] = await tx
 		.select({
 			id: device.id,
 			archivedAt: device.archivedAt,
 			useAuth: device.useAuth,
 			apiKey: device.apiKey,
-			sensorId: sensor.id,
 		})
 		.from(device)
-		.leftJoin(
-			sensor,
-			and(eq(sensor.deviceId, device.id), eq(sensor.id, sensorId)),
-		)
 		.where(eq(device.id, id))
 		.limit(1)
+		.for('share')
 
-	if (!row) return undefined
-
-	return {
-		id: row.id,
-		archivedAt: row.archivedAt,
-		useAuth: row.useAuth,
-		apiKey: row.apiKey,
-		sensors: row.sensorId ? [{ id: row.sensorId }] : [],
-	}
+	return currentDevice
 }
 
 export function getUserDevice({ id, userId }: Pick<Device, 'id' | 'userId'>) {

app/db/models/measurement.server.ts

@@ -1,5 +1,4 @@
 import { and, desc, eq, gt, gte, inArray, lt, lte, sql } from 'drizzle-orm'
-import { ArchivedDeviceError } from './device.server'
 import {
 	type LastMeasurement,
 	location,
@@ -9,9 +8,8 @@
 	measurements1hourView,
 	measurements1monthView,
 	measurements1yearView,
-	device,
 } from '~/db/schema'
-import { drizzleClient } from '~/db.server'
+import { drizzleClient, type DatabaseTransaction } from '~/db.server'
 import {
 	type MinimalDevice,
 	type MeasurementWithLocation,
@@ -178,12 +176,12 @@
 }
 
 export async function saveMeasurements(
+	tx: DatabaseTransaction,
 	minimalDevice: MinimalDevice,
 	measurements: MeasurementWithLocation[],
 	timing?: MeasurementTiming | null,
 ): Promise<void> {
-	if (!device) throw new Error('No device given!')
 	if (!Array.isArray(measurements)) throw new Error('Array expected')
 
 	const sensorIds = new Set(minimalDevice.sensors.map((s: any) => s.id))
 	const lastMeasurements: Record<string, NonNullable<LastMeasurement>> = {}
@@ -234,57 +232,35 @@
 		locationUpdateCount: deviceLocationUpdates.length,
 	})
 
-	await drizzleClient.transaction(async (tx) => {
-		const [currentDevice] = await tx
-			.select({
-				id: device.id,
-				archivedAt: device.archivedAt,
-			})
-			.from(device)
-			.where(eq(device.id, minimalDevice.id))
-			.limit(1)
-		timing?.mark('transactionDeviceLookup')
-
-		if (!currentDevice) {
-			const error = new Error('Device not found')
-			error.name = 'NotFoundError'
-			throw error
-		}
-
-		if (currentDevice.archivedAt) {
-			throw new ArchivedDeviceError(currentDevice.id)
-		}
-
-		const locations =
-			deviceLocationUpdates.length > 0
-				? await findOrCreateLocations(deviceLocationUpdates)
-				: []
-		timing?.mark('findOrCreateLocations', {
-			locationCount: locations.length,
-		})
-
-		if (deviceLocationUpdates.length > 0) {
-			await addLocationUpdates(
-				deviceLocationUpdates,
-				minimalDevice.id,
-				locations,
-			)
-		}
-		timing?.mark('addLocationUpdates')
+	const locations =
+		deviceLocationUpdates.length > 0
+			? await findOrCreateLocations(deviceLocationUpdates, tx)
+			: []
+	timing?.mark('findOrCreateLocations', {
+		locationCount: locations.length,
+	})
 
-		await insertMeasurementsWithLocation(
-			measurements,
-			locations,
+	if (deviceLocationUpdates.length > 0) {
+		await addLocationUpdates(
+			deviceLocationUpdates,
 			minimalDevice.id,
+			locations,
 			tx,
-			{ shouldReturn: false },
-			timing,
 		)
-		timing?.mark('insertMeasurements')
-		await updateLastMeasurements(lastMeasurements, tx, timing)
-		timing?.mark('updateLastMeasurements')
-	})
-	timing?.mark('transaction')
+	}
+	timing?.mark('addLocationUpdates')
+
+	await insertMeasurementsWithLocation(
+		measurements,
+		locations,
+		minimalDevice.id,
+		tx,
+		{ shouldReturn: false },
+		timing,
+	)
+	timing?.mark('insertMeasurements')
+	await updateLastMeasurements(lastMeasurements, tx, timing)
+	timing?.mark('updateLastMeasurements')
 }
 
 export async function insertMeasurements(measurements: any[]): Promise<void> {