Add database helper for creating triggers to assign sharding keys

What does this MR do and why?

Adds a database helper that translates the desired_sharding_key configuration for a table into a database function and trigger that populates the sharding key column when a record is inserted or updated. This will be used in combination with a batched background migration to ensure that newly added sharding keys are set for all records going forward.

Background

As part of the prerequisite work for Cells, tables that have been identified as cell local are required to have a sharding key. Those that don't yet have a valid sharding key can be assigned a desired sharding key, to be added at some point in the future. We plan to use GitLab Housekeeper to automate adding and populating these desired sharding keys.

Previously we tried to assign these fields in the application using ActiveRecord callbacks, however this was problematic because not all operations trigger callbacks (for example, update_column or bulk insert/updates). This means that cases can be missed, and the column never receives a value. Sometimes this caused spec failures, which required manual work to resolve. In the worst case, without explicit spec coverage this can cause incidents.

Proposal

The desired_sharding_key is populated using a database trigger, instead of being set in application code.

For example, given a desired_sharding_key configuration such as:

desired_sharding_key:
  project_id:
    references: projects
    backfill_via:
      parent:
        foreign_key: agent_id
        table: cluster_agents
        sharding_key: project_id
        belongs_to: agent

We can create a migration using the new helper:

class AddClusterAgentTokensShardingKeyTrigger < Gitlab::Database::Migration[2.2]
  milestone '16.10'

  def up
    install_sharding_key_assignment_trigger(
      table: :cluster_agent_tokens,
      sharding_key: :project_id,
      parent_table: :cluster_agents,
      parent_sharding_key: :project_id,
      foreign_key: :agent_id
    )
  end

  def down
    remove_sharding_key_assignment_trigger(
      table: :cluster_agent_tokens,
      sharding_key: :project_id,
      parent_table: :cluster_agents,
      parent_sharding_key: :project_id,
      foreign_key: :agent_id
    )
  end
end

This creates the following trigger and function:

CREATE FUNCTION trigger_82f3e193fff7() RETURNS trigger
    LANGUAGE plpgsql
    AS $$
    BEGIN
        IF NEW."project_id" IS NULL THEN
            SELECT "project_id"
            INTO NEW."project_id"
            FROM "cluster_agents"
            WHERE "cluster_agents"."id" = NEW."agent_id";
        END IF;

        RETURN NEW;
    END;
$$

...

CREATE TRIGGER trigger_82f3e193fff7 BEFORE INSERT OR UPDATE ON cluster_agent_tokens FOR EACH ROW EXECUTE FUNCTION trigger_82f3e193fff7();

Whenever a cluster_agent_token record is inserted or updated, the trigger fires and fetches the project_id from the parent record (cluster_agents), and assigns it to the project_id on the newly inserted/updated record.

Advantages

• The functionality operates independently of how records are managed in the application, eliminating the risk of missing a special case in application code.
• There should be a small performance improvement, as one less query is required on the application side if the associated record has not been loaded.
• Updates are handled entirely behind the scenes, requiring no specific domain knowledge from the author/reviewer (also a disadvantage, see below).

Disadvantages

• Introduces a tight coupling between migrations and database dictionary files, which is required to formulate the query used by the trigger. While this isn't strictly using application code in migrations which is discouraged, it could still cause headaches in the future.
• Handling updates outside of the application code reduces discoverability.

Open questions

• It it unclear whether these triggers would be intended as a permanent solution, or a temporary measure while we are automatically backfilling sharding keys on many tables. Aside from the downsides listed above, there's no blocker to leaving them in place permanently. However, if temporary we will encounter the same problems we faced initially when trying to solve using application code (though without the time pressure of holding up the Cells initiative).
• The IF FOUND THEN RETURN logic in the trigger aligns the functionality with the default suggested in NOT NULL constraints for multiple columns for tables with more than one possible sharding key. By default, we expect these tables to have exactly one key present (as opposed to at least one), and without an early return the trigger would attempt to set both (raising an error if both succeed). Is this unnecessary complexity?

How to set up and validate locally

1. For any table that has a desired_sharding_key configured, add the column that will eventually be used as the sharding key (project_id or namespace_id), either via migration or manually in the DB console.
2. Create and run a migration similar to the example above, using the table from the previous step
3. From a rails console, pick any record for the table and update any attribute (eg. using touch to set the updated at).
4. Reload the record, and verify that the sharding key is set.

Related to #443604 (closed)