Skip to content

Instantly share code, notes, and snippets.

@schpet
Created July 12, 2017 15:22
Show Gist options
  • Save schpet/144ffb71e33809e5ab56f2d333d7c3ae to your computer and use it in GitHub Desktop.
Save schpet/144ffb71e33809e5ab56f2d333d7c3ae to your computer and use it in GitHub Desktop.
short postgres primary keys rails migrations based on https://blog.andyet.com/2016/02/23/generating-shortids-in-postgres/
diff --git db/migrate/20170712145519_add_unique_short_id_function.rb db/migrate/20170712145519_add_unique_short_id_function.rb
new file mode 100644
index 0000000..4342d8b
--- /dev/null
+++ db/migrate/20170712145519_add_unique_short_id_function.rb
@@ -0,0 +1,76 @@
+class AddUniqueShortIdFunction < ActiveRecord::Migration[5.1]
+ def up
+ # if needed:
+ # enable_extension 'pgcrypto'
+ #
+ # short id function from:
+ # https://blog.andyet.com/2016/02/23/generating-shortids-in-postgres/
+ unique_short_id_function_sql = <<-EOS
+ -- Create a trigger function that takes no arguments.
+ -- Trigger functions automatically have OLD, NEW records
+ -- and TG_TABLE_NAME as well as others.
+ CREATE OR REPLACE FUNCTION unique_short_id()
+ RETURNS TRIGGER AS $$
+
+ -- Declare the variables we'll be using.
+ DECLARE
+ key TEXT;
+ qry TEXT;
+ found TEXT;
+ BEGIN
+
+ -- generate the first part of a query as a string with safely
+ -- escaped table name, using || to concat the parts
+ qry := 'SELECT id FROM ' || quote_ident(TG_TABLE_NAME) || ' WHERE id=';
+
+ -- This loop will probably only run once per call until we've generated
+ -- millions of ids.
+ LOOP
+
+ -- Generate our string bytes and re-encode as a base64 string.
+ key := encode(gen_random_bytes(6), 'base64');
+
+ -- Base64 encoding contains 2 URL unsafe characters by default.
+ -- The URL-safe version has these replacements.
+ key := replace(key, '/', '_'); -- url safe replacement
+ key := replace(key, '+', '-'); -- url safe replacement
+
+ -- Concat the generated key (safely quoted) with the generated query
+ -- and run it.
+ -- SELECT id FROM "test" WHERE id='blahblah' INTO found
+ -- Now "found" will be the duplicated id or NULL.
+ EXECUTE qry || quote_literal(key) INTO found;
+
+ -- Check to see if found is NULL.
+ -- If we checked to see if found = NULL it would always be FALSE
+ -- because (NULL = NULL) is always FALSE.
+ IF found IS NULL THEN
+
+ -- If we didn't find a collision then leave the LOOP.
+ EXIT;
+ END IF;
+
+ -- We haven't EXITed yet, so return to the top of the LOOP
+ -- and try again.
+ END LOOP;
+
+ -- NEW and OLD are available in TRIGGER PROCEDURES.
+ -- NEW is the mutated row that will actually be INSERTed.
+ -- We're replacing id, regardless of what it was before
+ -- with our key variable.
+ NEW.id = key;
+
+ -- The RECORD returned here is what will actually be INSERTed,
+ -- or what the next trigger will get if there is one.
+ RETURN NEW;
+ END;
+ $$ language 'plpgsql';
+ EOS
+
+ execute(unique_short_id_function_sql)
+ end
+
+ def down
+ execute "DROP FUNCTION unique_short_id();"
+ end
+end
diff --git db/migrate/20170712150121_apply_short_id_function_to_recipes.rb db/migrate/20170712150121_apply_short_id_function_to_recipes.rb
new file mode 100644
index 0000000..c510030
--- /dev/null
+++ db/migrate/20170712150121_apply_short_id_function_to_recipes.rb
@@ -0,0 +1,15 @@
+class ApplyShortIdFunctionToRecipes < ActiveRecord::Migration[5.1]
+ def up
+ execute(%Q{
+ -- We name the trigger "trigger_test_genid" so that we can remove
+ -- or replace it later.
+ -- If an INSERT contains multiple RECORDs, each one will call
+ -- unique_short_id individually.
+ CREATE TRIGGER trigger_recipes_genid BEFORE INSERT ON recipes FOR EACH ROW EXECUTE PROCEDURE unique_short_id();
+ })
+ end
+
+ def down
+ execute("DROP TRIGGER trigger_recipes_genid ON recipes")
+ end
+end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment