Working around SCOPE_IDENTITY not working with PDO

I discovered recently that due to the way the PDO_SQLSRV driver handles prepared statements, SCOPE_IDENTITY() fails to work. Generally one could use @@IDENTITY to work around this as most of the time they accomplish the same task, but triggers that perform an insert can get you into trouble there.

Triggers causing problems is exactly what lead me to this problem. I recently added some triggers to a few tables in an application which insert data into a third table whenever an insert or update is done on the original tables. After adding these triggers I discovered that the code for inserting data into the tables was no longer functioning properly. The data was being inserted but the ID being returned was incorrect. I'd run into this trigger problem before and knew the problem was using @@IDENTITY rather than SCOPE_IDENTITY() so I went into the PHP code and fixed the PDO::lastInsertId function to use SCOPE_IDENTITY() and figured that'd fix it all up. To my surprise it was still broken except now rather than the wrong ID I was getting no ID as SCOPE_IDENTITY() simply returned NULL.

After a lot of debugging and failed attempts at fixing the problem I decided that the problem must be that each query is being run in it's own scope and thus there was no way to recover the ID from a future query. After a lot of googling I once again though I'd found a fix, namely the PDO::SQLSRV_ATTR_DIRECT_QUERY attribute. Apparently by default the PDO_SQLSRV driver does wraps every query as a prepared statement and does RPC call to SQL Server. A side effect of this is that each query runs in it's own scope so things like variables, temporary tables, etc do not work.

If a query requires the context that was set in a previous query, you should execute your queries with PDO::SQLSRV_ATTR_DIRECT_QUERY set to True. For example, if you use temporary tables in your queries, PDO::SQLSRV_ATTR_DIRECT_QUERY must be set to True.

Alright, so surely just setting that to true will resolve the problem, right? Wrong! Ugh. As it turns out, setting that does help by making normal PDO::query calls execute directly, but it has no effect when using PDO::prepare and bound parameters. These queries still run as a RPC call and still have an isolated scope, thus SCOPE_IDENTITY() still returns NULL.

So at this point I've discovered the only way to get the ID is to request it within the same RPC call as the query. How can I do that without having to go through the application and re-writing every INSERT query though? The solution I came up with is to use a global temporary table to store the ID's and then append a second INSERT query to each query being run that attempts to insert SCOPE_IDENTITY() into this global temporary table. It's a pretty terrible hack, but it seems to work.

class DBConnection extends PDO {
    private $token;

    public function __construct($host, $user, $pass, $database){
        parent::__construct('sqlsrv:SERVER='.$host.';DATABASE='.$database, $user, $pass);
        $this->setAttribute(PDO::SQLSRV_ATTR_DIRECT_QUERY, true);
        $this->setAttribute(PDO::ATTR_EMULATE_PREPARES, false);

        $this->setupLastInsertIdHack();
    }

    function __destruct(){
        $this->exec("DELETE FROM ##lastInsertId WHERE token='{$this->token}'");
    }

    private function setupLastInsertIdHack(){
        $this->exec("
        IF OBJECT_ID('tempdb..##lastInsertId') IS NULL
        BEGIN
            CREATE TABLE ##lastInsertId (id INT, token VARCHAR(30) NOT NULL, ts DATETIME2 NOT NULL)
            CREATE CLUSTERED INDEX IX_token ON ##lastInsertId (token)
            CREATE NONCLUSTERED INDEX IX_ts ON ##lastInsertId (ts)
        END");

        $this->token = uniqid('token', true);

        $this->exec("DELETE FROM ##lastInsertId WHERE token='{$this->token}'");
    }

    public function prepare($sql, $opts = array()){
        $sql .= "\r\nSET NOCOUNT ON\r\nINSERT INTO ##lastInsertId (id, token, ts) SELECT SCOPE_IDENTITY(), '{$this->token}', SYSUTCDATETIME()";

        return parent::prepare($sql, $opts);
    }

    public function lastInsertId($name=null){
        $query = $this->query("
        SELECT COALESCE(SCOPE_IDENTITY(), t.id, @@IDENTITY)
        FROM ##lastInsertId t
        WHERE
            t.token = '{$this->token}'
            AND t.ts = (SELECT MAX(ts) FROM ##lastInsertId WHERE token='{$this->token}')
        ");

        return intval($query->fetchColumn(0));
    }
}

First off, we have to create the global temporary table if it doesn't already exist. We add this to the connection initialization so it's run on each request. We check if the table exists using OBJECT_ID. If it does not exist we create it and some indexes on it to ensure fast lookups. Rather than just a single column to store the result of SCOPE_IDENTITY() we also store the connection id and a timestamp. This prevents issues from concurrent requests and permits grabbing the most recent ID.

Next we need to generate a unique token for this session. This is to prevent problems from concurrent requests getting there ID data mixed up. Each connection will generate it's own unique token to associate with it's generated IDs. We also clear out any old entries for the token just incase the token is a repeat where old entries were left behind. In an ideal world there wouldn't be anything in the table with that token value as the tokens should be unique and we run a cleanup on shutdown also, but better safe than sorry.

Then we override the PDO::prepare method to append an query which inserts into this table. I chose to simply append the the query unconditionally which seems to work fine. One could try and sniff for INSERT queries and only modify those instead.

Finally we override the lastInsertId method to use the data from this table instead of it's default implementation of only using @@IDENTITY. In addition to our hack table, we still check the SCOPE_IDENTITY() and @@IDENTITY values as alternatives. We check SCOPE_IDENTITY() first since it should be the preferred value if one exists. If it's null we then fall back to our table. If nothing exist in our table or it is NULL then we fall back to @@IDENTITY as a last resort.

So there you have it, a workaround for this this annoying problem. Hopefully this saves someone time as it took me quite a while to figure out the problem and then find a solution. It should be noted that this is certainly a hack and there are better ways to work around the issue if you're willing to alter your INSERT queries and not use PDO::lastInsertId() at all. Eventually I may get around to a more proper fix for this but I needed a quick solution that worked with existing PDO::lastInsertId() reliant code.