logto/packages/schemas/alterations/1.0.0_rc.0-1674032095.4-add-id-column.ts

import { sql } from 'slonik';

import type { AlterationScript } from '../lib/types/alteration.js';

// https://github.com/viascom/nanoid-postgres/blob/main/nanoid.sql
// Removed `_-` from the default alphabet to match our JS ID generator.
const createNanoId = sql`
  /*
   * Copyright 2022 Viascom Ltd liab. Co
   *
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations
   * under the License.
   */

  CREATE EXTENSION IF NOT EXISTS pgcrypto;

  CREATE OR REPLACE FUNCTION nanoid(size int DEFAULT 21, alphabet text DEFAULT '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ')
      RETURNS text
      LANGUAGE plpgsql
      volatile
  AS
  $$
  DECLARE
      idBuilder     text := '';
      i             int  := 0;
      bytes         bytea;
      alphabetIndex int;
      mask          int;
      step          int;
  BEGIN
      mask := (2 << cast(floor(log(length(alphabet) - 1) / log(2)) as int)) - 1;
      step := cast(ceil(1.6 * mask * size / length(alphabet)) AS int);

      while true
          loop
              bytes := gen_random_bytes(size);
              while i < size
                  loop
                      alphabetIndex := (get_byte(bytes, i) & mask) + 1;
                      if alphabetIndex <= length(alphabet) then
                          idBuilder := idBuilder || substr(alphabet, alphabetIndex, 1);
                          if length(idBuilder) = size then
                              return idBuilder;
                          end if;
                      end if;
                      i = i + 1;
                  end loop;

              i := 0;
          end loop;
  END
  $$;
`;

const getId = (value: string) => sql.identifier([value]);

type TableInfo = {
  table: string;
  primaryKeys: string[];
};

const tablesWithId: TableInfo[] = [
  { table: 'applications_roles', primaryKeys: ['application_id', 'role_id'] },
  { table: 'custom_phrases', primaryKeys: ['language_tag'] },
  { table: 'roles_scopes', primaryKeys: ['role_id', 'scope_id'] },
  { table: 'users_roles', primaryKeys: ['user_id', 'role_id'] },
];

const alteration: AlterationScript = {
  up: async (pool) => {
    await pool.query(createNanoId);

    // Add id column to tables
    await Promise.all(
      tablesWithId.map(async ({ table, primaryKeys }) =>
        pool.query(sql`
          alter table ${getId(table)}
            add column id varchar(21) not null default nanoid(),
            drop constraint ${getId(table + '_pkey')},
            add primary key (id),
            add constraint ${getId(table + '__' + primaryKeys.join('_'))}
              unique (${sql.join(
                primaryKeys.map((key) => getId(key)),
                sql`, `
              )});
        `)
      )
    );

    // Nano IDs should be generated by JS to align with other tables, so we drop the default value here.
    // Also usually backend server has a better specification. :-)
    await Promise.all(
      tablesWithId.map(async ({ table }) =>
        pool.query(sql`
          alter table ${getId(table)} alter column id drop default;
        `)
      )
    );

    await pool.query(sql`drop function nanoid;`);
  },
  down: async (pool) => {
    // Add id column to tables
    await Promise.all(
      tablesWithId.map(async ({ table, primaryKeys }) =>
        pool.query(sql`
          alter table ${getId(table)}
            drop column id,
            drop constraint ${getId(table + '__' + primaryKeys.join('_'))},
            add primary key (${sql.join(
              primaryKeys.map((key) => getId(key)),
              sql`, `
            )});
        `)
      )
    );
  },
};

export default alteration;