[BUG]: calling $count on a replica using withReplicas throws a TypeError

[AndrewJo]

Report hasn't been filed before.

• I have verified that the bug I'm about to report hasn't been filed before.

What version of drizzle-orm are you using?

0.39.1

What version of drizzle-kit are you using?

0.30.4

Other packages

No response

Describe the Bug

When resolving $count on a client with replicas, it results in a TypeError with message: Cannot read properties of undefined (reading 'session')"

const db = withReplicas(primaryDb, [replicaDb]);
// This line throws TypeError
const postCount = await db.$count(posts);

The stack traces back to this line where this is undefined because the this context is likely not bound properly:

drizzle-orm/drizzle-orm/src/pg-core/db.ts

Line 152 in 5ffd0bd

return new PgCountBuilder({ source, filters, session: this.session });

I'm wondering if this is because calling withReplica returns a plain object without the associated session object.

drizzle-orm/drizzle-orm/src/pg-core/db.ts

Lines 671 to 689 in 5ffd0bd

	return {
	...primary,
	update,
	insert,
	delete: $delete,
	execute,
	transaction,
	refreshMaterializedView,
	$primary: primary,
	select,
	selectDistinct,
	selectDistinctOn,
	$count,
	$with,
	with: _with,
	get query() {
	return getReplica(replicas).query;
	},
	};

Or whether calling getReplicas.$count(...args) is not binding the underlying replica PgDatabase instance to this.

drizzle-orm/drizzle-orm/src/pg-core/db.ts

Line 659 in 5ffd0bd

const $count: Q['$count'] = (...args: [any]) => getReplica(replicas).$count(...args);

If this is indeed the issue, the bug most likely exists for all dialects and not just Postgres.

I'm beginning to wonder if withReplica needs to return a class that extends PgDatabase (or equivalent in each dialect) that contains $primary property among others.

This is something that didn't grab my attention when I submitted #3952.

@AndriiSherman, @Sukairo-02: I'll be happy to submit a PR based on the outcomes of the discussion here.