February 6th, 2017 by depesz | Tags: , , , , , , | 1 comment »
Did it help? If yes - maybe you can help me?

Couple of people asked, on irc, about how to generate unique, short, ids for rows.

Since I wrote code for this for explain.depesz.com, I figured that instead of pointing to sources, and letting everyone to extract the functions, I will, instead, write a set of functions that what is needed in reusable way.

First thing that we need is a function that will generate random string. It will need one or two arguments:

  1. number of characters that the string should have
  2. optionally – set of characters to use for generation

The function itself is pretty simple:

$ CREATE FUNCTION get_random_string(
        IN string_length integer,
        IN possible_chars TEXT DEFAULT '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
    ) RETURNS text
    LANGUAGE plpgsql
    AS $$
DECLARE
    output TEXT = '';
    i INT4;
    pos INT4;
BEGIN
    FOR i IN 1..string_length LOOP
        pos := 1 + cast( random() * ( length(possible_chars) - 1) as INT4 );
        output := output || substr(possible_chars, pos, 1);
    END LOOP;
    RETURN output;
END;
$$;
CREATE FUNCTION

Sample usage:

$ select get_random_string(10);
 get_random_string 
-------------------
 BOaM1K6wbS
(1 row)
 
$ select get_random_string(15);
 get_random_string 
-------------------
 fWejQRndBUyJ5Mh
(1 row)
 
$ select get_random_string(5, 'abc');
 get_random_string 
-------------------
 babbc
(1 row)

Then we need a way to insert data into table. Using a function (like explain.depesz.com) is not always convenient, so perhaps we can make something like:

$ create table test (
    id text default get_random_string(10) primary key,
    something text
);
CREATE TABLE

Obviously, it will work:

$ insert into test (something)
    select 'sth #' || i
    from generate_series(1,10) i;
INSERT 0 10
 
$ select * from test;
     id     | something 
------------+-----------
 QOGAsBmuXQ | sth #1
 yu24hNR6cx | sth #2
 BJBlvuKHlx | sth #3
 ZCMoMF18AY | sth #4
 Z9TbDByeHa | sth #5
 cSunFqiZ8U | sth #6
 XggtW39WBK | sth #7
 5kTYMfjLK1 | sth #8
 vwTqkihTIp | sth #9
 xoWei3hrZs | sth #10
(10 rows)

But there are couple of issues:

  • what will happen if the same string will get generated?
  • 10 characters is not exactly short

With 10 characters in id, chances of generating duplicate id are pretty slim, so let's cheat a bit:

$ alter table test alter column id set default get_random_string(2, 'ab');
ALTER TABLE
 
$ insert into test (something)
    select 'sth #' || i
    from generate_series(11,20) i;
ERROR:  duplicate key value violates unique constraint "test_pkey"
DETAIL:  Key (id)=(aa) already exists.

Obviously, it's not good. Inserts use default value to assume that it will work.

So, let's make another function, that will be used as default:

$ create or replace function generate_random_id(
    IN string_length integer,
    IN possible_chars TEXT DEFAULT '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
) returns text
    LANGUAGE plpgsql
    AS $$
DECLARE
    v_random_id text;
    v_temp text;
BEGIN
    LOOP
        v_random_id := get_random_string( string_length, possible_chars );
        select id INTO v_temp FROM test where id = v_random_id;
        exit when not found;
    END LOOP;
    return v_random_id;
END;
$$;
CREATE FUNCTION

Now, let's test it:

$ alter table test alter column id set default generate_random_id(2, 'ab');
ALTER TABLE
 
$ insert into test (something) values ('a');
INSERT 0 1
 
$ insert into test (something) values ('b');
INSERT 0 1
 
$ insert into test (something) values ('c');
INSERT 0 1
 
$ insert into test (something) values ('d');
INSERT 0 1
 
$ insert into test (something) values ('e');
Cancel request sent
ERROR:  canceling statement due to user request
CONTEXT:  SQL statement "select id             FROM test where id = v_random_id"
PL/pgSQL function generate_random_id(integer,text) line 8 at SQL statement

I stopped it, because it entered infinite loop. The problem is that once all possible combinations of possible_chars with given string_length are used, then it will simply spin in infinite loop.

Not so great. Luckily – it's trivial to fix. We'll just make the loop increase length of id if collision will be found:

$ create or replace function generate_random_id(
    IN string_length integer,
    IN possible_chars TEXT DEFAULT '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
) returns text
    LANGUAGE plpgsql
    AS $$
DECLARE
    v_random_id  text;
    v_temp       text;
    v_length     int4   :=  string_length;
BEGIN
    LOOP
        v_random_id := get_random_string( v_length, possible_chars );
        select id INTO v_temp FROM test where id = v_random_id;
        exit when not found;
        v_length := v_length + 1;
    END LOOP;
    return v_random_id;
END;
$$;
CREATE FUNCTION

So, now, even with default that generates strings with length “2", if I'll try to insert, I will get id with length of 3:

$ insert into test (something) values ('e') returning *;
 id  | something 
-----+-----------
 bab | e
(1 row)
INSERT 0 1

This is pretty good, but there are still problems.

First problem: what to do if I want to generate ID for table that is not named test?

That should be solvable too, let's pass table name as parameter. Since we know that tables can be in various schema, let's pass schema and table names separately. And we'll need column name too:

$ create or replace function generate_random_id(
    IN table_schema   TEXT,
    IN table_name     TEXT,
    IN column_name    TEXT,
    IN string_length  INTEGER,
    IN possible_chars TEXT DEFAULT '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
) returns text
    LANGUAGE plpgsql
    AS $$
DECLARE
    v_random_id  text;
    v_temp       text;
    v_length     int4   :=  string_length;
    v_sql        text;
BEGIN
    v_sql := format( 'SELECT %I FROM %I.%I WHERE %I = $1', column_name, table_schema, table_name, column_name );
    LOOP
        v_random_id := get_random_string( v_length, possible_chars );
        EXECUTE v_sql INTO v_temp USING v_random_id;
        exit when v_temp is null;
        v_length := v_length + 1;
    END LOOP;
    return v_random_id;
END;
$$ STRICT;
CREATE FUNCTION

I added STRICT designation, to avoid processing when any argument is NULL.

So, now I can:

$ alter table test alter column id set default generate_random_id('public', 'test', 'id', 2, 'ab');
ALTER TABLE
 
$ insert into test (something) values ('f'),('g'),('h') returning *;
 id  | something 
-----+-----------
 baa | f
 aab | g
 aaa | h
(3 rows)

Nice. Works. But does it really?

Let's consider, what will happen in this (unlikely, but we want to be safe) scenario:

  • we have two concurrent sessions (s1 and s2)
  • both sessions issue insert which calls our random generator
  • generator, in both cases, generates exactly the same random id
  • both functions (generate_random_id) check in underlying table if the id is used – and in both of them it returns false – the id is unused
  • both functions return the same id to insert command
  • one of the inserts will succeed, but the other one will fail.

Of course, the problem might seem far fetched – chances are slim, to put it lightly, but why take a chance anyway? We can fix it.

To fix it we can use advisory locks. New version of the function:

$ create or replace function generate_random_id(
    IN table_schema   TEXT,
    IN table_name     TEXT,
    IN column_name    TEXT,
    IN string_length  INTEGER,
    IN possible_chars TEXT DEFAULT '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
) returns text
    LANGUAGE plpgsql
    AS $$
DECLARE
    v_random_id   text;
    v_temp        text;
    v_length      int4   :=  string_length;
    v_sql         text;
    v_advisory_1  int4 := hashtext( format('%I:%I:%I', table_schema, table_name, column_name) );
    v_advisory_2  int4;
    v_advisory_ok bool;
BEGIN
    v_sql := format( 'SELECT %I FROM %I.%I WHERE %I = $1', column_name, table_schema, table_name, column_name );
    LOOP
        v_random_id := get_random_string( v_length, possible_chars );
        v_advisory_2 := hashtext( v_random_id );
        v_advisory_ok := pg_try_advisory_xact_lock( v_advisory_1, v_advisory_2 );
        IF v_advisory_ok THEN
            EXECUTE v_sql INTO v_temp USING v_random_id;
            exit when v_temp is null;
        END IF;
        v_length := v_length + 1;
    END LOOP;
    return v_random_id;
END;
$$ STRICT;
CREATE FUNCTION

Quick sanity check:

$ insert into test (something)
    select unnest('{i,j,k,l,m}'::text[]) returning *;
  id   | something 
-------+-----------
 baaa  | i
 abbbb | j
 abb   | k
 bbba  | l
 abbab | m
(5 rows)

Since getting advisory lock is cheaper than doing actual data check, I figured that getting the lock first is slightly better idea.

After some using of the function, you might find out that it always returns strings longer than the minimal requested – it could be because pool of possible combinations of possible_chars with given length is empty. In such case – you might want to change default expression to start with longer strings, to avoid one loop that will always fail.

Now, since I already wrote the code, there is no point in redoing it by others. The extension is available on GitHub and in PGXN. You can get it, install, use and/or modify. Have fun 🙂

  1. One comment

  2. # Richard
    Feb 6, 2017

    Very nice, thank you for sharing.

Leave a comment