Drizzle ORM's Returning() Not Working With Pg-mem

by Alex Johnson 50 views

Hey there, fellow developers! Let's dive into a peculiar issue I've bumped into while working with Drizzle ORM, specifically when using it with the pg-mem adapter. This is a heads-up if you're leveraging Drizzle to manage your database interactions and leaning on pg-mem for in-memory testing. The core problem revolves around the returning() function, and how it behaves (or rather, doesn't behave quite as expected) when paired with pg-mem. Let's break down the scenario and explore what's happening.

The Heart of the Matter: returning() and pg-mem

First off, Drizzle ORM is a fantastic tool for writing type-safe SQL queries in your TypeScript or JavaScript projects. It provides a clean, modern approach to database interactions, making your code more readable and maintainable. pg-mem, on the other hand, is an in-memory PostgreSQL emulator. It's a lifesaver for writing unit tests because it allows you to simulate a PostgreSQL database environment without actually connecting to a live database. This speeds up your tests and isolates them from external dependencies.

The expectation when using returning() with Drizzle is that after an INSERT, UPDATE, or DELETE operation, you can retrieve the inserted or modified rows' data directly. This is incredibly useful for validating that your database operations are successful and for accessing generated values like auto-incrementing IDs. The test case highlights that when you use returning() with pg-mem, the returned values are not populated as expected. The rows are returned, but the values inside are undefined or default values, not the actual values that should have been inserted or updated.

Now, the plot thickens because the raw SQL equivalent (...returning *) does work as intended within the same setup. This means that if you execute a raw SQL query with RETURNING *, you'll get the expected data back. This discrepancy is the core of the bug. It indicates that the issue lies somewhere in how Drizzle ORM translates and executes the returning() call with pg-mem.

Code Deep Dive: The Problem Unveiled

Let's take a closer look at the test case to understand how this is happening. The test starts by setting up a basic pg-mem database. Then, it defines a table named products with an id (serial, primary key) and name (text, not null) column. The setup code includes a patch to the query function so that it correctly handles prepared statements.

The test tries to insert a product using Drizzle's returning() function. Afterwards, it checks the returned value. The crucial part of the test is the assertions: expect(inserted[0].id).toBeUndefined() and expect(inserted[0].name).toBeUndefined(). These lines assert that the id and name properties of the returned object are undefined, which reveals the bug. The test goes on to demonstrate that raw SQL with RETURNING * works fine, confirming the issue's specificity to Drizzle's returning().

import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import { pgTable, serial, text } from 'drizzle-orm/pg-core';
import { drizzle } from 'drizzle-orm/node-postgres';
import { newDb } from 'pg-mem';

const products = pgTable('products', {
  id: serial('id').primaryKey(),
  name: text('name').notNull(),
});

describe('pg-mem returning() limitation', () => {
  let db: any;
  let pool: any;

  beforeAll(async () => {
    const pgMemDb = newDb();
    const pgAdapter = pgMemDb.adapters.createPg();
    pool = new pgAdapter.Pool();
    db = drizzle(pool);

    const originalQuery = pool.query.bind(pool);
    pool.query = function (text: any, params?: any) {
      // If it's a prepared statement object, extract the text
      if (typeof text === 'object' && text.text) {
        return originalQuery(text.text, params || text.values);
      }
      return originalQuery(text, params);
    } as any;

    await pool.query(`
      CREATE TABLE IF NOT EXISTS products (
        id serial PRIMARY KEY,
        name text NOT NULL
      )
    `);
  });

  afterAll(async () => {
    await pool.end();
  });

  it('should demonstrate that returning() does not work with pg-mem', async () => {
    // Try to insert a product using drizzle's returning()
    const inserted = await db
      .insert(products)
      .values({
        name: 'Test Product 1',
      })
      .returning();

    // Prove that returning() doesn't work - the returned values are undefined
    expect(inserted).toBeDefined();
    expect(inserted.length).toBeGreaterThan(0);
    expect(inserted[0].id).toBeUndefined();                 // ‼️ these two lines assert the problem
    expect(inserted[0].name).toBeUndefined();

    // Insert a product using raw sql and returning *
    const insertedWithRawSql = await pool.query(`
      INSERT INTO products (name) VALUES ('Test Product 2') RETURNING *
    `);

    expect(insertedWithRawSql.rows).toBeDefined();
    expect(insertedWithRawSql.rows.length).toBeGreaterThan(0);
    expect(insertedWithRawSql.rows[0].name).toBe('Test Product 2');
    expect(insertedWithRawSql.rows[0].id).toBeDefined();
    expect(typeof insertedWithRawSql.rows[0].id).toBe('number');

    // Verify the data WAS actually inserted by using raw SQL query
    const result = await pool.query(`
      SELECT * FROM products WHERE name LIKE '%Test Product%' LIMIT 2
    `);

    // The raw SQL query proves the data exists in the database
    expect(result.rows).toBeDefined();
    expect(result.rows.length).toBeGreaterThan(0);
    const [product1, product2] = result.rows;
    expect(product1.name).toMatch(/Test Product/);
    expect(product2.name).toMatch(/Test Product/);

    console.log('Inserted (drizzle returning()):', inserted[0]);
    console.log('Inserted (SQL returning *):', insertedWithRawSql.rows[0]);
    console.log('Queried (SQL):', result.rows);
  });
});

Potential Causes and Workarounds

While the exact root cause might require digging into the internals of Drizzle ORM and pg-mem, a few potential causes come to mind. It's possible that pg-mem has limitations in how it handles returning() clauses when used through a prepared statement generated by Drizzle. Another possibility is that there's a misinterpretation or incorrect translation of the returning() syntax in Drizzle when interacting with pg-mem.

For now, if you're encountering this issue, here are a few workarounds:

  1. Use Raw SQL with RETURNING *: As demonstrated in the test, you can bypass the issue by writing raw SQL queries with RETURNING *. This gives you the desired behavior, but it sacrifices some of the type safety and abstraction that Drizzle ORM offers.

  2. Query After Insert: After inserting the data, you can immediately query the database to retrieve the newly inserted row. This requires an extra database round trip, which can impact performance, but it provides the correct data.

  3. Investigate Drizzle ORM's Implementation: If you're feeling adventurous, you could delve into the source code of Drizzle ORM to see how it generates the SQL for returning() and how it interacts with the underlying database adapter. This might give you insights into the problem and potentially allow you to contribute a fix or a workaround.

Conclusion: Navigating the returning() Puzzle

So, there you have it – a bug related to Drizzle ORM's returning() function not working as expected with the pg-mem adapter. This is an excellent example of how different tools and libraries can interact in unexpected ways. It is essential to understand the limitations of each tool and to adapt your approach accordingly. Always verify your assumptions, especially when using in-memory databases for testing, and be prepared to use workarounds when necessary. This is an evolving landscape, and libraries will continue to update, so it is crucial to stay informed on the changes.

I hope this explanation has been helpful. Keep an eye on updates to both Drizzle ORM and pg-mem, as this issue may be resolved in future versions. Happy coding!

External Resources

For further reading and insights into Drizzle ORM and its capabilities, you may find the official documentation and community resources helpful. The Drizzle ORM GitHub repository is an excellent place to find examples, report issues, and follow the development of the project. Also, you can find the pg-mem GitHub repository to get the latest updates.